feat: added latest 30 and 18k models

mbari-org · Aug 20, 2024 · 673a7f6 · 673a7f6
1 parent 56b4e17
commit 673a7f6
Showing 1 changed file with 13 additions and 2 deletions.
diff --git a/sdcat/detect/commands.py b/sdcat/detect/commands.py
@@ -109,10 +109,21 @@ def run_detect(show: bool, image_dir: str, save_dir: str, model: str,
                 confidence_threshold=conf,
                 device=device,
             )
-        elif model == 'MBARI/uav-yolov5':
+        elif model == 'MBARI/uav-yolov5-30k':
             # Download model path
             from huggingface_hub import hf_hub_download
-            model_path = hf_hub_download(repo_id="MBARI-org/uav-yolov5", filename="best.pt")
+            model_path = hf_hub_download(repo_id="MBARI-org/yolov5x6-uav-30k", filename="yolov5x6-uav-30k.pt")
+            detection_model = AutoDetectionModel.from_pretrained(
+                model_type='yolov5',
+                model_path=model_path,
+                config_path=model_path,
+                confidence_threshold=conf,
+                device=device,
+            )
+        elif model == 'MBARI/uav-yolov5-18k':
+            # Download model path
+            from huggingface_hub import hf_hub_download
+            model_path = hf_hub_download(repo_id="MBARI-org/yolov5-uav-18k", filename="yolov5x6-uav-18k.pt")
             detection_model = AutoDetectionModel.from_pretrained(
                 model_type='yolov5',
                 model_path=model_path,