feat: Add model yolo-nas

build_static_site.py

@@ -2,7 +2,7 @@
 from pathlib import Path
 
 # Temporarily disabled from appearing on the board, e.g. if there's still some issues
-BLACKLIST = ["yolov9"]
+BLACKLIST = ["yolov9", "yolo-nas"]
 
 results_list = []
 for model_dir in Path("models/object_detection").iterdir():

models/object_detection/yolo-nas/requirements.txt

@@ -0,0 +1,4 @@
+super_gradients==3.7.1
+supervision>=0.24.0rc1
+torch==2.2.2
+tqdm

models/object_detection/yolo-nas/results_yolo_nas_l.json

@@ -0,0 +1,38 @@
+{
+    "metadata": {
+        "model": "YOLO-NAS L",
+        "license": "Apache-2.0",
+        "param_count": 66976392,
+        "run_date": "2024-09-12T15:24:51.884717+00:00"
+    },
+    "map50_95": 0.48298094432764754,
+    "map50": 0.6395166289719741,
+    "map75": 0.5280721995716876,
+    "small_objects": {
+        "map50_95": 0.20388303052343654,
+        "map50": 0.3176009687794093,
+        "map75": 0.22339241590855352
+    },
+    "medium_objects": {
+        "map50_95": 0.44532990793587546,
+        "map50": 0.6110120117547352,
+        "map75": 0.5011779178048194
+    },
+    "large_objects": {
+        "map50_95": 0.6240832744737295,
+        "map50": 0.751358076433659,
+        "map75": 0.6747371906034554
+    },
+    "iou_thresholds": [
+        0.5,
+        0.55,
+        0.6,
+        0.65,
+        0.7,
+        0.75,
+        0.8,
+        0.85,
+        0.8999999999999999,
+        0.95
+    ]
+}

models/object_detection/yolo-nas/results_yolo_nas_m.json

@@ -0,0 +1,38 @@
+{
+    "metadata": {
+        "model": "YOLO-NAS M",
+        "license": "Apache-2.0",
+        "param_count": 51182658,
+        "run_date": "2024-09-12T15:03:37.531535+00:00"
+    },
+    "map50_95": 0.4766481832173873,
+    "map50": 0.632842568125929,
+    "map75": 0.5229364873051732,
+    "small_objects": {
+        "map50_95": 0.198911024541053,
+        "map50": 0.31075430842955604,
+        "map75": 0.21979759840215746
+    },
+    "medium_objects": {
+        "map50_95": 0.43936196228345464,
+        "map50": 0.6034758757860981,
+        "map75": 0.49902495607127173
+    },
+    "large_objects": {
+        "map50_95": 0.6237642850164452,
+        "map50": 0.7525691527759143,
+        "map75": 0.6784723077498078
+    },
+    "iou_thresholds": [
+        0.5,
+        0.55,
+        0.6,
+        0.65,
+        0.7,
+        0.75,
+        0.8,
+        0.85,
+        0.8999999999999999,
+        0.95
+    ]
+}

models/object_detection/yolo-nas/results_yolo_nas_s.json

@@ -0,0 +1,38 @@
+{
+    "metadata": {
+        "model": "YOLO-NAS S",
+        "license": "Apache-2.0",
+        "param_count": 19053888,
+        "run_date": "2024-09-12T14:45:02.840446+00:00"
+    },
+    "map50_95": 0.4369354384266927,
+    "map50": 0.5950857479290619,
+    "map75": 0.47826919036030613,
+    "small_objects": {
+        "map50_95": 0.1682506254047077,
+        "map50": 0.2634323893258659,
+        "map75": 0.18181972444412323
+    },
+    "medium_objects": {
+        "map50_95": 0.39388058360539213,
+        "map50": 0.5587017365344934,
+        "map75": 0.43967107626652585
+    },
+    "large_objects": {
+        "map50_95": 0.5820850233048388,
+        "map50": 0.7216236849916268,
+        "map75": 0.6364020637637353
+    },
+    "iou_thresholds": [
+        0.5,
+        0.55,
+        0.6,
+        0.65,
+        0.7,
+        0.75,
+        0.8,
+        0.85,
+        0.8999999999999999,
+        0.95
+    ]
+}

models/object_detection/yolo-nas/run.py

@@ -0,0 +1,115 @@
+import argparse
+import sys
+from pathlib import Path
+from typing import List, Optional
+
+import super_gradients
+import super_gradients.training
+import super_gradients.training.models
+import supervision as sv
+import torch
+from tqdm import tqdm
+
+sys.path.append(str(Path(__file__).resolve().parent.parent))
+from utils import (
+    load_detections_dataset,
+    result_json_already_exists,
+    write_result_json,
+)
+
+LICENSE = "Apache-2.0"
+MODEL_DICT = {
+    "yolo_nas_s": {
+        "name": "YOLO-NAS S",
+    },
+    "yolo_nas_m": {
+        "name": "YOLO-NAS M",
+    },
+    "yolo_nas_l": {
+        "name": "YOLO-NAS L",
+    },
+}
+DATASET_DIR = "../../../data/coco-val-2017"
+CONFIDENCE_THRESHOLD = 0.001
+
+
+def run_on_image(model, image) -> sv.Detections:
+    model_params = dict(
+        conf=0.01,
+        iou=0.7,
+        nms_top_k=1000,
+        max_predictions=300,
+    )
+    result = model.predict(image, **model_params)
+    detections = sv.Detections.from_yolo_nas(result)
+    detections = detections[detections.confidence > CONFIDENCE_THRESHOLD]
+    return detections
+
+
+def run(
+    model_ids: List[str],
+    skip_if_result_exists=False,
+    dataset: Optional[sv.DetectionDataset] = None,
+) -> None:
+    """
+    Run the evaluation for the given models and dataset.
+
+    Arguments:
+        model_ids: List of model ids to evaluate. Evaluate all models if None.
+        skip_if_result_exists: If True, skip the evaluation if the result json already exists.
+        dataset: If provided, use this dataset for evaluation. Otherwise, load the dataset from the default directory.
+    """  # noqa: E501 // docs
+    if not model_ids:
+        model_ids = list(MODEL_DICT.keys())
+
+    for model_id in model_ids:
+        print(f"\nEvaluating model: {model_id}")
+        model_values = MODEL_DICT[model_id]
+
+        if skip_if_result_exists and result_json_already_exists(model_id):
+            print(f"Skipping {model_id}. Result already exists!")
+            continue
+
+        if dataset is None:
+            dataset = load_detections_dataset(DATASET_DIR)
+
+        model = super_gradients.training.models.get(model_id, pretrained_weights="coco")
+        if torch.cuda.is_available():
+            model = model.cuda()
+
+        predictions = []
+        targets = []
+        print("Evaluating...")
+        for _, image, target_detections in tqdm(dataset, total=len(dataset)):
+            # Run model
+            detections = run_on_image(model, image)
+            predictions.append(detections)
+            targets.append(target_detections)
+
+        mAP_metric = sv.metrics.MeanAveragePrecision()
+        mAP_result = mAP_metric.update(predictions, targets).compute()
+
+        write_result_json(
+            model_id=model_id,
+            model_name=model_values["name"],
+            model=model,
+            mAP_result=mAP_result,
+            license_name=LICENSE,
+        )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "model_ids",
+        nargs="*",
+        help="Model ids to evaluate. If not provided, evaluate all models.",
+    )
+    parser.add_argument(
+        "--skip_if_result_exists",
+        help="If specified, skip the evaluation if the result json already exists.",
+        action="store_true",
+    )
+    args = parser.parse_args()
+
+    run(args.model_ids, args.skip_if_result_exists)