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.unmasker = 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=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)