ultralytics 8.3.4__py3-none-any.whl → 8.3.6__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.4"
+__version__ = "8.3.6"
 
 import os
 

ultralytics/cfg/default.yaml

@@ -1,7 +1,7 @@
 # Ultralytics YOLO 🚀, AGPL-3.0 license
 # Default training settings and hyperparameters for medium-augmentation COCO training
 
-task: detect # (str) YOLO task, i.e. detect, segment, classify, pose
+task: detect # (str) YOLO task, i.e. detect, segment, classify, pose, obb
 mode: train # (str) YOLO mode, i.e. train, val, predict, export, track, benchmark
 
 # Train settings -------------------------------------------------------------------------------------------------------

ultralytics/cfg/solutions/default.yaml

@@ -0,0 +1,16 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+
+# Configuration for Ultralytics Solutions
+
+model: "yolo11n.pt" # The Ultralytics YOLO11 model to be used (e.g., yolo11n.pt for YOLO11 nano version)
+
+region: # Object counting, queue or speed estimation region points
+line_width: 2 # Thickness of the lines used to draw regions on the image/video frames
+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/base.py

@@ -90,13 +90,15 @@ class BaseDataset(Dataset):
         self.ims, self.im_hw0, self.im_hw = [None] * self.ni, [None] * self.ni, [None] * self.ni
         self.npy_files = [Path(f).with_suffix(".npy") for f in self.im_files]
         self.cache = cache.lower() if isinstance(cache, str) else "ram" if cache is True else None
-        if (self.cache == "ram" and self.check_cache_ram()) or self.cache == "disk":
-            if self.cache == "ram" and hyp.deterministic:
+        if self.cache == "ram" and self.check_cache_ram():
+            if hyp.deterministic:
                 LOGGER.warning(
                     "WARNING ⚠️ cache='ram' may produce non-deterministic training results. "
                     "Consider cache='disk' as a deterministic alternative if your disk space allows."
                 )
             self.cache_images()
+        elif self.cache == "disk" and self.check_cache_disk():
+            self.cache_images()
 
         # Transforms
         self.transforms = self.build_transforms(hyp=hyp)
@@ -206,25 +208,55 @@ class BaseDataset(Dataset):
         if not f.exists():
             np.save(f.as_posix(), cv2.imread(self.im_files[i]), allow_pickle=False)
 
+    def check_cache_disk(self, safety_margin=0.5):
+        """Check image caching requirements vs available disk space."""
+        import shutil
+
+        b, gb = 0, 1 << 30  # bytes of cached images, bytes per gigabytes
+        n = min(self.ni, 30)  # extrapolate from 30 random images
+        for _ in range(n):
+            im_file = random.choice(self.im_files)
+            im = cv2.imread(im_file)
+            if im is None:
+                continue
+            b += im.nbytes
+            if not os.access(Path(im_file).parent, os.W_OK):
+                self.cache = None
+                LOGGER.info(f"{self.prefix}Skipping caching images to disk, directory not writeable ⚠️")
+                return False
+        disk_required = b * self.ni / n * (1 + safety_margin)  # bytes required to cache dataset to disk
+        total, used, free = shutil.disk_usage(Path(self.im_files[0]).parent)
+        if disk_required > free:
+            self.cache = None
+            LOGGER.info(
+                f"{self.prefix}{disk_required / gb:.1f}GB disk space required, "
+                f"with {int(safety_margin * 100)}% safety margin but only "
+                f"{free / gb:.1f}/{total / gb:.1f}GB free, not caching images to disk ⚠️"
+            )
+            return False
+        return True
+
     def check_cache_ram(self, safety_margin=0.5):
         """Check image caching requirements vs available memory."""
         b, gb = 0, 1 << 30  # bytes of cached images, bytes per gigabytes
         n = min(self.ni, 30)  # extrapolate from 30 random images
         for _ in range(n):
             im = cv2.imread(random.choice(self.im_files))  # sample image
+            if im is None:
+                continue
             ratio = self.imgsz / max(im.shape[0], im.shape[1])  # max(h, w)  # ratio
             b += im.nbytes * ratio**2
         mem_required = b * self.ni / n * (1 + safety_margin)  # GB required to cache dataset into RAM
         mem = psutil.virtual_memory()
-        success = mem_required < mem.available  # to cache or not to cache, that is the question
-        if not success:
+        if mem_required > mem.available:
             self.cache = None
             LOGGER.info(
                 f"{self.prefix}{mem_required / gb:.1f}GB RAM required to cache images "
                 f"with {int(safety_margin * 100)}% safety margin but only "
                 f"{mem.available / gb:.1f}/{mem.total / gb:.1f}GB available, not caching images ⚠️"
             )
-        return success
+            return False
+        return True
 
     def set_rectangle(self):
         """Sets the shape of bounding boxes for YOLO detections as rectangles."""

ultralytics/data/utils.py

@@ -216,7 +216,7 @@ def polygons2masks_overlap(imgsz, segments, downsample_ratio=1):
     ms = []
     for si in range(len(segments)):
         mask = polygon2mask(imgsz, [segments[si].reshape(-1)], downsample_ratio=downsample_ratio, color=1)
-        ms.append(mask)
+        ms.append(mask.astype(masks.dtype))
         areas.append(mask.sum())
     areas = np.asarray(areas)
     index = np.argsort(-areas)
@@ -452,12 +452,12 @@ class HUBDatasetStats:
         path = Path(path).resolve()
         LOGGER.info(f"Starting HUB dataset checks for {path}....")
 
-        self.task = task  # detect, segment, pose, classify
+        self.task = task  # detect, segment, pose, classify, obb
         if self.task == "classify":
             unzip_dir = unzip_file(path)
             data = check_cls_dataset(unzip_dir)
             data["path"] = unzip_dir
-        else:  # detect, segment, pose
+        else:  # detect, segment, pose, obb
             _, data_dir, yaml_path = self._unzip(Path(path))
             try:
                 # Load YAML with checks

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/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/engine/validator.py

@@ -119,6 +119,8 @@ class BaseValidator:
             self.args.plots &= trainer.stopper.possible_stop or (trainer.epoch == trainer.epochs - 1)
             model.eval()
         else:
+            if str(self.args.model).endswith(".yaml"):
+                LOGGER.warning("WARNING ⚠️ validating an untrained model YAML will result in 0 mAP.")
             callbacks.add_integration_callbacks(self)
             model = AutoBackend(
                 weights=model or self.args.model,

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