Update model.py

model.py

@@ -1,9 +1,12 @@
 import os
 import logging
+import torch
 
 from transformers import AutoImageProcessor, AutoModelForObjectDetection
 from label_studio_ml.model import LabelStudioMLBase
 from lxml import etree
+from uuid import uuid4
+from PIL import Image
 
 
 class Model(LabelStudioMLBase):
@@ -26,23 +29,46 @@ class Model(LabelStudioMLBase):
         """
         predictions = []
         for task in tasks:
-            predictions.append({
-                'score': 0.987,  # prediction overall score, visible in the data manager columns
-                'model_version': 'delorean-20151021',  # all predictions will be differentiated by model version
-                'result': [{
-                    'from_name': self.from_name,
-                    'to_name': self.to_name,
-                    'type': 'choices',
-                    'score': 0.5,  # per-region score, visible in the editor 
-                    'value': {
-                        'choices': [self.labels[0]]
+
+
+            image_path = task["data"]["image"]
+            image = Image(image_path)
+                original_width, original_height = image.size
+            with torch.no_grad():
+                
+                inputs = image_processor(images=image, return_tensors="pt")
+                outputs = model(**inputs)
+                target_sizes = torch.tensor([image.size[::-1]])
+                results = image_processor.post_process_object_detection(outputs, threshold=0.5, target_sizes=target_sizes)[0]
+
+            result_list = []
+            for score, label, box in zip(results['scores'], results['labels'], scores['boxes']):
+                label_id = str(uuid4())[:4]
+                x, y, x2, y2 = tuple(box)
+                result_list.append(
+                    {
+                        'id': id
+                        'original_width': original_width,
+                        'original_height': original_height,
+                        'from_name': "label",
+                        'to_name': "image",
+                        'type': 'labels',
+                        'score': score,  # per-region score, visible in the editor 
+                        'value': {
+                            'x': x,
+                            'y': y,
+                            'width': x2-x,
+                            'height': y2-y,
+                            'rotation': 0
+                            'labels': [self.id2label[label]]
+                        } 
                     }
-                }]
+                )
+        
+            
+            predictions.append({
+                'score': results['scores'].mean(),  # prediction overall score, visible in the data manager columns
+                'model_version': 'diegokauer/conditional-detr-coe-int',  # all predictions will be differentiated by model version
+                'result': result_list
             })
-        return predictions
-
-    def fit(self, annotations, **kwargs):
-        """ This is where training happens: train your model given list of annotations, 
-            then returns dict with created links and resources
-        """
-        return {'path/to/created/model': 'my/model.bin'}
+        return predictions