davanstrien HF staff commited on
Commit
477cb68
1 Parent(s): f78ab1e

Update service.py

Browse files
Files changed (1) hide show
  1. service.py +4 -4
service.py CHANGED
@@ -8,12 +8,12 @@ class PretrainedModelRunnable(bentoml.Runnable):
8
  SUPPORTS_CPU_MULTI_THREADING = True
9
 
10
  def __init__(self):
11
- self.unmasker = pipeline('object-detection', model="biglam/detr-resnet-50_fine_tuned_nls_chapbooks",
12
  feature_extractor="biglam/detr-resnet-50_fine_tuned_nls_chapbooks")
13
 
14
- @bentoml.Runnable.method(batchable=True)
15
- def __call__(self, input_text):
16
- return self.unmasker(input_text)
17
 
18
 
19
  runner = bentoml.Runner(PretrainedModelRunnable, name="pretrained_illustration_detection")
 
8
  SUPPORTS_CPU_MULTI_THREADING = True
9
 
10
  def __init__(self):
11
+ self.object_detector = pipeline('object-detection', model="biglam/detr-resnet-50_fine_tuned_nls_chapbooks",
12
  feature_extractor="biglam/detr-resnet-50_fine_tuned_nls_chapbooks")
13
 
14
+ @bentoml.Runnable.method(batchable=False)
15
+ def __call__(self, input_image):
16
+ return self.object_detector(input_image)
17
 
18
 
19
  runner = bentoml.Runner(PretrainedModelRunnable, name="pretrained_illustration_detection")