supoort yolov5_sam

anylabeling/app_info.py

@@ -1,4 +1,4 @@
 __appname__ = "X-AnyLabeling"
 __appdescription__ = "Advanced Auto Labeling Solution with Added Features"
-__version__ = "0.2.2"
+__version__ = "0.2.3"
 __preferred_device__ = "CPU"  # GPU or CPU

anylabeling/configs/auto_labeling/models.yaml

@@ -1,5 +1,7 @@
 - model_name: "ch_ppocr_v4-r20230915"
   config_file: ":/ch_ppocr_v4.yaml"
+- model_name: "yolov5s_mobile_sam_vit_h-r20230920"
+  config_file: ":/yolov5s_mobile_sam_vit_h.yaml"
 - model_name: "sam_med2d_vit_b-r20230901"
   config_file: ":/sam_med2d_vit_b.yaml"
 - model_name: "segment_anything_vit_b_quant-r20230810"

anylabeling/configs/auto_labeling/yolov5s_mobile_sam_vit_h.yaml

@@ -0,0 +1,98 @@
+type: yolov5_sam
+name: yolov5s_mobile_sam_vit_h-r20230920
+display_name: YOLOv5s-MobileSAM (ViT-Huge)
+#SAM
+encoder_model_path: https://github.com/CVHub520/X-AnyLabeling/releases/download/v0.2.0/mobile_sam.encoder.onnx
+decoder_model_path: https://github.com/CVHub520/X-AnyLabeling/releases/download/v0.2.0/sam_vit_h_4b8939.decoder.onnx
+cache_size: 10
+target_size: 1024
+max_width: 1024
+max_height: 682
+#Det
+model_path: https://github.com/CVHub520/X-AnyLabeling/releases/download/v0.1.0/yolov5s.onnx
+input_width: 640
+input_height: 640
+stride: 32
+nms_threshold: 0.45
+confidence_threshold: 0.25
+classes:
+  - person
+  - bicycle
+  - car
+  - motorcycle
+  - airplane
+  - bus
+  - train
+  - truck
+  - boat
+  - traffic light
+  - fire hydrant
+  - stop sign
+  - parking meter
+  - bench
+  - bird
+  - cat
+  - dog
+  - horse
+  - sheep
+  - cow
+  - elephant
+  - bear
+  - zebra
+  - giraffe
+  - backpack
+  - umbrella
+  - handbag
+  - tie
+  - suitcase
+  - frisbee
+  - skis
+  - snowboard
+  - sports ball
+  - kite
+  - baseball bat
+  - baseball glove
+  - skateboard
+  - surfboard
+  - tennis racket
+  - bottle
+  - wine glass
+  - cup
+  - fork
+  - knife
+  - spoon
+  - bowl
+  - banana
+  - apple
+  - sandwich
+  - orange
+  - broccoli
+  - carrot
+  - hot dog
+  - pizza
+  - donut
+  - cake
+  - chair
+  - couch
+  - potted plant
+  - bed
+  - dining table
+  - toilet
+  - tv
+  - laptop
+  - mouse
+  - remote
+  - keyboard
+  - cell phone
+  - microwave
+  - oven
+  - toaster
+  - sink
+  - refrigerator
+  - book
+  - clock
+  - vase
+  - scissors
+  - teddy bear
+  - hair drier
+  - toothbrush

anylabeling/services/auto_labeling/model_manager.py

@@ -179,7 +179,8 @@ def load_custom_model(self, config_file):
                 "yolo_nas",
                 "yolox_dwpose",
                 "clrnet",
-                "ppocr_v4"
+                "ppocr_v4",
+                "yolov5_sam",
             ]
         ):
             self.new_model_status.emit(
@@ -435,6 +436,30 @@ def _load_model(self, model_id):
                     )
                 )
                 return
+        elif model_config["type"] == "yolov5_sam":
+            from .yolov5_sam import YOLOv5SegmentAnything
+
+            try:
+                model_config["model"] = YOLOv5SegmentAnything(
+                    model_config, on_message=self.new_model_status.emit
+                )
+                self.auto_segmentation_model_selected.emit()
+            except Exception as e:  # noqa
+                print(
+                    "Error in loading model: {error_message}".format(
+                        error_message=str(e)
+                    )
+                )
+                self.new_model_status.emit(
+                    self.tr(
+                        "Error in loading model: {error_message}".format(
+                            error_message=str(e)
+                        )
+                    )
+                )
+                return
+            # Request next files for prediction
+            self.request_next_files_requested.emit()
         elif model_config["type"] == "segment_anything":
             from .segment_anything import SegmentAnything
 
@@ -625,9 +650,14 @@ def set_auto_labeling_marks(self, marks):
         """Set auto labeling marks
         (For example, for segment_anything model, it is the marks for)
         """
+        marks_model_list = [
+            "segment_anything", 
+            "sam_med2d", 
+            "yolov5_sam",
+        ]
         if (
             self.loaded_model_config is None
-            or self.loaded_model_config["type"] not in ["segment_anything", "sam_med2d"]
+            or self.loaded_model_config["type"] not in marks_model_list
         ):
             return
         self.loaded_model_config["model"].set_auto_labeling_marks(marks)
@@ -714,7 +744,7 @@ def on_next_files_changed(self, next_files):
             return
 
         # Currently only segment_anything-like model supports this feature
-        if self.loaded_model_config["type"] not in ["segment_anything", "sam_med2d"]:
+        if self.loaded_model_config["type"] not in ["segment_anything", "sam_med2d", "yolov5_sam"]:
             return
 
         self.loaded_model_config["model"].on_next_files_changed(next_files)