ultralytics 8.3.5__py3-none-any.whl → 8.3.7__py3-none-any.whl

tests/test_solutions.py

@@ -19,8 +19,8 @@ def test_major_solutions():
     cap = cv2.VideoCapture("solutions_ci_demo.mp4")
     assert cap.isOpened(), "Error reading video file"
     region_points = [(20, 400), (1080, 404), (1080, 360), (20, 360)]
-    # counter = solutions.ObjectCounter(reg_pts=region_points, names=names, view_img=False)
-    heatmap = solutions.Heatmap(colormap=cv2.COLORMAP_PARULA, names=names, view_img=False)
+    counter = solutions.ObjectCounter(region=region_points, model="yolo11n.pt", show=False)
+    heatmap = solutions.Heatmap(colormap=cv2.COLORMAP_PARULA, model="yolo11n.pt", show=False)
     speed = solutions.SpeedEstimator(reg_pts=region_points, names=names, view_img=False)
     queue = solutions.QueueManager(names=names, reg_pts=region_points, view_img=False)
     while cap.isOpened():
@@ -29,8 +29,8 @@ def test_major_solutions():
             break
         original_im0 = im0.copy()
         tracks = model.track(im0, persist=True, show=False)
-        # _ = counter.start_counting(original_im0.copy(), tracks)
-        _ = heatmap.generate_heatmap(original_im0.copy(), tracks)
+        _ = counter.count(original_im0.copy())
+        _ = heatmap.generate_heatmap(original_im0.copy())
         _ = speed.estimate_speed(original_im0.copy(), tracks)
         _ = queue.process_queue(original_im0.copy(), tracks)
     cap.release()
@@ -41,16 +41,14 @@ def test_major_solutions():
 def test_aigym():
     """Test the workouts monitoring solution."""
     safe_download(url=WORKOUTS_SOLUTION_DEMO)
-    model = YOLO("yolo11n-pose.pt")
     cap = cv2.VideoCapture("solution_ci_pose_demo.mp4")
     assert cap.isOpened(), "Error reading video file"
-    gym_object = solutions.AIGym(line_thickness=2, pose_type="squat", kpts_to_check=[5, 11, 13])
+    gym = solutions.AIGym(line_width=2, kpts=[5, 11, 13])
     while cap.isOpened():
         success, im0 = cap.read()
         if not success:
             break
-        results = model.track(im0, verbose=False)
-        _ = gym_object.start_counting(im0, results)
+        _ = gym.monitor(im0)
     cap.release()
     cv2.destroyAllWindows()
 

ultralytics/__init__.py

@@ -1,6 +1,6 @@
 # Ultralytics YOLO 🚀, AGPL-3.0 license
 
-__version__ = "8.3.5"
+__version__ = "8.3.7"
 
 import os
 

ultralytics/cfg/solutions/default.yaml

@@ -10,3 +10,7 @@ show: True # Flag to control whether to display output image or not
 show_in: True # Flag to display objects moving *into* the defined region
 show_out: True # Flag to display objects moving *out of* the defined region
 classes: # To count specific classes
+up_angle: 145.0 # Workouts up_angle for counts, 145.0 is default value
+down_angle: 90 # Workouts down_angle for counts, 90 is default value
+kpts: [6, 8, 10] # Keypoints for workouts monitoring
+colormap: # Colormap for heatmap

ultralytics/data/converter.py

@@ -1,13 +1,18 @@
 # Ultralytics YOLO 🚀, AGPL-3.0 license
 
 import json
+import random
+import shutil
 from collections import defaultdict
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from pathlib import Path
 
 import cv2
 import numpy as np
+from PIL import Image
 
-from ultralytics.utils import LOGGER, TQDM
+from ultralytics.utils import DATASETS_DIR, LOGGER, NUM_THREADS, TQDM
+from ultralytics.utils.downloads import download
 from ultralytics.utils.files import increment_path
 
 
@@ -588,15 +593,13 @@ def yolo_bbox2segment(im_dir, save_dir=None, sam_model="sam_b.pt"):
 
             - im_dir
                 ├─ 001.jpg
-                ├─ ..
+                ├─ ...
                 └─ NNN.jpg
             - labels
                 ├─ 001.txt
