diff --git a/build_static_site.py b/build_static_site.py index 98917bf..b34d623 100644 --- a/build_static_site.py +++ b/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(): diff --git a/models/object_detection/yolo-nas/requirements.txt b/models/object_detection/yolo-nas/requirements.txt new file mode 100644 index 0000000..0d37b7e --- /dev/null +++ b/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 diff --git a/models/object_detection/yolo-nas/results_yolo_nas_l.json b/models/object_detection/yolo-nas/results_yolo_nas_l.json new file mode 100644 index 0000000..4f4f643 --- /dev/null +++ b/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 + ] +} diff --git a/models/object_detection/yolo-nas/results_yolo_nas_m.json b/models/object_detection/yolo-nas/results_yolo_nas_m.json new file mode 100644 index 0000000..2dee79b --- /dev/null +++ b/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 + ] +} diff --git a/models/object_detection/yolo-nas/results_yolo_nas_s.json b/models/object_detection/yolo-nas/results_yolo_nas_s.json new file mode 100644 index 0000000..0b17f10 --- /dev/null +++ b/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 + ] +} diff --git a/models/object_detection/yolo-nas/run.py b/models/object_detection/yolo-nas/run.py new file mode 100644 index 0000000..4bd9e5f --- /dev/null +++ b/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)