ultralytics 8.3.28 new Solutions CLI commands (#17233)

Signed-off-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Muhammad Rizwan Munawar 2024-11-07 04:44:05 +05:00 committed by GitHub
parent d049e22769
commit 3c976807b8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
17 changed files with 310 additions and 48 deletions

View file

@ -19,7 +19,6 @@ class AIGym(BaseSolution):
up_angle (float): Angle threshold for considering the 'up' position of an exercise.
down_angle (float): Angle threshold for considering the 'down' position of an exercise.
kpts (List[int]): Indices of keypoints used for angle calculation.
lw (int): Line width for drawing annotations.
annotator (Annotator): Object for drawing annotations on the image.
Methods:
@ -51,7 +50,6 @@ class AIGym(BaseSolution):
self.up_angle = float(self.CFG["up_angle"]) # Pose up predefined angle to consider up pose
self.down_angle = float(self.CFG["down_angle"]) # Pose down predefined angle to consider down pose
self.kpts = self.CFG["kpts"] # User selected kpts of workouts storage for further usage
self.lw = self.CFG["line_width"] # Store line_width for usage
def monitor(self, im0):
"""
@ -84,14 +82,14 @@ class AIGym(BaseSolution):
self.stage += ["-"] * new_human
# Initialize annotator
self.annotator = Annotator(im0, line_width=self.lw)
self.annotator = Annotator(im0, line_width=self.line_width)
# Enumerate over keypoints
for ind, k in enumerate(reversed(tracks.keypoints.data)):
# Get keypoints and estimate the angle
kpts = [k[int(self.kpts[i])].cpu() for i in range(3)]
self.angle[ind] = self.annotator.estimate_pose_angle(*kpts)
im0 = self.annotator.draw_specific_points(k, self.kpts, radius=self.lw * 3)
im0 = self.annotator.draw_specific_points(k, self.kpts, radius=self.line_width * 3)
# Determine stage and count logic based on angle thresholds
if self.angle[ind] < self.down_angle:

View file

@ -5,7 +5,7 @@ from collections import defaultdict
import cv2
from ultralytics import YOLO
from ultralytics.utils import DEFAULT_CFG_DICT, DEFAULT_SOL_DICT, LOGGER
from ultralytics.utils import ASSETS_URL, DEFAULT_CFG_DICT, DEFAULT_SOL_DICT, LOGGER
from ultralytics.utils.checks import check_imshow, check_requirements
@ -42,8 +42,12 @@ class BaseSolution:
>>> solution.display_output(image)
"""
def __init__(self, **kwargs):
"""Initializes the BaseSolution class with configuration settings and YOLO model for Ultralytics solutions."""
def __init__(self, IS_CLI=False, **kwargs):
"""
Initializes the `BaseSolution` class with configuration settings and the YOLO model for Ultralytics solutions.
IS_CLI (optional): Enables CLI mode if set.
"""
check_requirements("shapely>=2.0.0")
from shapely.geometry import LineString, Point, Polygon
@ -63,9 +67,20 @@ class BaseSolution:
) # Store line_width for usage
# Load Model and store classes names
self.model = YOLO(self.CFG["model"] if self.CFG["model"] else "yolov8n.pt")
if self.CFG["model"] is None:
self.CFG["model"] = "yolo11n.pt"
self.model = YOLO(self.CFG["model"])
self.names = self.model.names
if IS_CLI: # for CLI, download the source and init video writer
if self.CFG["source"] is None:
d_s = "solutions_ci_demo.mp4" if "-pose" not in self.CFG["model"] else "solution_ci_pose_demo.mp4"
LOGGER.warning(f"⚠️ WARNING: source not provided. using default source {ASSETS_URL}/{d_s}")
from ultralytics.utils.downloads import safe_download
safe_download(f"{ASSETS_URL}/{d_s}") # download source from ultralytics assets
self.CFG["source"] = d_s # set default source
# Initialize environment and region setup
self.env_check = check_imshow(warn=True)
self.track_history = defaultdict(list)