ultralytics 8.0.183 RayTune and yolo checks fixes (#5002)
Co-authored-by: Kapil Raj <103250862+raj-kapil@users.noreply.github.com> Co-authored-by: Muhammad Rizwan Munawar <62513924+RizwanMunawar@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Glenn Jocher
GitHub
parent
7f78fad8ba
commit
3223e71fea
7 changed files with 124 additions and 64 deletions
|
|
@ -4,34 +4,37 @@ from pathlib import Path
|
|||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from shapely.geometry import Polygon
|
||||
from shapely.geometry.point import Point
|
||||
|
||||
from ultralytics import YOLO
|
||||
|
||||
track_history = defaultdict(lambda: [])
|
||||
|
||||
from ultralytics.utils.files import increment_path
|
||||
from ultralytics.utils.plotting import Annotator, colors
|
||||
|
||||
# Region utils
|
||||
track_history = defaultdict(lambda: [])
|
||||
|
||||
current_region = None
|
||||
counting_regions = [{
|
||||
'name': 'YOLOv8 Region A',
|
||||
'roi': (50, 100, 240, 300),
|
||||
'counts': 0,
|
||||
'dragging': False,
|
||||
'region_color': (0, 255, 0)}, {
|
||||
'name': 'YOLOv8 Region B',
|
||||
'roi': (200, 250, 240, 300),
|
||||
counting_regions = [
|
||||
{
|
||||
'name': 'YOLOv8 Polygon Region',
|
||||
'polygon': Polygon([(50, 80), (250, 20), (450, 80), (400, 350), (100, 350)]), # Polygon points
|
||||
'counts': 0,
|
||||
'dragging': False,
|
||||
'region_color': (255, 144, 31)}]
|
||||
'region_color': (255, 42, 4), # BGR Value
|
||||
'text_color': (255, 255, 255) # Region Text Color
|
||||
},
|
||||
{
|
||||
'name': 'YOLOv8 Rectangle Region',
|
||||
'polygon': Polygon([(200, 250), (440, 250), (440, 550), (200, 550)]), # Polygon points
|
||||
'counts': 0,
|
||||
'dragging': False,
|
||||
'region_color': (37, 255, 225), # BGR Value
|
||||
'text_color': (0, 0, 0), # Region Text Color
|
||||
}, ]
|
||||
|
||||
|
||||
def is_inside_roi(box, roi):
|
||||
"""Compare bbox with region box."""
|
||||
x, y, _, _ = box
|
||||
roi_x, roi_y, roi_w, roi_h = roi
|
||||
return roi_x < x < roi_x + roi_w and roi_y < y < roi_y + roi_h
|
||||
def is_inside_polygon(point, polygon):
|
||||
return polygon.contains(Point(point))
|
||||
|
||||
|
||||
def mouse_callback(event, x, y, flags, param):
|
||||
|
|
@ -41,18 +44,21 @@ def mouse_callback(event, x, y, flags, param):
|
|||
# Mouse left button down event
|
||||
if event == cv2.EVENT_LBUTTONDOWN:
|
||||
for region in counting_regions:
|
||||
roi_x, roi_y, roi_w, roi_h = region['roi']
|
||||
if roi_x < x < roi_x + roi_w and roi_y < y < roi_y + roi_h:
|
||||
if is_inside_polygon((x, y), region['polygon']):
|
||||
current_region = region
|
||||
current_region['dragging'] = True
|
||||
current_region['offset_x'] = x - roi_x
|
||||
current_region['offset_y'] = y - roi_y
|
||||
current_region['offset_x'] = x
|
||||
current_region['offset_y'] = y
|
||||
|
||||
# Mouse move event
|
||||
elif event == cv2.EVENT_MOUSEMOVE:
|
||||
if current_region is not None and current_region['dragging']:
|
||||
current_region['roi'] = (x - current_region['offset_x'], y - current_region['offset_y'],
|
||||
current_region['roi'][2], current_region['roi'][3])
|
||||
dx = x - current_region['offset_x']
|
||||
dy = y - current_region['offset_y']
|
||||
current_region['polygon'] = Polygon([
|
||||
(p[0] + dx, p[1] + dy) for p in current_region['polygon'].exterior.coords])
|
||||
current_region['offset_x'] = x
|
||||
current_region['offset_y'] = y
|
||||
|
||||
# Mouse left button up event
|
||||
elif event == cv2.EVENT_LBUTTONUP:
|
||||
|
|
@ -60,26 +66,33 @@ def mouse_callback(event, x, y, flags, param):
|
|||
current_region['dragging'] = False
|
||||
|
||||
|
||||
def run(weights='yolov8n.pt',
|
||||
source='test.mp4',
|
||||
view_img=False,
|
||||
save_img=False,
|
||||
exist_ok=False,
|
||||
line_thickness=2,
|
||||
region_thickness=2):
|
||||
def run(
|
||||
weights='yolov8n.pt',
|
||||
source=None,
|
||||
device='cpu',
|
||||
view_img=False,
|
||||
save_img=False,
|
||||
exist_ok=False,
|
||||
line_thickness=2,
|
||||
track_thickness=2,
|
||||
region_thickness=2,
|
||||
):
|
||||
"""
|
||||
Run Region counting on a video using YOLOv8 and ByteTrack.
|
||||
|
||||
Supports movable region for real time counting inside specific area.
|
||||
Supports multiple regions counting.
|
||||
Regions can be Polygons or rectangle in shape
|
||||
|
||||
Args:
|
||||
weights (str): Model weights path.
|
||||
source (str): Video file path.
|
||||
device (str): processing device cpu, 0, 1
|
||||
view_img (bool): Show results.
|
||||
save_img (bool): Save results.
|
||||
exist_ok (bool): Overwrite existing files.
|
||||
line_thickness (int): Bounding box thickness.
|
||||
track_thickness (int): Tracking line thickness
|
||||
region_thickness (int): Region thickness.
|
||||
"""
|
||||
vid_frame_count = 0
|
||||
|
|
@ -90,6 +103,7 @@ def run(weights='yolov8n.pt',
|
|||
|
||||
# Setup Model
|
||||
model = YOLO(f'{weights}')
|
||||
model.to('cuda') if device == '0' else model.to('cpu')
|
||||
|
||||
# Video setup
|
||||
videocapture = cv2.VideoCapture(source)
|
||||
|
|
@ -122,40 +136,43 @@ def run(weights='yolov8n.pt',
|
|||
label = str(names[cls])
|
||||
xyxy = (x - w / 2), (y - h / 2), (x + w / 2), (y + h / 2)
|
||||
|
||||
# Bounding box
|
||||
# Bounding box plot
|
||||
bbox_color = colors(cls, True)
|
||||
annotator.box_label(xyxy, label, color=bbox_color)
|
||||
|
||||
# Tracking Lines
|
||||
# Tracking Lines plot
|
||||
track = track_history[track_id]
|
||||
track.append((float(x), float(y)))
|
||||
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=line_thickness)
|
||||
cv2.polylines(frame, [points], isClosed=False, color=bbox_color, thickness=track_thickness)
|
||||
|
||||
# Check If detection inside region
|
||||
# Check if detection inside region
|
||||
for region in counting_regions:
|
||||
if is_inside_roi(box, region['roi']):
|
||||
if is_inside_polygon((x, y), region['polygon']):
|
||||
region['counts'] += 1
|
||||
|
||||
# Draw region boxes
|
||||
# Draw regions (Polygons/Rectangles)
|
||||
for region in counting_regions:
|
||||
region_label = str(region['counts'])
|
||||
roi_x, roi_y, roi_w, roi_h = region['roi']
|
||||
region_color = region['region_color']
|
||||
center_x = roi_x + roi_w // 2
|
||||
center_y = roi_y + roi_h // 2
|
||||
text_margin = 15
|
||||
region_text_color = region['text_color']
|
||||
|
||||
# Region plotting
|
||||
cv2.rectangle(frame, (roi_x, roi_y), (roi_x + roi_w, roi_y + roi_h), region_color, region_thickness)
|
||||
t_size, _ = cv2.getTextSize(region_label, cv2.FONT_HERSHEY_SIMPLEX, fontScale=1.0, thickness=line_thickness)
|
||||
text_x = center_x - t_size[0] // 2 - text_margin
|
||||
text_y = center_y + t_size[1] // 2 + text_margin
|
||||
cv2.rectangle(frame, (text_x - text_margin, text_y - t_size[1] - text_margin),
|
||||
(text_x + t_size[0] + text_margin, text_y + text_margin), region_color, -1)
|
||||
cv2.putText(frame, region_label, (text_x, text_y), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 0), line_thickness)
|
||||
polygon_coords = np.array(region['polygon'].exterior.coords, dtype=np.int32)
|
||||
centroid_x, centroid_y = int(region['polygon'].centroid.x), int(region['polygon'].centroid.y)
|
||||
|
||||
text_size, _ = cv2.getTextSize(region_label,
|
||||
cv2.FONT_HERSHEY_SIMPLEX,
|
||||
fontScale=0.7,
|
||||
thickness=line_thickness)
|
||||
text_x = centroid_x - text_size[0] // 2
|
||||
text_y = centroid_y + text_size[1] // 2
|
||||
cv2.rectangle(frame, (text_x - 5, text_y - text_size[1] - 5), (text_x + text_size[0] + 5, text_y + 5),
|
||||
region_color, -1)
|
||||
cv2.putText(frame, region_label, (text_x, text_y), cv2.FONT_HERSHEY_SIMPLEX, 0.7, region_text_color,
|
||||
line_thickness)
|
||||
cv2.polylines(frame, [polygon_coords], isClosed=True, color=region_color, thickness=region_thickness)
|
||||
|
||||
if view_img:
|
||||
if vid_frame_count == 1:
|
||||
|
|
@ -182,12 +199,15 @@ def parse_opt():
|
|||
"""Parse command line arguments."""
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--weights', type=str, default='yolov8n.pt', help='initial weights path')
|
||||
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
|
||||
parser.add_argument('--source', type=str, required=True, help='video file path')
|
||||
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('--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')
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue