From 8f1c3f3d1e7d2a30800d65c087618a5618a78df0 Mon Sep 17 00:00:00 2001 From: Muhammad Rizwan Munawar Date: Fri, 24 Nov 2023 17:45:35 +0500 Subject: [PATCH] Add `--classes` arg in YOLOv8 Region Counter + optimize (#6568) Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- examples/YOLOv8-Region-Counter/readme.md | 6 ++++ .../yolov8_region_counter.py | 29 +++++++++---------- 2 files changed, 20 insertions(+), 15 deletions(-) diff --git a/examples/YOLOv8-Region-Counter/readme.md b/examples/YOLOv8-Region-Counter/readme.md index da303f1946..9c0ad168be 100644 --- a/examples/YOLOv8-Region-Counter/readme.md +++ b/examples/YOLOv8-Region-Counter/readme.md @@ -48,6 +48,9 @@ python yolov8_region_counter.py --source "path/to/video.mp4" --save-img --view-i # If you want to change model file python yolov8_region_counter.py --source "path/to/video.mp4" --save-img --weights "path/to/model.pt" +# If you want to detect specific class (first class and third class) +python yolov8_region_counter.py --source "path/to/video.mp4" --classes 0 2 --weights "path/to/model.pt" + # If you dont want to save results python yolov8_region_counter.py --source "path/to/video.mp4" --view-img ``` @@ -58,6 +61,7 @@ python yolov8_region_counter.py --source "path/to/video.mp4" --view-img - `--device`: Specifies the device `cpu` or `0` - `--save-img`: Flag to save the detection results as images. - `--weights`: Specifies a different YOLOv8 model file (e.g., `yolov8n.pt`, `yolov8s.pt`, `yolov8m.pt`, `yolov8l.pt`, `yolov8x.pt`). +- `--classes`: Specifies the class to be detected - `--line-thickness`: Specifies the bounding box thickness - `--region-thickness`: Specifies the region boxes thickness - `--track-thickness`: Specifies the track line thickness @@ -73,6 +77,8 @@ Region counting is a computational method utilized to ascertain the quantity of The Region Counter offers the capability to create regions in various formats, such as polygons and rectangles. You have the flexibility to modify region attributes, including coordinates, colors, and other details, as demonstrated in the following code: ```python +from shapely.geometry import Polygon + counting_regions = [ { "name": "YOLOv8 Polygon Region", diff --git a/examples/YOLOv8-Region-Counter/yolov8_region_counter.py b/examples/YOLOv8-Region-Counter/yolov8_region_counter.py index a3e52ef34f..5379fd3bb7 100644 --- a/examples/YOLOv8-Region-Counter/yolov8_region_counter.py +++ b/examples/YOLOv8-Region-Counter/yolov8_region_counter.py @@ -69,6 +69,7 @@ def run( view_img=False, save_img=False, exist_ok=False, + classes=None, line_thickness=2, track_thickness=2, region_thickness=2, @@ -87,6 +88,7 @@ def run( view_img (bool): Show results. save_img (bool): Save results. exist_ok (bool): Overwrite existing files. + classes (list): classes to detect and track line_thickness (int): Bounding box thickness. track_thickness (int): Tracking line thickness region_thickness (int): Region thickness. @@ -101,6 +103,9 @@ def run( model = YOLO(f'{weights}') model.to('cuda') if device == '0' else model.to('cpu') + # Extract classes names + names = model.model.names + # Video setup videocapture = cv2.VideoCapture(source) frame_width, frame_height = int(videocapture.get(3)), int(videocapture.get(4)) @@ -119,36 +124,29 @@ def run( vid_frame_count += 1 # Extract the results - results = model.track(frame, persist=True) + results = model.track(frame, persist=True, classes=classes) if results[0].boxes.id is not None: - boxes = results[0].boxes.xywh.cpu() + boxes = results[0].boxes.xyxy.cpu() track_ids = results[0].boxes.id.int().cpu().tolist() clss = results[0].boxes.cls.cpu().tolist() - names = results[0].names annotator = Annotator(frame, line_width=line_thickness, example=str(names)) for box, track_id, cls in zip(boxes, track_ids, clss): - x, y, w, h = box - label = str(names[cls]) - xyxy = (x - w / 2), (y - h / 2), (x + w / 2), (y + h / 2) - - # Bounding box plot - bbox_color = colors(cls, True) - annotator.box_label(xyxy, label, color=bbox_color) + annotator.box_label(box, str(names[cls]), color=colors(cls, True)) + bbox_center = (box[0] + box[2]) / 2, (box[1] + box[3]) / 2 # Bbox center - # Tracking Lines plot - track = track_history[track_id] - track.append((float(x), float(y))) + track = track_history[track_id] # Tracking Lines plot + track.append((float(bbox_center[0]), float(bbox_center[1]))) if len(track) > 30: track.pop(0) points = np.hstack(track).astype(np.int32).reshape((-1, 1, 2)) - cv2.polylines(frame, [points], isClosed=False, color=bbox_color, thickness=track_thickness) + cv2.polylines(frame, [points], isClosed=False, color=colors(cls, True), thickness=track_thickness) # Check if detection inside region for region in counting_regions: - if region['polygon'].contains(Point((x, y))): + if region['polygon'].contains(Point((bbox_center[0], bbox_center[1]))): region['counts'] += 1 # Draw regions (Polygons/Rectangles) @@ -202,6 +200,7 @@ def parse_opt(): parser.add_argument('--view-img', action='store_true', help='show results') parser.add_argument('--save-img', action='store_true', help='save results') parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment') + parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --classes 0, or --classes 0 2 3') parser.add_argument('--line-thickness', type=int, default=2, help='bounding box thickness') parser.add_argument('--track-thickness', type=int, default=2, help='Tracking line thickness') parser.add_argument('--region-thickness', type=int, default=4, help='Region thickness')