diff --git a/aipipeline/projects/uav/cluster_pipeline.py b/aipipeline/projects/uav/cluster_pipeline.py index 3e4bda7..5ce5245 100644 --- a/aipipeline/projects/uav/cluster_pipeline.py +++ b/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 # ,,, # /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,10 +101,6 @@ 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: ( @@ -115,7 +108,7 @@ def run_pipeline(argv=None): | "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) ) diff --git a/aipipeline/projects/uav/config/config.yml b/aipipeline/projects/uav/config/config.yml index 0e2574e..43aa2d1 100644 --- a/aipipeline/projects/uav/config/config.yml +++ b/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: diff --git a/aipipeline/projects/uav/config/sdcat.ini b/aipipeline/projects/uav/config/sdcat.ini index 0595503..3a50042 100644 --- a/aipipeline/projects/uav/config/sdcat.ini +++ b/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]