From 5f97fc461d633b3d2fa2671320b63590d283e940 Mon Sep 17 00:00:00 2001
From: danellecline <dcline@mbari.org>
Date: Tue, 21 Jan 2025 17:40:28 -0800
Subject: [PATCH] refactor(uav): move all arguments that may change to
 config.yml for detect/classify pipelines

---
 aipipeline/projects/uav/cluster_pipeline.py |  8 +++++---
 aipipeline/projects/uav/config/config.yml   | 18 ++++++++++++++++--
 aipipeline/projects/uav/detect_pipeline.py  | 21 +++------------------
 3 files changed, 24 insertions(+), 23 deletions(-)

diff --git a/aipipeline/projects/uav/cluster_pipeline.py b/aipipeline/projects/uav/cluster_pipeline.py
index 7fb5774..a8f563f 100644
--- a/aipipeline/projects/uav/cluster_pipeline.py
+++ b/aipipeline/projects/uav/cluster_pipeline.py
@@ -49,6 +49,7 @@ def process_mission(element):
 
     det_dir = base_path / mission_name / "detections" / "combined" / model / "det_filtered"
     save_dir = base_path / mission_name / "detections" / "combined" / model / "clusters"
+    add_args = config_dict["sdcat"]["cluster_args"]
 
     if not det_dir.exists():
         logger.error(f"Could not find directory: {det_dir}")
@@ -67,11 +68,12 @@ def process_mission(element):
         "--save-dir",
         str(save_dir),
         "--device",
-        str(gpu_device),
-        "--skip-visualization",
-        "--use-vits"
+        str(gpu_device)
     ]
 
+    if add_args:
+        args += add_args
+
     container = run_docker(
         image=config_dict["docker"]["sdcat"],
         name=f"sdcat-clu-{mission_name}",
diff --git a/aipipeline/projects/uav/config/config.yml b/aipipeline/projects/uav/config/config.yml
index b5dd871..9d862a2 100644
--- a/aipipeline/projects/uav/config/config.yml
+++ b/aipipeline/projects/uav/config/config.yml
@@ -36,9 +36,23 @@ data:
     - name: "near_duplicates"
 
 sdcat:
-  model: "/mnt/DeepSea-AI/models/MBARI-org--yolov5x6-uavs-oneclass"
   ini: "sdcat.ini"
-
+  detect_args:
+    - "--scale-percent"
+    - "100"
+    - "--model"
+    - "/mnt/DeepSea-AI/models/UAV/mbari-uav-yolo11x-jelly/"
+    - "--model-type"
+    - "yolov11
+    - "--slice-size-width"
+    - "640"
+    - "--slice-size-height"
+    - "640"
+    - "--conf"
+    - "0.1"
+  cluster_args:
+    - "--skip-visualization"
+    - "--use-vits"
 vss:
   project: "901902-uavs"
   threshold: 0.1
diff --git a/aipipeline/projects/uav/detect_pipeline.py b/aipipeline/projects/uav/detect_pipeline.py
index 6643bb2..a52fcdd 100644
--- a/aipipeline/projects/uav/detect_pipeline.py
+++ b/aipipeline/projects/uav/detect_pipeline.py
@@ -3,25 +3,18 @@
 # Description: Batch process missions with sdcat detection
 import multiprocessing
 import os
-import uuid
 from datetime import datetime
 from typing import Any
 
 import apache_beam as beam
-import pandas as pd
 from apache_beam.options.pipeline_options import PipelineOptions
 from apache_beam.io import ReadFromText
 from pathlib import Path
-from multiprocessing import Pool
 import logging
-import io
-import tqdm
 
 from aipipeline.docker.utils import run_docker
 from aipipeline.projects.uav.args_common import parse_args, POSSIBLE_PLATFORMS, parse_mission_string
 from aipipeline.config_setup import setup_config, SDCAT_KEY
-from aipipeline.prediction.library import run_vss
-from aipipeline.prediction.utils import crop_square_image
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
@@ -47,7 +40,7 @@ def run_mission_detect(element) -> Any:
     gpu_device, mission_name, mission_dir, section, start_image, end_image = parse_mission_string(line)
 
     base_path = Path(config_dict["data"]["processed_path_sdcat"]) / "seedDetections"
-    model = config_dict["sdcat"]["model"]
+    add_args = config_dict["sdcat"]["detect_args"]
 
     if not mission_name:
         logger.error(f"Could not find mission name in path: {mission_dir} that starts with {POSSIBLE_PLATFORMS}")
@@ -67,21 +60,13 @@ def run_mission_detect(element) -> Any:
         str(gpu_device),
         "--config-ini",
         conf_files[SDCAT_KEY],
-        "--scale-percent",
-        "50",
-        "--model",
-        model,
-        "--slice-size-width",
-        "1280",
-        "--slice-size-height",
-        "1280",
-        "--conf",
-        "0.1",
         "--save-dir",
         str(save_dir),
         "--image-dir",
         mission_dir,
     ]
+    if add_args:
+        args += add_args
 
     if start_image:
         args += ["--start-image", start_image]