bento / service.py
davanstrien's picture
davanstrien HF staff
Update service.py
477cb68
raw
history blame contribute delete
898 Bytes
import bentoml
from bentoml.io import Image, JSON
from transformers import pipeline
class PretrainedModelRunnable(bentoml.Runnable):
SUPPORTED_RESOURCES = ("cpu",)
SUPPORTS_CPU_MULTI_THREADING = True
def __init__(self):
self.object_detector = pipeline('object-detection', model="biglam/detr-resnet-50_fine_tuned_nls_chapbooks",
feature_extractor="biglam/detr-resnet-50_fine_tuned_nls_chapbooks")
@bentoml.Runnable.method(batchable=False)
def __call__(self, input_image):
return self.object_detector(input_image)
runner = bentoml.Runner(PretrainedModelRunnable, name="pretrained_illustration_detection")
svc = bentoml.Service('pretrained_illustration_detection', runners=[runner])
@svc.api(input=Image(), output=JSON())
async def object_detect(input_series: str) -> list:
return await runner.async_run(input_series)