ultralytics 8.1.29__py3-none-any.whl → 8.3.62__py3-none-any.whl

ultralytics/solutions/ai_gym.py

@@ -1,158 +1,111 @@
-# Ultralytics YOLO 🚀, AGPL-3.0 license
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
-import cv2
-
-from ultralytics.utils.checks import check_imshow
+from ultralytics.solutions.solutions import BaseSolution
 from ultralytics.utils.plotting import Annotator
 
 
-class AIGym:
-    """A class to manage the gym steps of people in a real-time video stream based on their poses."""
-
-    def __init__(self):
-        """Initializes the AIGym with default values for Visual and Image parameters."""
-
-        # Image and line thickness
-        self.im0 = None
-        self.tf = None
-
-        # Keypoints and count information
-        self.keypoints = None
-        self.poseup_angle = None
-        self.posedown_angle = None
-        self.threshold = 0.001
-
-        # Store stage, count and angle information
-        self.angle = None
-        self.count = None
-        self.stage = None
-        self.pose_type = "pushup"
-        self.kpts_to_check = None
-
-        # Visual Information
-        self.view_img = False
-        self.annotator = None
-
-        # Check if environment support imshow
-        self.env_check = check_imshow(warn=True)
-
-    def set_args(
-        self,
-        kpts_to_check,
-        line_thickness=2,
-        view_img=False,
-        pose_up_angle=145.0,
-        pose_down_angle=90.0,
-        pose_type="pullup",
-    ):
+class AIGym(BaseSolution):
+    """
+    A class to manage gym steps of people in a real-time video stream based on their poses.
+
+    This class extends BaseSolution to monitor workouts using YOLO pose estimation models. It tracks and counts
+    repetitions of exercises based on predefined angle thresholds for up and down positions.
+
+    Attributes:
+        count (List[int]): Repetition counts for each detected person.
+        angle (List[float]): Current angle of the tracked body part for each person.
+        stage (List[str]): Current exercise stage ('up', 'down', or '-') for each person.
+        initial_stage (str | None): Initial stage of the exercise.
+        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.
+        annotator (Annotator): Object for drawing annotations on the image.
+
+    Methods:
+        monitor: Processes a frame to detect poses, calculate angles, and count repetitions.
+
+    Examples:
+        >>> gym = AIGym(model="yolov8n-pose.pt")
+        >>> image = cv2.imread("gym_scene.jpg")
+        >>> processed_image = gym.monitor(image)
+        >>> cv2.imshow("Processed Image", processed_image)
+        >>> cv2.waitKey(0)
+    """
+
+    def __init__(self, **kwargs):
+        """Initializes AIGym for workout monitoring using pose estimation and predefined angles."""
+        # 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
+
+    def monitor(self, im0):
         """
-        Configures the AIGym line_thickness, save image and view image parameters.
+        Monitors workouts using Ultralytics YOLO Pose Model.
 
-        Args:
-            kpts_to_check (list): 3 keypoints for counting
-            line_thickness (int): Line thickness for bounding boxes.
-            view_img (bool): display the im0
-            pose_up_angle (float): Angle to set pose position up
-            pose_down_angle (float): Angle to set pose position down
-            pose_type (str): "pushup", "pullup" or "abworkout"
-        """
-        self.kpts_to_check = kpts_to_check
-        self.tf = line_thickness
-        self.view_img = view_img
-        self.poseup_angle = pose_up_angle
-        self.posedown_angle = pose_down_angle
-        self.pose_type = pose_type
-
-    def start_counting(self, im0, results, frame_count):
-        """
-        Function used to count the gym steps.
+        This function processes an input image to track and analyze human poses for workout monitoring. It uses
+        the YOLO Pose model to detect keypoints, estimate angles, and count repetitions based on predefined
+        angle thresholds.
 
         Args:
-            im0 (ndarray): Current frame from the video stream.
-            results (list): Pose estimation data
-            frame_count (int): store current frame count
-        """
-        self.im0 = im0
-        if frame_count == 1:
-            self.count = [0] * len(results[0])
-            self.angle = [0] * len(results[0])
-            self.stage = ["-" for _ in results[0]]
-        self.keypoints = results[0].keypoints.data
-        self.annotator = Annotator(im0, line_width=2)
-
-        num_keypoints = len(results[0])
-
-        # Resize self.angle, self.count, and self.stage if the number of keypoints has changed
-        if len(self.angle) != num_keypoints:
-            self.angle = [0] * num_keypoints
-            self.count = [0] * num_keypoints
-            self.stage = ["-" for _ in range(num_keypoints)]
-
-        for ind, k in enumerate(reversed(self.keypoints)):
-            if self.pose_type in ["pushup", "pullup"]:
-                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)
+            im0 (ndarray): Input image for processing.
 
