perf(uav): switch to one class detection model and --use-vits in clus…

…tering MBARI-org/mbari-uav-vit-b-16

aipipeline/projects/uav/cluster_pipeline.py

@@ -27,7 +27,7 @@
 ENVIRONMENT = os.getenv("ENVIRONMENT") if os.getenv("ENVIRONMENT") else None
 
 
-def process_mission(element, run_vss=False):
+def process_mission(element):
     # Data is in the format
     # <path>,<tator section>,<start image>,<end image>
     # /mnt/UAV/Level-1/trinity-2_20240702T153433_NewBrighton/SONY_DSC-RX1RM2,2024/07/NewBrighton,DSC00100.JPG,DSC00301.JPG
@@ -69,12 +69,9 @@ def process_mission(element, run_vss=False):
         "--device",
         "cuda:0",
         "--skip-visualization",
+        "--use-vits"
     ]
 
-    if run_vss:
-        vss_url = config_dict["vss"]["url"]
-        args.extend(["--vss-url", vss_url])
-
     container = run_docker(
         image=config_dict["docker"]["sdcat"],
         name=f"sdcat-clu-{mission_name}",
@@ -104,18 +101,14 @@ def run_pipeline(argv=None):
     options = PipelineOptions(beam_args)
     conf_files, config_dict = setup_config(args.config)
 
-    run_vss = False
-    if '--vss' in beam_args:
-        run_vss = True
-
     logger.info("Starting cluster pipeline...")
     with beam.Pipeline(options=options) as p:
         (
             p
             | "Read missions" >> beam.io.ReadFromText(args.missions)
             | "Filter comments" >> beam.Filter(lambda line: not line.startswith("#"))
             | "Create elements" >> beam.Map(lambda line: (line, config_dict))
-            | "Process missions (cluster)" >> beam.Map(process_mission, run_vss=run_vss)
+            | "Process missions (cluster)" >> beam.Map(process_mission)
         )
 
 

aipipeline/projects/uav/config/config.yml

@@ -30,7 +30,7 @@ data:
   download_args: ["--verified"]
 
 sdcat:
-  model: "MBARI/uav-yolov5-30k"
+  model: "MBARI/yolov5x6-uavs-oneclass"
   ini: "sdcat.ini"
 
 vss:

aipipeline/projects/uav/config/sdcat.ini

@@ -37,9 +37,10 @@ min_similarity = 0.90
 # dino models were pretrained on ImageNet which contains 1.3 M images with labels from 1000 classes
 # Smaller block_size means more patches and more accurate fine-grained clustering on smaller objects
 # Larger block_size means fewer patches and faster processing
-model = google/vit-base-patch16-224
+;model = google/vit-base-patch16-224
 ;model = facebook/dino-vits8
 ;model = facebook/dino-vits16
+model = MBARI-org/mbari-uav-vit-b-16
 
 
 [detect]