From 4a3eb772f98699acabb6e098652df2c46516f0da Mon Sep 17 00:00:00 2001 From: agosl Date: Sun, 13 Aug 2023 14:05:15 -0300 Subject: [PATCH 1/3] yolo nas demo --- demos/yolo_nas/Dockerfile | 20 ++++ demos/yolo_nas/requirements.txt | 1 + demos/yolo_nas/run_gpu.sh | 8 ++ demos/yolo_nas/src/demo.py | 164 ++++++++++++++++++++++++++++++++ 4 files changed, 193 insertions(+) create mode 100644 demos/yolo_nas/Dockerfile create mode 100644 demos/yolo_nas/requirements.txt create mode 100644 demos/yolo_nas/run_gpu.sh create mode 100644 demos/yolo_nas/src/demo.py diff --git a/demos/yolo_nas/Dockerfile b/demos/yolo_nas/Dockerfile new file mode 100644 index 00000000..1a1a0631 --- /dev/null +++ b/demos/yolo_nas/Dockerfile @@ -0,0 +1,20 @@ +# https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch/tags +FROM nvcr.io/nvidia/pytorch:22.08-py3 + +RUN apt-get update && \ + apt-get install -y libgl1 && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* + +RUN pip install --upgrade pip + +RUN pip install -r https://raw.githubusercontent.com/Deci-AI/super-gradients/master/requirements.txt && \ + # For some reason it doesn't work with latest version of OpenCV \ --> AttributeError: partially initialized module 'cv2' has no attribute '_registerMatType' (most likely due to a circular import) + pip install opencv-python==4.5.5.64 + +RUN pip install git+https://github.com/tryolabs/norfair.git@master#egg=norfair + +RUN pip install super-gradients + + +WORKDIR /demo/src/ \ No newline at end of file diff --git a/demos/yolo_nas/requirements.txt b/demos/yolo_nas/requirements.txt new file mode 100644 index 00000000..6a95e21b --- /dev/null +++ b/demos/yolo_nas/requirements.txt @@ -0,0 +1 @@ +super-gradient==3.1.1 \ No newline at end of file diff --git a/demos/yolo_nas/run_gpu.sh b/demos/yolo_nas/run_gpu.sh new file mode 100644 index 00000000..00d7e2d1 --- /dev/null +++ b/demos/yolo_nas/run_gpu.sh @@ -0,0 +1,8 @@ +#!/usr/bin/env -S bash -e +docker build . -t norfair-yolonas +docker run -it --rm \ + --gpus all \ + --shm-size=1gb \ + -v `realpath .`:/demo \ + norfair-yolonas \ + bash diff --git a/demos/yolo_nas/src/demo.py b/demos/yolo_nas/src/demo.py new file mode 100644 index 00000000..5750408e --- /dev/null +++ b/demos/yolo_nas/src/demo.py @@ -0,0 +1,164 @@ +import argparse +from typing import List, Optional, Union + +import numpy as np +import torch + +import norfair +from norfair import Detection, Paths, Tracker, Video + +DISTANCE_THRESHOLD_BBOX: float = 0.7 +DISTANCE_THRESHOLD_CENTROID: int = 30 +MAX_DISTANCE: int = 10000 + + +import super_gradients + + +class YOLO_NAS: + def __init__(self, model_name: str, device: Optional[str] = None): + if device is not None and "cuda" in device and not torch.cuda.is_available(): + raise Exception( + "Selected device='cuda', but cuda is not available to Pytorch." + ) + # automatically set device if its None + elif device is None: + device = "cuda:0" if torch.cuda.is_available() else "cpu" + + # load model + else: + self.model = super_gradients.training.models.get( + "yolo_nas_l", pretrained_weights="coco" + ).cuda() + + def __call__( + self, + img: Union[str, np.ndarray], + conf_threshold: float = 0.35, + iou_threshold: float = 0.45, + image_size: int = 720, + classes: Optional[List[int]] = None, + ) -> torch.tensor: + + if classes is not None: + self.model.classes = classes + + detections = self.model.predict(img, iou_threshold, conf_threshold) + return detections + + +def yolo_detections_to_norfair_detections( + yolo_detections: torch.tensor, track_points: str = "centroid" # bbox or centroid +) -> List[Detection]: + """convert detections_as_xywh to norfair detections""" + norfair_detections: List[Detection] = [] + + if track_points == "centroid": + detections_as_xywh = yolo_detections.xywh[0] + for detection_as_xywh in detections_as_xywh: + centroid = np.array( + [detection_as_xywh[0].item(), detection_as_xywh[1].item()] + ) + scores = np.array([detection_as_xywh[4].item()]) + norfair_detections.append( + Detection( + points=centroid, + scores=scores, + label=int(detection_as_xywh[-1].item()), + ) + ) + + elif track_points == "bbox": + + ## yolo_nas detections + detections_as_xyxy = yolo_detections[0] + class_names = detections_as_xyxy.class_names + labels = detections_as_xyxy.prediction.labels + confidence = detections_as_xyxy.prediction.confidence + bboxes = detections_as_xyxy.prediction.bboxes_xyxy + + for i, (label, conf, bbox_yolo) in enumerate(zip(labels, confidence, bboxes)): + bbox = np.array( + [ + [bbox_yolo[0], bbox_yolo[1]], + [bbox_yolo[2], bbox_yolo[3]], + ] + ) + + scores = np.array([conf, conf]) + norfair_detections.append( + Detection(points=bbox, scores=scores, label=class_names[int(label)]) + ) + + return norfair_detections + + +parser = argparse.ArgumentParser(description="Track objects in a video.") +parser.add_argument("files", type=str, nargs="+", help="Video files to process") +parser.add_argument( + "--model-name", type=str, default="yolovnas", help="YOLOv5 model name" +) +parser.add_argument( + "--img-size", type=int, default="720", help="YOLO_nas inference size (pixels)" +) +parser.add_argument( + "--conf-threshold", + type=float, + default="0.25", + help="YOLOv5 object confidence threshold", +) +parser.add_argument( + "--iou-threshold", type=float, default="0.45", help="YOLOv5 IOU threshold for NMS" +) +parser.add_argument( + "--classes", + nargs="+", + type=int, + help="Filter by class: --classes 0, or --classes 0 2 3", +) +parser.add_argument( + "--device", type=str, default="cuda", help="Inference device: 'cpu' or 'cuda'" +) +parser.add_argument( + "--track-points", + type=str, + default="bbox", + help="Track points: 'centroid' or 'bbox'", +) +args = parser.parse_args() + +model = YOLO_NAS(args.model_name, device=args.device) + +for input_path in args.files: + video = Video(input_path=input_path) + distance_function = "iou" if args.track_points == "bbox" else "euclidean" + distance_threshold = ( + DISTANCE_THRESHOLD_BBOX + if args.track_points == "bbox" + else DISTANCE_THRESHOLD_CENTROID + ) + + tracker = Tracker( + distance_function=distance_function, + distance_threshold=distance_threshold, + ) + + for frame in video: + yolo_detections = model( + frame, + conf_threshold=args.conf_threshold, + iou_threshold=args.iou_threshold, + image_size=args.img_size, + ) + + detections = yolo_detections_to_norfair_detections( + yolo_detections, track_points=args.track_points + ) + tracked_objects = tracker.update(detections=detections) + if args.track_points == "centroid": + norfair.draw_points(frame, detections) + norfair.draw_tracked_objects(frame, tracked_objects) + elif args.track_points == "bbox": + norfair.draw_boxes(frame, detections) + norfair.draw_boxes(frame, tracked_objects, draw_ids=True) + video.write(frame) From e1e80ac9d41e7b21b4c7de4d00dd30c399d64c33 Mon Sep 17 00:00:00 2001 From: agosl Date: Mon, 28 Aug 2023 08:27:27 -0300 Subject: [PATCH 2/3] include readme --- demos/yolo_nas/Dockerfile | 4 ++-- demos/yolo_nas/README.md | 19 +++++++++++++++++++ demos/yolo_nas/run_gpu.sh | 0 demos/yolo_nas/src/demo.py | 11 ++++------- 4 files changed, 25 insertions(+), 9 deletions(-) create mode 100644 demos/yolo_nas/README.md mode change 100644 => 100755 demos/yolo_nas/run_gpu.sh diff --git a/demos/yolo_nas/Dockerfile b/demos/yolo_nas/Dockerfile index 1a1a0631..1add14a1 100644 --- a/demos/yolo_nas/Dockerfile +++ b/demos/yolo_nas/Dockerfile @@ -14,7 +14,7 @@ RUN pip install -r https://raw.githubusercontent.com/Deci-AI/super-gradients/mas RUN pip install git+https://github.com/tryolabs/norfair.git@master#egg=norfair -RUN pip install super-gradients - +COPY requirements.txt requirements.txt +RUN pip install -r requirements.txt WORKDIR /demo/src/ \ No newline at end of file diff --git a/demos/yolo_nas/README.md b/demos/yolo_nas/README.md new file mode 100644 index 00000000..dbbd206c --- /dev/null +++ b/demos/yolo_nas/README.md @@ -0,0 +1,19 @@ +# YOLO nas example + +Simplest possible example of tracking. Based on [YOLO-NAS-L](https://github.com/Deci-AI/super-gradients/blob/master/YOLONAS.md). + +## Instructions + +1. Build and run the Docker container with `./run_gpu.sh`. +2. Copy a video to the `src` folder. +3. Within the container, run with the default parameters: + + ```bash + python demo.py