-            if self.pose_type == "abworkout":
-                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)
-                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
-                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])],
-                    line_thickness=self.tf,
-                )
+        Returns:
+            (ndarray): Processed image with annotations for workout monitoring.
 
-            if self.pose_type == "pushup":
-                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"
-                    self.count[ind] += 1
-                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])],
-                    line_thickness=self.tf,
-                )
-            if self.pose_type == "pullup":
-                if self.angle[ind] > self.poseup_angle:
+        Examples:
+            >>> gym = AIGym()
+            >>> image = cv2.imread("workout.jpg")
+            >>> processed_image = gym.monitor(image)
+        """
+        # Extract tracks
+        tracks = self.model.track(source=im0, persist=True, classes=self.CFG["classes"], **self.track_add_args)[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.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.line_width * 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"
-                if self.angle[ind] < self.posedown_angle and self.stage[ind] == "down":
+                elif self.angle[ind] > self.up_angle:
                     self.stage[ind] = "up"
-                    self.count[ind] += 1
+
+                # 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])],
-                    line_thickness=self.tf,
+                    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
                 )
 
-            self.annotator.kpts(k, shape=(640, 640), radius=1, kpt_line=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__":
-    AIGym()
+        self.display_output(im0)  # Display output image, if environment support display
+        return im0  # return an image for writing or further usage

ultralytics/solutions/analytics.py

@@ -0,0 +1,247 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from itertools import cycle
+
+import cv2
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
+from matplotlib.figure import Figure
+
+from ultralytics.solutions.solutions import BaseSolution  # Import a parent class
+
+
+class Analytics(BaseSolution):
+    """
+    A class for creating and updating various types of charts for visual analytics.
+
+    This class extends BaseSolution to provide functionality for generating line, bar, pie, and area charts
+    based on object detection and tracking data.
+
+    Attributes:
+        type (str): The type of analytics chart to generate ('line', 'bar', 'pie', or 'area').
+        x_label (str): Label for the x-axis.
+        y_label (str): Label for the y-axis.
+        bg_color (str): Background color of the chart frame.
+        fg_color (str): Foreground color of the chart frame.
+        title (str): Title of the chart window.
+        max_points (int): Maximum number of data points to display on the chart.
+        fontsize (int): Font size for text display.
+        color_cycle (cycle): Cyclic iterator for chart colors.
+        total_counts (int): Total count of detected objects (used for line charts).
+        clswise_count (Dict[str, int]): Dictionary for class-wise object counts.
+        fig (Figure): Matplotlib figure object for the chart.
+        ax (Axes): Matplotlib axes object for the chart.
+        canvas (FigureCanvas): Canvas for rendering the chart.
+
+    Methods:
+        process_data: Processes image data and updates the chart.
+        update_graph: Updates the chart with new data points.
+
+    Examples:
+        >>> analytics = Analytics(analytics_type="line")
+        >>> frame = cv2.imread("image.jpg")
+        >>> processed_frame = analytics.process_data(frame, frame_number=1)
+        >>> cv2.imshow("Analytics", processed_frame)
+    """
+
+    def __init__(self, **kwargs):
+        """Initialize Analytics class with various chart types for visual data representation."""
+        super().__init__(**kwargs)
+
+        self.type = self.CFG["analytics_type"]  # extract type of analytics
+        self.x_label = "Classes" if self.type in {"bar", "pie"} else "Frame#"
+        self.y_label = "Total Counts"
+
+        # Predefined data
+        self.bg_color = "#F3F3F3"  # background color of frame
+        self.fg_color = "#111E68"  # foreground color of frame
+        self.title = "Ultralytics Solutions"  # window name
+        self.max_points = 45  # maximum points to be drawn on window
+        self.fontsize = 25  # text font size for display
+        figsize = (19.2, 10.8)  # Set output image size 1920 * 1080
+        self.color_cycle = cycle(["#DD00BA", "#042AFF", "#FF4447", "#7D24FF", "#BD00FF"])
+
+        self.total_counts = 0  # count variable for storing total counts i.e. for line
+        self.clswise_count = {}  # dictionary for class-wise counts
+
+        # Ensure line and area chart
+        if self.type in {"line", "area"}:
+            self.lines = {}
+            self.fig = Figure(facecolor=self.bg_color, figsize=figsize)
+            self.canvas = FigureCanvas(self.fig)  # Set common axis properties
+            self.ax = self.fig.add_subplot(111, facecolor=self.bg_color)
+            if self.type == "line":
+                (self.line,) = self.ax.plot([], [], color="cyan", linewidth=self.line_width)
+        elif self.type in {"bar", "pie"}:
+            # Initialize bar or pie plot
+            self.fig, self.ax = plt.subplots(figsize=figsize, facecolor=self.bg_color)
+            self.canvas = FigureCanvas(self.fig)  # Set common axis properties
+            self.ax.set_facecolor(self.bg_color)
+            self.color_mapping = {}
+
+            if self.type == "pie":  # Ensure pie chart is circular
+                self.ax.axis("equal")
+
+    def process_data(self, im0, frame_number):
+        """
+        Processes image data and runs object tracking to update analytics charts.
+
+        Args:
+            im0 (np.ndarray): Input image for processing.
+            frame_number (int): Video frame number for plotting the data.
+
+        Returns:
+            (np.ndarray): Processed image with updated analytics chart.
+
+        Raises:
+            ModuleNotFoundError: If an unsupported chart type is specified.
+
+        Examples:
+            >>> analytics = Analytics(analytics_type="line")
+            >>> frame = np.zeros((480, 640, 3), dtype=np.uint8)
+            >>> processed_frame = analytics.process_data(frame, frame_number=1)
+        """
+        self.extract_tracks(im0)  # Extract tracks
+
+        if self.type == "line":
+            for _ in self.boxes:
+                self.total_counts += 1
+            im0 = self.update_graph(frame_number=frame_number)
+            self.total_counts = 0
+        elif self.type in {"pie", "bar", "area"}:
+            self.clswise_count = {}
+            for box, cls in zip(self.boxes, self.clss):
+                if self.names[int(cls)] in self.clswise_count:
+                    self.clswise_count[self.names[int(cls)]] += 1
+                else:
+                    self.clswise_count[self.names[int(cls)]] = 1
+            im0 = self.update_graph(frame_number=frame_number, count_dict=self.clswise_count, plot=self.type)
+        else:
+            raise ModuleNotFoundError(f"{self.type} chart is not supported ❌")
+        return im0
+
+    def update_graph(self, frame_number, count_dict=None, plot="line"):
+        """
+        Updates the graph with new data for single or multiple classes.
+
+        Args:
+            frame_number (int): The current frame number.
+            count_dict (Dict[str, int] | None): Dictionary with class names as keys and counts as values for multiple
+                classes. If None, updates a single line graph.
+            plot (str): Type of the plot. Options are 'line', 'bar', 'pie', or 'area'.
+
+        Returns:
+            (np.ndarray): Updated image containing the graph.
+
+        Examples:
+            >>> analytics = Analytics()
+            >>> frame_number = 10
+            >>> count_dict = {"person": 5, "car": 3}
+            >>> updated_image = analytics.update_graph(frame_number, count_dict, plot="bar")
+        """
+        if count_dict is None:
+            # Single line update
+            x_data = np.append(self.line.get_xdata(), float(frame_number))
+            y_data = np.append(self.line.get_ydata(), float(self.total_counts))
+
+            if len(x_data) > self.max_points:
+                x_data, y_data = x_data[-self.max_points :], y_data[-self.max_points :]
+
+            self.line.set_data(x_data, y_data)
+            self.line.set_label("Counts")
+            self.line.set_color("#7b0068")  # Pink color
+            self.line.set_marker("*")
+            self.line.set_markersize(self.line_width * 5)
+        else:
+            labels = list(count_dict.keys())
+            counts = list(count_dict.values())
+            if plot == "area":
+                color_cycle = cycle(["#DD00BA", "#042AFF", "#FF4447", "#7D24FF", "#BD00FF"])
+                # Multiple lines or area update
+                x_data = self.ax.lines[0].get_xdata() if self.ax.lines else np.array([])
+                y_data_dict = {key: np.array([]) for key in count_dict.keys()}
+                if self.ax.lines:
+                    for line, key in zip(self.ax.lines, count_dict.keys()):
+                        y_data_dict[key] = line.get_ydata()
+
+                x_data = np.append(x_data, float(frame_number))
+                max_length = len(x_data)
+                for key in count_dict.keys():
+                    y_data_dict[key] = np.append(y_data_dict[key], float(count_dict[key]))
+                    if len(y_data_dict[key]) < max_length:
+                        y_data_dict[key] = np.pad(y_data_dict[key], (0, max_length - len(y_data_dict[key])))
+                if len(x_data) > self.max_points:
+                    x_data = x_data[1:]
+                    for key in count_dict.keys():
+                        y_data_dict[key] = y_data_dict[key][1:]
+
+                self.ax.clear()
+                for key, y_data in y_data_dict.items():
+                    color = next(color_cycle)
+                    self.ax.fill_between(x_data, y_data, color=color, alpha=0.7)
+                    self.ax.plot(
+                        x_data,
+                        y_data,
+                        color=color,
+                        linewidth=self.line_width,
+                        marker="o",
+                        markersize=self.line_width * 5,
+                        label=f"{key} Data Points",
+                    )
+            if plot == "bar":
+                self.ax.clear()  # clear bar data
+                for label in labels:  # Map labels to colors
+                    if label not in self.color_mapping:
+                        self.color_mapping[label] = next(self.color_cycle)
+                colors = [self.color_mapping[label] for label in labels]
+                bars = self.ax.bar(labels, counts, color=colors)
+                for bar, count in zip(bars, counts):
+                    self.ax.text(
+                        bar.get_x() + bar.get_width() / 2,
+                        bar.get_height(),
+                        str(count),
+                        ha="center",
+                        va="bottom",
+                        color=self.fg_color,
+                    )
+                # Create the legend using labels from the bars
+                for bar, label in zip(bars, labels):
+                    bar.set_label(label)  # Assign label to each bar
+                self.ax.legend(loc="upper left", fontsize=13, facecolor=self.fg_color, edgecolor=self.fg_color)
+            if plot == "pie":
+                total = sum(counts)
+                percentages = [size / total * 100 for size in counts]
+                start_angle = 90
+                self.ax.clear()
+
+                # Create pie chart and create legend labels with percentages
+                wedges, autotexts = self.ax.pie(
+                    counts, labels=labels, startangle=start_angle, textprops={"color": self.fg_color}, autopct=None
+                )
+                legend_labels = [f"{label} ({percentage:.1f}%)" for label, percentage in zip(labels, percentages)]
+
+                # Assign the legend using the wedges and manually created labels
+                self.ax.legend(wedges, legend_labels, title="Classes", loc="center left", bbox_to_anchor=(1, 0, 0.5, 1))
+                self.fig.subplots_adjust(left=0.1, right=0.75)  # Adjust layout to fit the legend
+
+        # Common plot settings
+        self.ax.set_facecolor("#f0f0f0")  # Set to light gray or any other color you like
+        self.ax.set_title(self.title, color=self.fg_color, fontsize=self.fontsize)
+        self.ax.set_xlabel(self.x_label, color=self.fg_color, fontsize=self.fontsize - 3)
+        self.ax.set_ylabel(self.y_label, color=self.fg_color, fontsize=self.fontsize - 3)
+
+        # Add and format legend
+        legend = self.ax.legend(loc="upper left", fontsize=13, facecolor=self.bg_color, edgecolor=self.bg_color)
+        for text in legend.get_texts():
+            text.set_color(self.fg_color)
+
+        # Redraw graph, update view, capture, and display the updated plot
+        self.ax.relim()
+        self.ax.autoscale_view()
+        self.canvas.draw()
+        im0 = np.array(self.canvas.renderer.buffer_rgba())
+        im0 = cv2.cvtColor(im0[:, :, :3], cv2.COLOR_RGBA2BGR)
+        self.display_output(im0)
+
+        return im0  # Return the image