Fix model path + Add ONNX validation

Author: vietanhdev

anylabeling/configs/auto_labeling/segment_anything_vit_b_quant.yaml

@@ -1,5 +1,5 @@
 type: segment_anything
-name: segment_anything_vit_b-r20230416
+name: segment_anything_vit_b_quant-r20230416
 display_name: Segment Anything (ViT-B Quant)
 encoder_model_path: https://github.com/vietanhdev/anylabeling-assets/releases/download/v0.2.0/segment_anything_vit_b_encoder_quant.onnx
 decoder_model_path: https://github.com/vietanhdev/anylabeling-assets/releases/download/v0.2.0/segment_anything_vit_b_decoder_quant.onnx

anylabeling/services/auto_labeling/model.py

@@ -5,6 +5,7 @@
 from abc import abstractmethod
 
 import yaml
+import onnx
 
 from PyQt5.QtCore import QFile
 from PyQt5.QtGui import QImage
@@ -86,7 +87,17 @@ def get_model_abs_path(self, model_path, model_folder_name):
             )
         )
         if os.path.exists(model_abs_path):
-            return model_abs_path
+            if model_abs_path.lower().endswith(".onnx"):
+                try:
+                    onnx.checker.check_model(model_abs_path)
+                except onnx.checker.ValidationError as e:
+                    print("The model is invalid: %s" % e)
+                    print("Action: Delete and redownload...")
+                    os.remove(model_abs_path)
+                else:
+                    return model_abs_path
+            else:
+                return model_abs_path
         pathlib.Path(model_abs_path).parent.mkdir(parents=True, exist_ok=True)
 
         # Download model from url

anylabeling/services/auto_labeling/segment_anything.py

@@ -32,7 +32,6 @@ class Meta:
             "button_add_point",
             "button_remove_point",
             "button_add_rect",
-            # "button_undo", # Dont support undo now
             "button_clear",
             "button_finish_object",
         ]
@@ -81,6 +80,7 @@ def __init__(self, config_path, on_message) -> None:
 
         # Pre-inference worker
         self.pre_inference_thread = None
+        self.stop_inference = False
 
     def set_auto_labeling_marks(self, marks):
         """Set auto labeling marks"""
@@ -216,6 +216,7 @@ def run_decoder(self, image_embedding, resized_ratio):
             ),
         }
         masks, _, _ = self.decoder_session.run(None, decoder_inputs)
+        masks = masks[0, 0, :, :]  # Only get 1 mask
         masks = masks > 0.0
         masks = masks.reshape(self.size_after_apply_max_width_height)
         return masks
@@ -279,11 +280,15 @@ def predict_shapes(self, image, filename=None) -> AutoLabelingResult:
             else:
                 cv_image = qt_img_to_cv_img(image)
                 encoder_inputs, resized_ratio = self.pre_process(cv_image)
+                if self.stop_inference:
+                    return AutoLabelingResult([], replace=False)
                 image_embedding = self.run_encoder(encoder_inputs)
                 self.image_embedding_cache.put(
                     filename,
                     (resized_ratio, image_embedding),
                 )
+            if self.stop_inference:
+                return AutoLabelingResult([], replace=False)
             masks = self.run_decoder(image_embedding, resized_ratio)
             shapes = self.post_process(masks, resized_ratio)
         except Exception as e:  # noqa
@@ -295,6 +300,9 @@ def predict_shapes(self, image, filename=None) -> AutoLabelingResult:
         return result
 
     def unload(self):
+        self.stop_inference = True
+        self.pre_inference_thread.quit()
+        self.pre_inference_thread.wait()
         if self.encoder_session:
             self.encoder_session = None
         if self.decoder_session:
@@ -311,6 +319,8 @@ def preload_worker(self, files):
             image = self.load_image_from_filename(filename)
             if image is None:
                 continue
+            if self.stop_inference:
+                return
             cv_image = qt_img_to_cv_img(image)
             encoder_inputs, resized_ratio = self.pre_process(cv_image)
             image_embedding = self.run_encoder(encoder_inputs)