main

  1import datetime
  2import logging
  3import os
  4import signal
  5import sys
  6import time
  7import warnings
  8from typing import Dict, List
  9
 10import torch
 11
 12IGNORE_FUTURE_WARNINGS = True
 13if IGNORE_FUTURE_WARNINGS:
 14    warnings.simplefilter("ignore", category=FutureWarning)
 15
 16import gc
 17
 18import fiftyone as fo
 19import torch.multiprocessing as mp
 20from tqdm import tqdm
 21
 22from config.config import (
 23    SELECTED_DATASET,
 24    SELECTED_WORKFLOW,
 25    V51_ADDRESS,
 26    V51_PORT,
 27    V51_REMOTE,
 28    WORKFLOWS,
 29)
 30from utils.anomaly_detection_data_preparation import AnomalyDetectionDataPreparation
 31from utils.data_loader import FiftyOneTorchDatasetCOCO, TorchToHFDatasetCOCO
 32from utils.dataset_loader import load_dataset
 33from utils.logging import configure_logging
 34from utils.mp_distribution import ZeroShotDistributer
 35from utils.sidebar_groups import arrange_fields_in_groups
 36from utils.wandb_helper import wandb_close, wandb_init
 37from workflows.anomaly_detection import Anodec
 38from workflows.auto_labeling import (
 39    CustomCoDETRObjectDetection,
 40    HuggingFaceObjectDetection,
 41    UltralyticsObjectDetection,
 42    ZeroShotObjectDetection,
 43)
 44from workflows.aws_download import AwsDownloader
 45from workflows.class_mapping import ClassMapper
 46from workflows.embedding_selection import EmbeddingSelection
 47from workflows.ensemble_selection import EnsembleSelection
 48from workflows.auto_label_mask import AutoLabelMask
 49from workflows.class_mapping import ClassMapper
 50
 51wandb_run = None  # Init globally to make sure it is available
 52
 53def signal_handler(sig, frame):
 54    """Handle Ctrl+C signal by cleaning up resources and exiting."""
 55    logging.error("You pressed Ctrl+C!")
 56    try:
 57        wandb_close(exit_code=1)
 58        cleanup_memory()
 59    except:
 60        pass
 61    sys.exit(0)
 62
 63
 64def workflow_aws_download(parameters, wandb_activate=True):
 65    """Download and process data from AWS S3 bucket."""
 66    dataset = None
 67    dataset_name = None
 68    wandb_exit_code = 0
 69    files_to_be_downloaded = 0
 70    try:
 71        # Config
 72        bucket = parameters["bucket"]
 73        prefix = parameters["prefix"]
 74        download_path = parameters["download_path"]
 75        test_run = parameters["test_run"]
 76
 77        # Logging
 78        now = datetime.datetime.now()
 79        datetime_str = now.strftime("%Y-%m-%d_%H-%M-%S")
 80        log_dir = f"logs/tensorboard/aws_{datetime_str}"
 81
 82        dataset_name = f"annarbor_rolling_{datetime_str}"
 83
 84        # Weights and Biases
 85        wandb_run = wandb_init(
 86            run_name=dataset_name,
 87            project_name="AWS Download",
 88            dataset_name=dataset_name,
 89            log_dir=log_dir,
 90            wandb_activate=wandb_activate,
 91        )
 92
 93        # Workflow
 94        aws_downloader = AwsDownloader(
 95            bucket=bucket,
 96            prefix=prefix,
 97            download_path=download_path,
 98            test_run=test_run,
 99        )
100
101        (
102            sub_folder,
103            files,
104            files_to_be_downloaded,
105            DOWNLOAD_NUMBER_SUCCESS,
106            DOWNLOAD_SIZE_SUCCESS,
107        ) = aws_downloader.download_files(log_dir=log_dir)
108
109        dataset = aws_downloader.decode_data(
110            sub_folder=sub_folder,
111            files=files,
112            log_dir=log_dir,
113            dataset_name=dataset_name,
114        )
115
116    except Exception as e:
117        logging.error(f"AWS Download and Extraction failed: {e}")
118        wandb_exit_code = 1
119
120    finally:
121        wandb_close(wandb_exit_code)
122
123    return dataset, dataset_name, files_to_be_downloaded
124
125
126def workflow_anomaly_detection(
127    dataset_normal,
128    dataset_ano_dec,
129    dataset_info,
130    eval_metrics,
131    run_config,
132    wandb_activate=True,
133):
134    """Run anomaly detection workflow using specified models and configurations."""
135    try:
136        # Weights and Biases
137        wandb_exit_code = 0
138        wandb_run, log_dir = wandb_init(
139            run_name=run_config["model_name"],
140            project_name="Selection by Anomaly Detection",
141            dataset_name=dataset_info["name"],
142            config=run_config,
143            wandb_activate=wandb_activate,
144        )
145
146        # Workflow
147
148        SUPPORTED_MODES = ["train", "inference"]
149        # Check if all selected modes are supported
150        for mode in run_config["mode"]:
151            if mode not in SUPPORTED_MODES:
152                logging.error(f"Selected mode {mode} is not supported.")
153        if SUPPORTED_MODES[0] in run_config["mode"]:
154            ano_dec = Anodec(
155                dataset=dataset_ano_dec,
156                eval_metrics=eval_metrics,
157                dataset_info=dataset_info,
158                config=run_config,
159                tensorboard_output=log_dir,
160            )
161            ano_dec.train_and_export_model()
162            ano_dec.run_inference(mode=SUPPORTED_MODES[0])
163            ano_dec.eval_v51()
164        if SUPPORTED_MODES[1] in run_config["mode"]:
165            ano_dec = Anodec(
166                dataset=dataset_normal,
167                eval_metrics=eval_metrics,
168                dataset_info=dataset_info,
169                config=run_config,
170                tensorboard_output=log_dir,
171            )
172            ano_dec.run_inference(mode=SUPPORTED_MODES[1])
173
174    except Exception as e:
175        logging.error(
176            f"Error in Anomaly Detection for model {run_config['model_name']}: {e}"
177        )
178        wandb_exit_code = 1
179    finally:
180        wandb_close(exit_code=wandb_exit_code)
181
182    return True
183
184
185def workflow_embedding_selection(
186    dataset, dataset_info, MODEL_NAME, config, wandb_activate=True
187):
188    """Compute embeddings and find representative and rare images for dataset selection."""
189    try:
190        wandb_exit_code = 0
191        wandb_run, log_dir = wandb_init(
192            run_name=MODEL_NAME,
193            project_name="Selection by Embedding",
194            dataset_name=dataset_info["name"],
195            config=config,
196            wandb_activate=wandb_activate,
197        )
198        embedding_selector = EmbeddingSelection(
199            dataset, dataset_info, MODEL_NAME, log_dir
200        )
201
202        if embedding_selector.model_already_used == False:
203
204            embedding_selector.compute_embeddings(config["mode"])
205            embedding_selector.compute_similarity()
206
207            # Find representative and unique samples as center points for further selections
208            thresholds = config["parameters"]
209            embedding_selector.compute_representativeness(
210                thresholds["compute_representativeness"]
211            )
212            embedding_selector.compute_unique_images_greedy(
213                thresholds["compute_unique_images_greedy"]
214            )
215            embedding_selector.compute_unique_images_deterministic(
216                thresholds["compute_unique_images_deterministic"]
217            )
218
219            # Select samples similar to the center points to enlarge the dataset
220            embedding_selector.compute_similar_images(
221                thresholds["compute_similar_images"], thresholds["neighbour_count"]
222            )
223        else:
224            logging.warning(
225                f"Skipping model {embedding_selector.model_name_key}. It was already used for sample selection."
226            )
227
228    except Exception as e:
229        logging.error(f"An error occurred with model {MODEL_NAME}: {e}")
230        wandb_exit_code = 1
231    finally:
232        wandb_close(wandb_exit_code)
233
234    return True
235
236
237def workflow_auto_labeling_ultralytics(dataset, run_config, wandb_activate=True):
238    """Auto-labeling workflow using Ultralytics models with optional training and inference."""
239    try:
240        wandb_exit_code = 0
241        wandb_run = wandb_init(
242            run_name=run_config["model_name"],
243            project_name="Auto Labeling Ultralytics",
244            dataset_name=run_config["v51_dataset_name"],
245            config=run_config,
246            wandb_activate=wandb_activate,
247        )
248
249        detector = UltralyticsObjectDetection(dataset=dataset, config=run_config)
250
251        # Check if all selected modes are supported
252        SUPPORTED_MODES = ["train", "inference"]
253        for mode in run_config["mode"]:
254            if mode not in SUPPORTED_MODES:
255                logging.error(f"Selected mode {mode} is not supported.")
256
257        if SUPPORTED_MODES[0] in run_config["mode"]:
258            logging.info(f"Training model {run_config['model_name']}")
259            detector.train()
260        if SUPPORTED_MODES[1] in run_config["mode"]:
261            logging.info(f"Running inference for model {run_config['model_name']}")
262            detector.inference()
263
264    except Exception as e:
265        logging.error(f"An error occurred with model {run_config['model_name']}: {e}")
266        wandb_exit_code = 1
267
268    finally:
269        wandb_close(wandb_exit_code)
270
271    return True
272
273
274def workflow_auto_labeling_hf(dataset, hf_dataset, run_config, wandb_activate=True):
275    """Auto-labeling using Hugging Face models on a dataset, including training and/or inference based on the provided configuration."""
276    try:
277        wandb_exit_code = 0
278        wandb_run = wandb_init(
279            run_name=run_config["model_name"],
280            project_name="Auto Labeling Hugging Face",
281            dataset_name=run_config["v51_dataset_name"],
282            config=run_config,
283            wandb_activate=wandb_activate,
284        )
285
286        detector = HuggingFaceObjectDetection(
287            dataset=dataset,
288            config=run_config,
289        )
290        SUPPORTED_MODES = ["train", "inference"]
291
292        # Check if all selected modes are supported
293        for mode in run_config["mode"]:
294            if mode not in SUPPORTED_MODES:
295                logging.error(f"Selected mode {mode} is not supported.")
296        if SUPPORTED_MODES[0] in run_config["mode"]:
297            logging.info(f"Training model {run_config['model_name']}")
298            detector.train(hf_dataset)
299        if SUPPORTED_MODES[1] in run_config["mode"]:
300            logging.info(f"Running inference for model {run_config['model_name']}")
301            detector.inference(inference_settings=run_config["inference_settings"])
302
303    except Exception as e:
304        logging.error(f"An error occurred with model {run_config['model_name']}: {e}")
305        wandb_exit_code = 1
306
307    finally:
308        wandb_close(wandb_exit_code)
309
310    return True
311
312
313def workflow_auto_labeling_custom_codetr(
314    dataset, dataset_info, run_config, wandb_activate=True
315):
316    """Auto labeling workflow using Co-DETR model supporting training and inference modes."""
317
318    try:
319        wandb_exit_code = 0
320        wandb_run = wandb_init(
321            run_name=run_config["config"],
322            project_name="Co-DETR Auto Labeling",
323            dataset_name=dataset_info["name"],
324            config=run_config,
325            wandb_activate=wandb_activate,
326        )
327
328        mode = run_config["mode"]
329
330        detector = CustomCoDETRObjectDetection(dataset, dataset_info, run_config)
331        detector.convert_data()
332        if "train" in mode:
333            detector.update_config_file(
334                dataset_name=dataset_info["name"],
335                config_file=run_config["config"],
336                max_epochs=run_config["epochs"],
337            )
338            detector.train(
339                run_config["config"], run_config["n_gpus"], run_config["container_tool"]
340            )
341        if "inference" in mode:
342            detector.run_inference(
343                dataset,
344                run_config["config"],
345                run_config["n_gpus"],
346                run_config["container_tool"],
347                run_config["inference_settings"],
348            )
349    except Exception as e:
350        logging.error(f"Error during CoDETR training: {e}")
351        wandb_exit_code = 1
352    finally:
353        wandb_close(wandb_exit_code)
354
355    return True
356
357
358def workflow_zero_shot_object_detection(dataset, dataset_info, config):
359    """Run zero-shot object detection on a dataset using models from Huggingface, supporting both single and multi-GPU inference."""
360    # Set multiprocessing mode for CUDA multiprocessing
361    try:
362        mp.set_start_method("spawn", force=True)
363        logging.debug("Successfully set multiprocessing start method to 'spawn'")
364    except RuntimeError as e:
365        # This is expected if the start method was already set
366        logging.debug(f"Multiprocessing start method was already set: {e}")
367    except Exception as e:
368        # Handle any other unexpected errors
369        logging.error(f"Failed to set multiprocessing start method: {e}")
370
371    # Zero-shot object detector models from Huggingface
372    # Optimized for parallel multi-GPU inference, also supports single GPU
373    dataset_torch = FiftyOneTorchDatasetCOCO(dataset)
374    detector = ZeroShotObjectDetection(
375        dataset_torch=dataset_torch, dataset_info=dataset_info, config=config
376    )
377
378    # Check if model detections are already stored in V51 dataset or on disk
379    models_splits_dict = detector.exclude_stored_predictions(
380        dataset_v51=dataset, config=config
381    )
382    if len(models_splits_dict) > 0:
383        config["hf_models_zeroshot_objectdetection"] = models_splits_dict
384        distributor = ZeroShotDistributer(
385            config=config,
386            n_samples=len(dataset_torch),
387            dataset_info=dataset_info,
388            detector=detector,
389        )
390        distributor.distribute_and_run()
391    else:
392        logging.info(
393            "All zero shot models already have predictions stored in the dataset."
394        )
395
396    # To make new fields available to follow-up processes
397    dataset.reload()
398    dataset.save()
399
400    return True
401
402
403def workflow_auto_label_mask(dataset, dataset_info, config):
404    try:
405        depth_config = config["depth_estimation"]
406        seg_config = config["semantic_segmentation"]
407
408        for architecture_name, architecture_info in depth_config.items():
409            auto_labeler = AutoLabelMask(
410                dataset=dataset,
411                dataset_info=dataset_info,
412                model_name=architecture_name,
413                task_type="depth_estimation",
414                model_config=architecture_info,
415            )
416            auto_labeler.run_inference()
417
418        for architecture_name, architecture_info in seg_config.items():
419            auto_labeler = AutoLabelMask(
420                dataset=dataset,
421                dataset_info=dataset_info,
422                model_name=architecture_name,
423                task_type="semantic_segmentation",
424                model_config=architecture_info,
425            )
426            auto_labeler.run_inference()
427
428    except Exception as e:
429        logging.error(f"Auto-labeling mask workflow failed: {e}")
430        raise
431
432
433def workflow_ensemble_selection(dataset, dataset_info, run_config, wandb_activate=True):
434    """Runs ensemble selection workflow on given dataset using provided configuration."""
435    try:
436        wandb_exit_code = 0
437
438        wandb_run = wandb_init(
439            run_name="Selection by Ensemble",
440            project_name="Ensemble Selection",
441            dataset_name=dataset_info["name"],
442            config=run_config,
443            wandb_activate=wandb_activate,
444        )
445        ensemble_selecter = EnsembleSelection(dataset, run_config)
446        ensemble_selecter.ensemble_selection()
447    except Exception as e:
448        logging.error(f"An error occured during Ensemble Selection: {e}")
449        wandb_exit_code = 1
450
451    finally:
452        wandb_close(wandb_exit_code)
453
454    return True
455
456
457def workflow_class_mapping(
458    dataset,
459    dataset_info,
460    run_config,
461    wandb_activate=True,
462    test_dataset_source=None,
463    test_dataset_target=None,
464):
465    """Runs class mapping workflow to align labels between the source dataset and target dataset."""
466    try:
467        wandb_exit_code = 0
468        # Initialize a wandb run for class mapping
469        wandb_run = wandb_init(
470            run_name="class_mapping",
471            project_name="Class Mapping",
472            dataset_name=dataset_info["name"],
473            config=run_config,
474            wandb_activate=wandb_activate,
475        )
476        class_mapping_models = run_config["hf_models_zeroshot_classification"]
477
478        for model_name in (
479            pbar := tqdm(class_mapping_models, desc="Processing Class Mapping")
480        ):
481            pbar.set_description(f"Zero Shot Classification model {model_name}")
482            mapper = ClassMapper(dataset, model_name, run_config)
483            try:
484                stats = mapper.run_mapping(test_dataset_source, test_dataset_target)
485                # Display statistics only if mapping was successful
486                source_class_counts = stats["total_processed"]
487                logging.info("\nClassification Results for Source Dataset:")
488                for source_class, count in stats["source_class_counts"].items():
489                    percentage = (
490                        (count / source_class_counts) * 100
491                        if source_class_counts > 0
492                        else 0
493                    )
494                    logging.info(
495                        f"{source_class}: {count} samples processed ({percentage:.1f}%)"
496                    )
497
498                # Display statistics for tags added to Target Dataset
499                logging.info("\nTag Addition Results (Target Dataset Tags):")
500                logging.info(f"Total new tags added: {stats['changes_made']}")
501                for target_class, tag_count in stats["tags_added_per_category"].items():
502                    logging.info(f"{target_class} tags added: {tag_count}")
503
504            except Exception as e:
505                logging.error(f"Error during mapping with model {model_name}: {e}")
506
507    except Exception as e:
508        logging.error(f"Error in class_mapping workflow: {e}")
509        wandb_exit_code = 1
510    finally:
511        wandb_close(wandb_exit_code)
512
513    return True
514
515
516def cleanup_memory(do_extensive_cleanup=False):
517    """Clean up memory after workflow execution. 'do_extensive_cleanup' recommended for multiple training sessions in a row."""
518    logging.info("Starting memory cleanup")
519    # Clear CUDA cache
520    if torch.cuda.is_available():
521        torch.cuda.empty_cache()
522
523    # Force garbage collection
524    gc.collect()
525
526    if do_extensive_cleanup:
527
528        # Clear any leftover tensors
529        n_deleted_torch_objects = 0
530        for obj in tqdm(
531            gc.get_objects(), desc="Deleting objects from Python Garbage Collector"
532        ):
533            try:
534                if torch.is_tensor(obj):
535                    del obj
536                    n_deleted_torch_objects += 1
537            except:
538                pass
539
540        logging.info(f"Deleted {n_deleted_torch_objects} torch objects")
541
542        # Final garbage collection
543        gc.collect()
544
545
546class WorkflowExecutor:
547    """Orchestrates the execution of multiple data processing workflows in sequence."""
548
549    def __init__(
550        self,
551        workflows: List[str],
552        selected_dataset: str,
553        dataset: fo.Dataset,
554        dataset_info: Dict,
555    ):
556        """Initializes with specified workflows, dataset selection, and dataset metadata."""
557        self.workflows = workflows
558        self.selected_dataset = selected_dataset
559        self.dataset = dataset
560        self.dataset_info = dataset_info
561
562    def execute(self) -> bool:
563        """Execute all configured workflows in sequence and handle errors."""
564        if len(self.workflows) == 0:
565            logging.error("No workflows selected.")
566            return False
567
568        logging.info(f"Selected workflows: {self.workflows}")
569        for workflow in self.workflows:
570            logging.info(
571                f"Running workflow {workflow} for dataset {self.selected_dataset}"
572            )
573            try:
574                if workflow == "aws_download":
575                    parameter_group = "mcity"
576                    parameters = WORKFLOWS["aws_download"].get(parameter_group, None)
577                    if parameter_group == "mcity":
578                        dataset, dataset_name, _ = workflow_aws_download(parameters)
579                    else:
580                        logging.error(
581                            f"The parameter group {parameter_group} is not supported. As AWS are highly specific, please provide a separate set of parameters and a workflow."
582                        )
583
584                    # Select downloaded dataset for further workflows if configured
585                    if dataset is not None:
586                        if parameters["selected_dataset_overwrite"] == True:
587
588                            dataset_info = {
589                                "name": dataset_name,
590                                "v51_type": "FiftyOneDataset",
591                                "splits": [],
592                            }
593
594                            self.dataset = dataset
595                            self.dataset_info = dataset_info
596                            logging.warning(
597                                f"Overwritting selected dataset {self.selected_dataset} with {dataset_name}"
598                            )
599                            self.selected_dataset = dataset_name
600
601                elif workflow == "embedding_selection":
602                    embedding_models = WORKFLOWS["embedding_selection"][
603                        "embedding_models"
604                    ]
605
606                    for MODEL_NAME in (
607                        pbar := tqdm(embedding_models, desc="Selection by Embeddings")
608                    ):
609
610                        # Status
611                        pbar.set_description(
612                            f"Selection by embeddings with model {MODEL_NAME}."
613                        )
614
615                        # Config
616                        mode = WORKFLOWS["embedding_selection"]["mode"]
617                        parameters = WORKFLOWS["embedding_selection"]["parameters"]
618                        config = {"mode": mode, "parameters": parameters}
619
620                        # Workflow
621                        workflow_embedding_selection(
622                            self.dataset,
623                            self.dataset_info,
624                            MODEL_NAME,
625                            config,
626                        )
627
628                elif workflow == "anomaly_detection":
629
630                    # Config
631                    ano_dec_config = WORKFLOWS["anomaly_detection"]
632                    anomalib_image_models = ano_dec_config["anomalib_image_models"]
633                    eval_metrics = ano_dec_config["anomalib_eval_metrics"]
634
635                    dataset_ano_dec = None
636                    data_root = None
637                    if "train" in ano_dec_config["mode"]:
638                        try:
639                            data_preparer = AnomalyDetectionDataPreparation(
640                                self.dataset, self.selected_dataset
641                            )
642                            dataset_ano_dec = data_preparer.dataset_ano_dec
643                            data_root = data_preparer.export_root
644                        except Exception as e:
645                            logging.error(
646                                f"Error during data preparation for Anomaly Detection: {e}"
647                            )
648
649                    for MODEL_NAME in (
650                        pbar := tqdm(anomalib_image_models, desc="Anomalib")
651                    ):
652                        # Status
653                        pbar.set_description(f"Anomalib model {MODEL_NAME}")
654
655                        # Config
656                        run_config = {
657                            "model_name": MODEL_NAME,
658                            "image_size": anomalib_image_models[MODEL_NAME].get(
659                                "image_size", None
660                            ),
661                            "batch_size": anomalib_image_models[MODEL_NAME].get(
662                                "batch_size", 1
663                            ),
664                            "epochs": ano_dec_config["epochs"],
665                            "early_stop_patience": ano_dec_config[
666                                "early_stop_patience"
667                            ],
668                            "data_root": data_root,
669                            "mode": ano_dec_config["mode"],
670                        }
671
672                        # Workflow
673                        workflow_anomaly_detection(
674                            self.dataset,
675                            dataset_ano_dec,
676                            self.dataset_info,
677                            eval_metrics,
678                            run_config,
679                        )
680
681                elif workflow == "auto_labeling":
682
683                    # Config
684                    SUPPORTED_MODEL_SOURCES = [
685                        "hf_models_objectdetection",
686                        "ultralytics",
687                        "custom_codetr",
688                    ]
689
690                    # Common parameters between models
691                    config_autolabel = WORKFLOWS["auto_labeling"]
692                    mode = config_autolabel["mode"]
693                    epochs = config_autolabel["epochs"]
694                    selected_model_source = config_autolabel["model_source"]
695
696                    # Check if all selected modes are supported
697                    for model_source in selected_model_source:
698                        if model_source not in SUPPORTED_MODEL_SOURCES:
699                            logging.error(
700                                f"Selected model source {model_source} is not supported."
701                            )
702
703                    if SUPPORTED_MODEL_SOURCES[0] in selected_model_source:
704                        # Hugging Face Models
705                        # Single GPU mode (https://github.com/huggingface/transformers/issues/28740)
706                        os.environ["CUDA_VISIBLE_DEVICES"] = "0"
707                        hf_models = config_autolabel["hf_models_objectdetection"]
708
709                        # Dataset Conversion
710                        try:
711                            logging.info("Converting dataset into Hugging Face format.")
712                            pytorch_dataset = FiftyOneTorchDatasetCOCO(self.dataset)
713                            pt_to_hf_converter = TorchToHFDatasetCOCO(pytorch_dataset)
714                            hf_dataset = pt_to_hf_converter.convert()
715                        except Exception as e:
716                            logging.error(f"Error during dataset conversion: {e}")
717
718                        for MODEL_NAME in (
719                            pbar := tqdm(hf_models, desc="Auto Labeling Models")
720                        ):
721                            # Status Update
722                            pbar.set_description(
723                                f"Processing Hugging Face model {MODEL_NAME}"
724                            )
725
726                            # Config
727                            config_model = config_autolabel[
728                                "hf_models_objectdetection"
729                            ][MODEL_NAME]
730
731                            run_config = {
732                                "mode": mode,
733                                "model_name": MODEL_NAME,
734                                "v51_dataset_name": self.selected_dataset,
735                                "epochs": epochs,
736                                "early_stop_threshold": config_autolabel[
737                                    "early_stop_threshold"
738                                ],
739                                "early_stop_patience": config_autolabel[
740                                    "early_stop_patience"
741                                ],
742                                "learning_rate": config_autolabel["learning_rate"],
743                                "weight_decay": config_autolabel["weight_decay"],
744                                "max_grad_norm": config_autolabel["max_grad_norm"],
745                                "batch_size": config_model.get("batch_size", 1),
746                                "image_size": config_model.get("image_size", None),
747                                "n_worker_dataloader": config_autolabel[
748                                    "n_worker_dataloader"
749                                ],
750                                "inference_settings": config_autolabel[
751                                    "inference_settings"
752                                ],
753                            }
754
755                            # Workflow
756                            workflow_auto_labeling_hf(
757                                self.dataset,
758                                hf_dataset,
759                                run_config,
760                            )
761
762                    if SUPPORTED_MODEL_SOURCES[1] in selected_model_source:
763                        # Ultralytics Models
764                        config_ultralytics = config_autolabel["ultralytics"]
765                        models_ultralytics = config_ultralytics["models"]
766                        export_dataset_root = config_ultralytics["export_dataset_root"]
767
768                        # Export data into necessary format
769                        if "train" in mode:
770                            try:
771                                UltralyticsObjectDetection.export_data(
772                                    self.dataset,
773                                    self.dataset_info,
774                                    export_dataset_root,
775                                )
776                            except Exception as e:
777                                logging.error(
778                                    f"Error during Ultralytics dataset export: {e}"
779                                )
780
781                        for model_name in (
782                            pbar := tqdm(
783                                models_ultralytics, desc="Ultralytics training"
784                            )
785                        ):
786                            pbar.set_description(f"Ultralytics model {model_name}")
787                            run_config = {
788                                "mode": mode,
789                                "model_name": model_name,
790                                "v51_dataset_name": self.dataset_info["name"],
791                                "epochs": epochs,
792                                "patience": config_autolabel["early_stop_patience"],
793                                "batch_size": models_ultralytics[model_name][
794                                    "batch_size"
795                                ],
796                                "img_size": models_ultralytics[model_name]["img_size"],
797                                "export_dataset_root": export_dataset_root,
798                                "inference_settings": config_autolabel[
799                                    "inference_settings"
800                                ],
801                                "multi_scale": config_ultralytics["multi_scale"],
802                                "cos_lr": config_ultralytics["cos_lr"],
803                            }
804
805                            workflow_auto_labeling_ultralytics(self.dataset, run_config)
806
807                    if SUPPORTED_MODEL_SOURCES[2] in selected_model_source:
808                        # Custom Co-DETR
809                        config_codetr = config_autolabel["custom_codetr"]
810                        run_config = {
811                            "export_dataset_root": config_codetr["export_dataset_root"],
812                            "container_tool": config_codetr["container_tool"],
813                            "n_gpus": config_codetr["n_gpus"],
814                            "mode": config_autolabel["mode"],
815                            "epochs": config_autolabel["epochs"],
816                            "inference_settings": config_autolabel[
817                                "inference_settings"
818                            ],
819                            "config": None,
820                        }
821                        codetr_configs = config_codetr["configs"]
822
823                        for config in (
824                            pbar := tqdm(
825                                codetr_configs, desc="Processing Co-DETR configurations"
826                            )
827                        ):
828                            pbar.set_description(f"Co-DETR model {config}")
829                            run_config["config"] = config
830                            workflow_auto_labeling_custom_codetr(
831                                self.dataset, self.dataset_info, run_config
832                            )
833
834                elif workflow == "auto_labeling_zero_shot":
835                    config = WORKFLOWS["auto_labeling_zero_shot"]
836                    workflow_zero_shot_object_detection(
837                        self.dataset, self.dataset_info, config
838                    )
839
840                elif workflow == "ensemble_selection":
841                    # Config
842                    run_config = WORKFLOWS["ensemble_selection"]
843
844                    # Workflow
845                    workflow_ensemble_selection(
846                        self.dataset, self.dataset_info, run_config
847                    )
848
849                elif workflow == "auto_label_mask":
850                    config = WORKFLOWS["auto_label_mask"]
851                    workflow_auto_label_mask(self.dataset, self.dataset_info, config)
852
853                elif workflow == "class_mapping":
854                    # Config
855                    run_config = WORKFLOWS["class_mapping"]
856
857                    # Workflow
858                    workflow_class_mapping(
859                        self.dataset,
860                        self.dataset_info,
861                        run_config,
862                        test_dataset_source=None,
863                        test_dataset_target=None,
864                    )
865
866                else:
867                    logging.error(
868                        f"Workflow {workflow} not found. Check available workflows in config.py."
869                    )
870                    return False
871
872                cleanup_memory()  # Clean after each workflow
873                logging.info(f"Completed workflow {workflow} and cleaned up memory")
874
875            except Exception as e:
876                logging.error(f"Workflow {workflow}: An error occurred: {e}")
877                wandb_close(exit_code=1)
878                cleanup_memory()  # Clean up even after failure
879
880        return True
881
882
883def main():
884    """Executes the data processing workflow, loads dataset, and launches Voxel51 visualization interface."""
885    time_start = time.time()
886    configure_logging()
887
888    # Signal handler for CTRL + C
889    signal.signal(signal.SIGINT, signal_handler)
890
891    # Execute workflows
892    dataset, dataset_info = load_dataset(SELECTED_DATASET)
893
894    executor = WorkflowExecutor(
895        SELECTED_WORKFLOW, SELECTED_DATASET["name"], dataset, dataset_info
896    )
897    executor.execute()
898
899    # Launch V51 session
900    dataset.reload()
901    dataset.save()
902    arrange_fields_in_groups(dataset)
903    logging.info(f"Launching Voxel51 session for dataset {dataset_info['name']}.")
904
905    # Dataset stats
906    logging.debug(dataset)
907    logging.debug(dataset.stats(include_media=True))
908
909    # V51 UI launch
910    session = fo.launch_app(
911        dataset, address=V51_ADDRESS, port=V51_PORT, remote=V51_REMOTE
912    )
913
914    time_stop = time.time()
915    logging.info(f"Elapsed time: {time_stop - time_start:.2f} seconds")
916
917
918if __name__ == "__main__":
919    cleanup_memory()
920    main()