Update model.py

model.py

@@ -5,7 +5,7 @@ import datetime
 import requests
 
 from google.cloud import storage
-from transformers import AutoImageProcessor, AutoModelForObjectDetection
+from transformers import AutoImageProcessor, AutoModelForObjectDetection, ViTImageProcessor, Swinv2ForImageClassification
 from label_studio_ml.model import LabelStudioMLBase
 from lxml import etree
 from uuid import uuid4
@@ -43,7 +43,10 @@ class Model(LabelStudioMLBase):
     os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = get_credentials()
     image_processor = AutoImageProcessor.from_pretrained("diegokauer/conditional-detr-coe-int")
     model = AutoModelForObjectDetection.from_pretrained("diegokauer/conditional-detr-coe-int")
+    seg_image_processor = ViTImageProcessor.from_pretrained("diegokauer/int-pet-classifier")
+    seg_model = SwinForImageClassification.from_pretrained("diegokauer/int-pet-classifier")
     id2label = model.config.id2label
+    seg_id2label = seg_model.config.id2label
 
     def predict(self, tasks, **kwargs):
         """ This is where inference happens: model returns 
@@ -71,6 +74,26 @@ class Model(LabelStudioMLBase):
             for score, label, box in zip(results['scores'], results['labels'], results['boxes']):
                 label_id = str(uuid4())
                 x, y, x2, y2 = tuple(box)
+
+                if self.id2label[label.item()] == 'Propuesta':
+                    pred_label_id = str(uuid4())
+                    image = image.crop((x, y, x2, y2))
+                    input = self.seg_image_processor(images=image, return_tensors="pt")
+                    logits = self.seg_model(**inputs).logits
+                    logits = torch.exp(logits)
+                    preds = logits > 0.5
+                    preds = [self.seg_id2label[i] for i, pred in enumerate(preds)]
+                    preds = ["No Reportado"] if "No Reportado" in preds else preds
+                    result_list.append({
+                      "value": {
+                        "choices": preds
+                      },
+                      "id": pred_label_id,
+                      "from_name": "propuesta",
+                      "to_name": "image",
+                      "type": "choices"
+                    })
+                
                 result_list.append({
                     'id': label_id,
                     'original_width': original_width,
@@ -107,7 +130,7 @@ class Model(LabelStudioMLBase):
             
             predictions.append({
                 'score': results['scores'].mean().item(),  # prediction overall score, visible in the data manager columns
-                'model_version': 'cdetr_v1',  # all predictions will be differentiated by model version
+                'model_version': 'cdetr_v2',  # all predictions will be differentiated by model version
                 'result': result_list
             })
         print(predictions)