service.py

import bentoml

from bentoml.io import Image, JSON
from transformers import pipeline
from transformers import AutoFeatureExtractor, AutoModelForObjectDetection
import transformers

transformers.logging.set_verbosity_info()


class PretrainedModelRunnable(bentoml.Runnable):
    SUPPORTED_RESOURCES = ("cpu",)
    SUPPORTS_CPU_MULTI_THREADING = True

    def __init__(self):
        self.unmasker = pipeline('object-detection', model="davanstrien/detr-resnet-50_fine_tuned_nls_chapbooks",
                                 feature_extractor="davanstrien/detr-resnet-50_fine_tuned_nls_chapbooks")

    @bentoml.Runnable.method(batchable=True)
    def __call__(self, input_text):
        return self.unmasker(input_text)


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)