SDK Reference
sieve.Model
def Model(
name: str,
gpu: bool = False,
python_packages: List[str] = [],
python_version: str = "3.8",
system_packages: List[str] = [],
cuda_version: str = None,
machine_type: str = None,
iterator_input: bool = False,
run_commands: List[str] = []
)
A sieve.Model
is a Python class that defines a model in Sieve. A model is a function that has a __setup__
and a __predict__
step. The __setup__
step is run once when the model initializes, and the __predict__
step is run every time the model is called.
YOLO Object Detection Example
import sieve
from typing import List, Dict
@sieve.Model(
name="yolo",
gpu = True,
python_packages=[
"torch==1.8.1",
"pandas==1.5.2",
"opencv-python-headless==4.5.4.60",
"ipython==8.4.0",
"torch==1.8.1",
"torchvision==0.9.1",
"psutil==5.8.0",
"seaborn==0.11.2"
],
system_packages=["libgl1-mesa-glx", "libglib2.0-0", "ffmpeg"],
python_version="3.8"
)
class Yolo:
def __setup__(self):
import torch
self.yolo_model = torch.hub.load('ultralytics/yolov5', 'yolov5l')
def __predict__(self, img: sieve.Image) -> List:
results = self.yolo_model(img.array)
outputs = []
for pred in reversed(results.pred):
for *box, conf, cls in reversed(pred):
cls_name = results.names[int(cls)]
box = [float(i) for i in box]
score = float(conf)
if score < 0.7:
continue
outputs.append({
"box": box,
"class_name": cls_name,
"score": score,
"frame_number": None if not hasattr(img, "frame_number") else img.frame_number
})
return outputs
Arguments
name
Name of the function.gpu
Whether the function should run on a GPU. Defaults toFalse
.python_packages
List of Python packages to install. Defaults to[]
.python_version
Python version to use. Defaults to"3.8"
.system_packages
List of system packages to install. Defaults to[]
.cuda_version
CUDA version to use. Defaults toNone
.machine_type
Machine type to use. Defaults toNone
.iterator_input
Whether the function takes in an iterator. Defaults toFalse
.run_commands
List of commands to run when the function is built. Defaults to[]
.
name
def name(self) -> str:
Name of the function.
gpu
def gpu(self) -> bool:
Whether the function should run on a GPU.
python_packages
def python_packages(self) -> List[str]:
List of Python packages to install.
python_version
def python_version(self) -> str:
Python version to use.
system_packages
def system_packages(self) -> List[str]:
List of system packages to install.
cuda_version
def cuda_version(self) -> str:
CUDA version to use.
machine_type
def machine_type(self) -> str:
Machine type to use.
iterator_input
def iterator_input(self) -> bool:
Whether the function takes in an iterator.
run_commands
def run_commands(self) -> List[str]:
List of commands to run when the function is built.