Merge pull request #307 from mcity/roboflowmodels

Added Roboflow's RFDETR models to the Auto Labeling Workflow

Author: rajanikant-patnaik

config/config.py

@@ -2,7 +2,7 @@
 import psutil
 
 #: Select workflow list from 'WORKFLOWS = {...}' dictionary
-SELECTED_WORKFLOW = ["data_ingest"]  # Choose from WORKFLOWS keys
+SELECTED_WORKFLOW = ["auto_labeling"]  # Choose from WORKFLOWS keys
 
 #: Select dataset from config/datasets.yaml
 SELECTED_DATASET = {
@@ -66,8 +66,9 @@
         "mode": ["train","inference"],  # "train" and "inference" supported
         "model_source": [
             # "hf_models_objectdetection",
-            "ultralytics",
+            #"ultralytics",
             # "custom_codetr",
+            "roboflow",
         ],
         "n_worker_dataloader": 8,
         "epochs": 1,
@@ -129,6 +130,24 @@
             "n_gpus": "1",
             "container_tool": "docker",
         },
+        "roboflow": {  # Roboflow RF-DETR configuration
+            "export_dataset_root": "output/datasets/roboflow_data/",
+            "configs": [
+                "rfdetr_nano",
+                "rfdetr_small",
+                #"rfdetr_medium",
+                #"rfdetr_large",
+            ],
+            # RF-DETR specific parameters only
+            "batch_size": 4,                     # Override default batch size
+            "grad_accum_steps": 4,                # Gradient accumulation steps
+            "lr_encoder": None,                   # Encoder-specific learning rate (optional)
+            "resolution": None,                   # Image resolution, must be divisible by 56 (optional)
+            "use_ema": True,                      # Exponential moving average
+            "gradient_checkpointing": False,      # Memory optimization
+            "early_stopping_min_delta": 0.001,    # Minimum improvement for early stopping
+            "early_stopping_use_ema": True,       # Use EMA model for early stopping
+        },
         "ultralytics": {
             "export_dataset_root": "output/datasets/ultralytics_data/",
             "multi_scale": False,

main.py

@@ -37,6 +37,7 @@
 from workflows.anomaly_detection import Anodec
 from workflows.auto_labeling import (
     CustomCoDETRObjectDetection,
+    CustomRFDETRObjectDetection,
     HuggingFaceObjectDetection,
     UltralyticsObjectDetection,
     ZeroShotObjectDetection,
@@ -686,6 +687,7 @@ def execute(self) -> bool:
                         "hf_models_objectdetection",
                         "ultralytics",
                         "custom_codetr",
+                        "roboflow",
                     ]
 
                     # Common parameters between models
@@ -832,6 +834,79 @@ def execute(self) -> bool:
                                 self.dataset, self.dataset_info, run_config
                             )
 
+                    if SUPPORTED_MODEL_SOURCES[3] in selected_model_source:
+
+                        config_rfdetr = config_autolabel["roboflow"]
+
+                        # Shared config parameters
+                        shared_config = {
+                            "epochs": config_autolabel["epochs"],
+                            "learning_rate": config_autolabel["learning_rate"],
+                            "weight_decay": config_autolabel["weight_decay"],
+                            "early_stop_patience": config_autolabel["early_stop_patience"],
+                            "early_stop_threshold": config_autolabel["early_stop_threshold"],
+                        }
+
+                        run_config = {
+                            "export_dataset_root": config_rfdetr["export_dataset_root"],
+                            "mode": config_autolabel["mode"],
+                            "inference_settings": config_autolabel["inference_settings"],
+                            "config": None,
+                            # RF-DETR specific parameters
+                            "batch_size": config_rfdetr["batch_size"],
+                            "grad_accum_steps": config_rfdetr["grad_accum_steps"],
+                            "lr_encoder": config_rfdetr["lr_encoder"],
+                            "resolution": config_rfdetr["resolution"],
+                            "use_ema": config_rfdetr["use_ema"],
+                            "gradient_checkpointing": config_rfdetr["gradient_checkpointing"],
+                            "early_stopping_min_delta": config_rfdetr["early_stopping_min_delta"],
+                            "early_stopping_use_ema": config_rfdetr["early_stopping_use_ema"],
+                        }
+
+                        rfdetr_configs = config_rfdetr["configs"]
+
+                        for config in (
+                            pbar := tqdm(rfdetr_configs, desc="Processing RF-DETR configurations")
+                        ):
+                            pbar.set_description(f"RF-DETR model {config}")
+                            run_config["config"] = config
+
+                            try:
+                                wandb_exit_code = 0
+                                wandb_run = wandb_init(
+                                    run_name=config,
+                                    project_name="RF-DETR Auto Labeling",
+                                    dataset_name=self.dataset_info["name"],
+                                    config=run_config,
+                                    wandb_activate=True,
+                                )
+
+                                detector = CustomRFDETRObjectDetection(
+                                    self.dataset, self.dataset_info, run_config
+                                )
+
+                                # Convert data to RF-DETR format
+                                detector.convert_data()
+
+                                # Training
+                                if "train" in mode:
+                                    logging.info(f"Training RF-DETR model: {config}")
+                                    detector.train(run_config, shared_config)
+
+                                # Inference
+                                if "inference" in mode:
+                                    logging.info(f"Running inference for RF-DETR model: {config}")
+                                    detector.inference(
+                                        inference_settings=config_autolabel["inference_settings"]
+                                    )
+
+                            except Exception as e:
+                                logging.error(f"Error during RF-DETR workflow with {config}: {e}")
+                                wandb_exit_code = 1
+                            finally:
+                                wandb_close(wandb_exit_code)
+
+
                 elif workflow == "auto_labeling_zero_shot":
                     config = WORKFLOWS["auto_labeling_zero_shot"]
                     workflow_zero_shot_object_detection(

requirements.txt

@@ -324,6 +324,7 @@ requirements-parser==0.11.0
 retrying==1.3.4
 rfc3339-validator==0.1.4
 rfc3986-validator==0.1.1
+rfdetr==1.3.0
 rich==13.9.1
 rich-argparse==1.5.2
 rpds-py==0.20.0