-                ├─ ..
+                ├─ ...
                 └─ NNN.txt
     """
-    from tqdm import tqdm
-
     from ultralytics import SAM
     from ultralytics.data import YOLODataset
     from ultralytics.utils import LOGGER
@@ -610,7 +613,7 @@ def yolo_bbox2segment(im_dir, save_dir=None, sam_model="sam_b.pt"):
 
     LOGGER.info("Detection labels detected, generating segment labels by SAM model!")
     sam_model = SAM(sam_model)
-    for label in tqdm(dataset.labels, total=len(dataset.labels), desc="Generating segment labels"):
+    for label in TQDM(dataset.labels, total=len(dataset.labels), desc="Generating segment labels"):
         h, w = label["shape"]
         boxes = label["bboxes"]
         if len(boxes) == 0:  # skip empty labels
@@ -635,3 +638,58 @@ def yolo_bbox2segment(im_dir, save_dir=None, sam_model="sam_b.pt"):
             with open(txt_file, "a") as f:
                 f.writelines(text + "\n" for text in texts)
     LOGGER.info(f"Generated segment labels saved in {save_dir}")
+
+
+def create_synthetic_coco_dataset():
+    """
+    Creates a synthetic COCO dataset with random images and existing labels.
+
+    This function downloads COCO labels, creates synthetic images for train2017 and val2017 subsets, and organizes
+    them in the COCO dataset structure. It uses multithreading to generate images efficiently.
+
+    Examples:
+        >>> create_synthetic_coco_dataset()
+
+    Notes:
+        - Requires internet connection to download label files.
+        - Generates random RGB images of varying sizes (480x480 to 640x640 pixels).
+        - Existing test2017 directory is removed as it's not needed.
+        - If label directories don't exist, image creation for that subset is skipped.
+    """
+
+    def create_synthetic_image(image_file):
+        """Generates synthetic images with random sizes and colors for dataset augmentation or testing purposes."""
+        if not image_file.exists():
+            size = (random.randint(480, 640), random.randint(480, 640))
+            Image.new(
+                "RGB",
+                size=size,
+                color=(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)),
+            ).save(image_file)
+
+    # Download labels
+    dir = DATASETS_DIR / "coco"
+    url = "https://github.com/ultralytics/assets/releases/download/v0.0.0/"
+    label_zip = "coco2017labels-segments.zip"
+    download([url + label_zip], dir=dir.parent)
+
+    # Create synthetic images
+    shutil.rmtree(dir / "labels" / "test2017", ignore_errors=True)  # Remove test2017 directory as not needed
+    with ThreadPoolExecutor(max_workers=NUM_THREADS) as executor:
+        for subset in ["train2017", "val2017"]:
+            subset_dir = dir / "images" / subset
+            subset_dir.mkdir(parents=True, exist_ok=True)
+
+            label_dir = dir / "labels" / subset
+            if label_dir.exists():
+                label_files = list(label_dir.glob("*.txt"))
+                image_files = [subset_dir / f"{label_file.stem}.jpg" for label_file in label_files]
+
+                # Submit all tasks
+                futures = [executor.submit(create_synthetic_image, image_file) for image_file in image_files]
+                for _ in TQDM(as_completed(futures), total=len(futures), desc=f"Generating images for {subset}"):
+                    pass  # The actual work is done in the background
+            else:
+                print(f"Warning: Label directory {label_dir} does not exist. Skipping image creation for {subset}.")
+
+    print("Synthetic COCO dataset created successfully.")

ultralytics/data/explorer/gui/dash.py

@@ -39,24 +39,11 @@ def init_explorer_form(data=None, model=None):
     else:
         ds = [data]
 
+    prefixes = ["yolov8", "yolo11"]
+    sizes = ["n", "s", "m", "l", "x"]
+    tasks = ["", "-seg", "-pose"]
     if model is None:
-        models = [
-            "yolov8n.pt",
-            "yolov8s.pt",
-            "yolov8m.pt",
-            "yolov8l.pt",
-            "yolov8x.pt",
-            "yolov8n-seg.pt",
-            "yolov8s-seg.pt",
-            "yolov8m-seg.pt",
-            "yolov8l-seg.pt",
-            "yolov8x-seg.pt",
-            "yolov8n-pose.pt",
-            "yolov8s-pose.pt",
-            "yolov8m-pose.pt",
-            "yolov8l-pose.pt",
-            "yolov8x-pose.pt",
-        ]
+        models = [f"{p}{s}{t}" for p in prefixes for s in sizes for t in tasks]
     else:
         models = [model]
 

ultralytics/engine/exporter.py

@@ -183,11 +183,10 @@ class Exporter:
 
             # Get the closest match if format is invalid
             matches = difflib.get_close_matches(fmt, fmts, n=1, cutoff=0.6)  # 60% similarity required to match
-            if matches:
-                LOGGER.warning(f"WARNING ⚠️ Invalid export format='{fmt}', updating to format='{matches[0]}'")
-                fmt = matches[0]
-            else:
+            if not matches:
                 raise ValueError(f"Invalid export format='{fmt}'. Valid formats are {fmts}")
+            LOGGER.warning(f"WARNING ⚠️ Invalid export format='{fmt}', updating to format='{matches[0]}'")
+            fmt = matches[0]
         flags = [x == fmt for x in fmts]
         if sum(flags) != 1:
             raise ValueError(f"Invalid export format='{fmt}'. Valid formats are {fmts}")

ultralytics/engine/model.py

@@ -544,6 +544,8 @@ class Model(nn.Module):
 
         if not self.predictor:
             self.predictor = predictor or self._smart_load("predictor")(overrides=args, _callbacks=self.callbacks)
+            if predictor:
+                self.predictor.args = get_cfg(self.predictor.args, args)
             self.predictor.setup_model(model=self.model, verbose=is_cli)
         else:  # only update args if predictor is already setup
             self.predictor.args = get_cfg(self.predictor.args, args)

ultralytics/engine/trainer.py

@@ -469,11 +469,11 @@ class BaseTrainer:
 
         if RANK in {-1, 0}:
             # Do final val with best.pt
-            LOGGER.info(
-                f"\n{epoch - self.start_epoch + 1} epochs completed in "
-                f"{(time.time() - self.train_time_start) / 3600:.3f} hours."
-            )
+            epochs = epoch - self.start_epoch + 1  # total training epochs
+            seconds = time.time() - self.train_time_start  # total training seconds
+            LOGGER.info(f"\n{epochs} epochs completed in {seconds / 3600:.3f} hours.")
             self.final_eval()
+            self.validator.metrics.training = {"epochs": epochs, "seconds": seconds}  # add training speed
             if self.args.plots:
                 self.plot_metrics()
             self.run_callbacks("on_train_end")

ultralytics/solutions/ai_gym.py

@@ -1,127 +1,79 @@
 # Ultralytics YOLO 🚀, AGPL-3.0 license
 
-import cv2
-
-from ultralytics.utils.checks import check_imshow
+from ultralytics.solutions.solutions import BaseSolution  # Import a parent class
 from ultralytics.utils.plotting import Annotator
 
 
-class AIGym:
+class AIGym(BaseSolution):
     """A class to manage the gym steps of people in a real-time video stream based on their poses."""
 
-    def __init__(
-        self,
-        kpts_to_check,
-        line_thickness=2,
-        view_img=False,
-        pose_up_angle=145.0,
-        pose_down_angle=90.0,
-        pose_type="pullup",
-    ):
+    def __init__(self, **kwargs):
+        """Initialization function for AiGYM class, a child class of BaseSolution class, can be used for workouts
+        monitoring.
         """
-        Initializes the AIGym class with the specified parameters.
-
-        Args:
-            kpts_to_check (list): Indices of keypoints to check.
-            line_thickness (int, optional): Thickness of the lines drawn. Defaults to 2.
-            view_img (bool, optional): Flag to display the image. Defaults to False.
-            pose_up_angle (float, optional): Angle threshold for the 'up' pose. Defaults to 145.0.
-            pose_down_angle (float, optional): Angle threshold for the 'down' pose. Defaults to 90.0.
-            pose_type (str, optional): Type of pose to detect ('pullup', 'pushup', 'abworkout'). Defaults to "pullup".
+        # Check if the model name ends with '-pose'
+        if "model" in kwargs and "-pose" not in kwargs["model"]:
+            kwargs["model"] = "yolo11n-pose.pt"
+        elif "model" not in kwargs:
+            kwargs["model"] = "yolo11n-pose.pt"
+
+        super().__init__(**kwargs)
+        self.count = []  # List for counts, necessary where there are multiple objects in frame
+        self.angle = []  # List for angle, necessary where there are multiple objects in frame
+        self.stage = []  # List for stage, necessary where there are multiple objects in frame
+
+        # Extract details from CFG single time for usage later
+        self.initial_stage = None
+        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):
         """
