diegokauer commited on
Commit
fcb8bb5
·
1 Parent(s): 0e19682

Update model.py

Browse files
Files changed (1) hide show
  1. model.py +3 -12
model.py CHANGED
@@ -14,15 +14,6 @@ class Model(LabelStudioMLBase):
14
  image_processor = AutoImageProcessor.from_pretrained("diegokauer/conditional-detr-coe-int")
15
  model = AutoModelForObjectDetection.from_pretrained("diegokauer/conditional-detr-coe-int")
16
 
17
- def __init__(self, **kwargs):
18
- # don't forget to call base class constructor
19
- super(Model, self).__init__(**kwargs)
20
-
21
- # you can preinitialize variables with keys needed to extract info from tasks and annotations and form predictions
22
- self.model = model
23
- self.tokenizer = image_processor
24
- self.id2label = model.config.id2label
25
-
26
  def predict(self, tasks, **kwargs):
27
  """ This is where inference happens: model returns
28
  the list of predictions based on input list of tasks
@@ -36,10 +27,10 @@ class Model(LabelStudioMLBase):
36
  original_width, original_height = image.size
37
  with torch.no_grad():
38
 
39
- inputs = self.image_processor(images=image, return_tensors="pt")
40
- outputs = self.model(**inputs)
41
  target_sizes = torch.tensor([image.size[::-1]])
42
- results = self.image_processor.post_process_object_detection(outputs, threshold=0.5, target_sizes=target_sizes)[0]
43
 
44
  result_list = []
45
  for score, label, box in zip(results['scores'], results['labels'], results['boxes']):
 
14
  image_processor = AutoImageProcessor.from_pretrained("diegokauer/conditional-detr-coe-int")
15
  model = AutoModelForObjectDetection.from_pretrained("diegokauer/conditional-detr-coe-int")
16
 
 
 
 
 
 
 
 
 
 
17
  def predict(self, tasks, **kwargs):
18
  """ This is where inference happens: model returns
19
  the list of predictions based on input list of tasks
 
27
  original_width, original_height = image.size
28
  with torch.no_grad():
29
 
30
+ inputs = image_processor(images=image, return_tensors="pt")
31
+ outputs = model(**inputs)
32
  target_sizes = torch.tensor([image.size[::-1]])
33
+ results = image_processor.post_process_object_detection(outputs, threshold=0.5, target_sizes=target_sizes)[0]
34
 
35
  result_list = []
36
  for score, label, box in zip(results['scores'], results['labels'], results['boxes']):