segmentation-backend

Sleeping

diegokauer commited on Jan 3, 2024

Commit

e63e5de

1 Parent(s): 307cad5

Update model.py

Files changed (1) hide show

model.py CHANGED Viewed

@@ -44,7 +44,7 @@ class Model(LabelStudioMLBase):
     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

     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 = Swinv2ForImageClassification.from_pretrained("diegokauer/int-pet-classifier")
     id2label = model.config.id2label
     seg_id2label = seg_model.config.id2label