-        # Image and line thickness
-        self.im0 = None
-        self.tf = line_thickness
-
-        # Keypoints and count information
-        self.keypoints = None
-        self.poseup_angle = pose_up_angle
-        self.posedown_angle = pose_down_angle
-        self.threshold = 0.001
-
-        # Store stage, count and angle information
-        self.angle = None
-        self.count = None
-        self.stage = None
-        self.pose_type = pose_type
-        self.kpts_to_check = kpts_to_check
-
-        # Visual Information
-        self.view_img = view_img
-        self.annotator = None
-
-        # Check if environment supports imshow
-        self.env_check = check_imshow(warn=True)
-        self.count = []
-        self.angle = []
-        self.stage = []
-
-    def start_counting(self, im0, results):
-        """
-        Function used to count the gym steps.
+        Monitor the workouts using Ultralytics YOLOv8 Pose Model: https://docs.ultralytics.com/tasks/pose/.
 
         Args:
-            im0 (ndarray): Current frame from the video stream.
-            results (list): Pose estimation data.
+            im0 (ndarray): The input image that will be used for processing
+        Returns
+            im0 (ndarray): The processed image for more usage
         """
-        self.im0 = im0
-
-        if not len(results[0]):
-            return self.im0
-
-        if len(results[0]) > len(self.count):
-            new_human = len(results[0]) - len(self.count)
-            self.count += [0] * new_human
-            self.angle += [0] * new_human
-            self.stage += ["-"] * new_human
-
-        self.keypoints = results[0].keypoints.data
-        self.annotator = Annotator(im0, line_width=self.tf)
-
-        for ind, k in enumerate(reversed(self.keypoints)):
-            # Estimate angle and draw specific points based on pose type
-            if self.pose_type in {"pushup", "pullup", "abworkout", "squat"}:
-                self.angle[ind] = self.annotator.estimate_pose_angle(
-                    k[int(self.kpts_to_check[0])].cpu(),
-                    k[int(self.kpts_to_check[1])].cpu(),
-                    k[int(self.kpts_to_check[2])].cpu(),
-                )
-                self.im0 = self.annotator.draw_specific_points(k, self.kpts_to_check, shape=(640, 640), radius=10)
-
-                # Check and update pose stages and counts based on angle
-                if self.pose_type in {"abworkout", "pullup"}:
-                    if self.angle[ind] > self.poseup_angle:
-                        self.stage[ind] = "down"
-                    if self.angle[ind] < self.posedown_angle and self.stage[ind] == "down":
-                        self.stage[ind] = "up"
-                        self.count[ind] += 1
-
-                elif self.pose_type in {"pushup", "squat"}:
-                    if self.angle[ind] > self.poseup_angle:
-                        self.stage[ind] = "up"
-                    if self.angle[ind] < self.posedown_angle and self.stage[ind] == "up":
-                        self.stage[ind] = "down"
+        # Extract tracks
+        tracks = self.model.track(source=im0, persist=True, classes=self.CFG["classes"])[0]
+
+        if tracks.boxes.id is not None:
+            # Extract and check keypoints
+            if len(tracks) > len(self.count):
+                new_human = len(tracks) - len(self.count)
+                self.angle += [0] * new_human
+                self.count += [0] * new_human
+                self.stage += ["-"] * new_human
+
+            # Initialize annotator
+            self.annotator = Annotator(im0, line_width=self.lw)
+
+            # 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)
+
+                # Determine stage and count logic based on angle thresholds
+                if self.angle[ind] < self.down_angle:
+                    if self.stage[ind] == "up":
                         self.count[ind] += 1
+                    self.stage[ind] = "down"
+                elif self.angle[ind] > self.up_angle:
+                    self.stage[ind] = "up"
 
+                # Display angle, count, and stage text
                 self.annotator.plot_angle_and_count_and_stage(
-                    angle_text=self.angle[ind],
-                    count_text=self.count[ind],
-                    stage_text=self.stage[ind],
-                    center_kpt=k[int(self.kpts_to_check[1])],
+                    angle_text=self.angle[ind],  # angle text for display
+                    count_text=self.count[ind],  # count text for workouts
+                    stage_text=self.stage[ind],  # stage position text
+                    center_kpt=k[int(self.kpts[1])],  # center keypoint for display
                 )
 
-            # Draw keypoints
-            self.annotator.kpts(k, shape=(640, 640), radius=1, kpt_line=True)
-
-        # Display the image if environment supports it and view_img is True
-        if self.env_check and self.view_img:
-            cv2.imshow("Ultralytics YOLOv8 AI GYM", self.im0)
-            if cv2.waitKey(1) & 0xFF == ord("q"):
-                return
-
-        return self.im0
-
-
-if __name__ == "__main__":
-    kpts_to_check = [0, 1, 2]  # example keypoints
-    aigym = AIGym(kpts_to_check)
+        self.display_output(im0)  # Display output image, if environment support display
+        return im0  # return an image for writing or further usage