Reference for ultralytics/solutions/solutions.py

def __call__(self, *args: Any, **kwargs: Any):
    """Allow instances to be called like a function with flexible arguments."""
    with self.profilers[1]:
        result = self.process(*args, **kwargs)  # Call the subclass-specific process method
    track_or_predict = "predict" if type(self).__name__ == "ObjectCropper" else "track"
    track_or_predict_speed = self.profilers[0].dt * 1e3
    solution_speed = (self.profilers[1].dt - self.profilers[0].dt) * 1e3  # solution time = process - track
    result.speed = {track_or_predict: track_or_predict_speed, "solution": solution_speed}
    if self.CFG["verbose"]:
        self.frame_no += 1
        LOGGER.info(
            f"{self.frame_no}: {result.plot_im.shape[0]}x{result.plot_im.shape[1]} {solution_speed:.1f}ms\n"
            f"Speed: {track_or_predict_speed:.1f}ms {track_or_predict}, "
            f"{solution_speed:.1f}ms solution per image at shape "
            f"(1, {getattr(self.model, 'ch', 3)}, {result.plot_im.shape[0]}, {result.plot_im.shape[1]})\n"
        )
    return result

def adjust_box_label(self, cls: int, conf: float, track_id: Optional[int] = None) -> Optional[str]:
    """
    Generate a formatted label for a bounding box.

    This method constructs a label string for a bounding box using the class index and confidence score.
    Optionally includes the track ID if provided. The label format adapts based on the display settings
    defined in `self.show_conf` and `self.show_labels`.

    Args:
        cls (int): The class index of the detected object.
        conf (float): The confidence score of the detection.
        track_id (int, optional): The unique identifier for the tracked object.

    Returns:
        (str | None): The formatted label string if `self.show_labels` is True; otherwise, None.
    """
    name = ("" if track_id is None else f"{track_id} ") + self.names[cls]
    return (f"{name} {conf:.2f}" if self.show_conf else name) if self.show_labels else None

def display_output(self, plot_im: np.ndarray) -> None:
    """
    Display the results of the processing, which could involve showing frames, printing counts, or saving results.

    This method is responsible for visualizing the output of the object detection and tracking process. It displays
    the processed frame with annotations, and allows for user interaction to close the display.

    Args:
        plot_im (np.ndarray): The image or frame that has been processed and annotated.

    Examples:
        >>> solution = BaseSolution()
        >>> frame = cv2.imread("path/to/image.jpg")
        >>> solution.display_output(frame)

    Notes:
        - This method will only display output if the 'show' configuration is set to True and the environment
          supports image display.
        - The display can be closed by pressing the 'q' key.
    """
    if self.CFG.get("show") and self.env_check:
        cv2.imshow("Ultralytics Solutions", plot_im)
        if cv2.waitKey(1) & 0xFF == ord("q"):
            cv2.destroyAllWindows()  # Closes current frame window
            return

def extract_tracks(self, im0: np.ndarray) -> None:
    """
    Apply object tracking and extract tracks from an input image or frame.

    Args:
        im0 (np.ndarray): The input image or frame.

    Examples:
        >>> solution = BaseSolution()
        >>> frame = cv2.imread("path/to/image.jpg")
        >>> solution.extract_tracks(frame)
    """
    with self.profilers[0]:
        self.tracks = self.model.track(
            source=im0, persist=True, classes=self.classes, verbose=False, **self.track_add_args
        )[0]
    is_obb = self.tracks.obb is not None
    self.track_data = self.tracks.obb if is_obb else self.tracks.boxes  # Extract tracks for OBB or object detection

    if self.track_data and self.track_data.is_track:
        self.boxes = (self.track_data.xyxyxyxy if is_obb else self.track_data.xyxy).cpu()
        self.clss = self.track_data.cls.cpu().tolist()
        self.track_ids = self.track_data.id.int().cpu().tolist()
        self.confs = self.track_data.conf.cpu().tolist()
    else:
        self.LOGGER.warning("no tracks found!")
        self.boxes, self.clss, self.track_ids, self.confs = [], [], [], []

def initialize_region(self) -> None:
    """Initialize the counting region and line segment based on configuration settings."""
    if self.region is None:
        self.region = [(10, 200), (540, 200), (540, 180), (10, 180)]
    self.r_s = (
        self.Polygon(self.region) if len(self.region) >= 3 else self.LineString(self.region)
    )  # region or line

def process(self, *args: Any, **kwargs: Any):
    """Process method should be implemented by each Solution subclass."""

def store_tracking_history(self, track_id: int, box) -> None:
    """
    Store the tracking history of an object.

    This method updates the tracking history for a given object by appending the center point of its
    bounding box to the track line. It maintains a maximum of 30 points in the tracking history.

    Args:
        track_id (int): The unique identifier for the tracked object.
        box (List[float]): The bounding box coordinates of the object in the format [x1, y1, x2, y2].

    Examples:
        >>> solution = BaseSolution()
        >>> solution.store_tracking_history(1, [100, 200, 300, 400])
    """
    # Store tracking history
    self.track_line = self.track_history[track_id]
    self.track_line.append(tuple(box.mean(dim=0)) if box.numel() > 4 else (box[:4:2].mean(), box[1:4:2].mean()))
    if len(self.track_line) > 30:
        self.track_line.pop(0)

def circle_label(
    self,
    box: Tuple[float, float, float, float],
    label: str = "",
    color: Tuple[int, int, int] = (128, 128, 128),
    txt_color: Tuple[int, int, int] = (255, 255, 255),
    margin: int = 2,
):
    """
    Draw a label with a background circle centered within a given bounding box.

    Args:
        box (Tuple[float, float, float, float]): The bounding box coordinates (x1, y1, x2, y2).
        label (str): The text label to be displayed.
        color (Tuple[int, int, int]): The background color of the circle (B, G, R).
        txt_color (Tuple[int, int, int]): The color of the text (R, G, B).
        margin (int): The margin between the text and the circle border.
    """
    if len(label) > 3:
        LOGGER.warning(f"Length of label is {len(label)}, only first 3 letters will be used for circle annotation.")
        label = label[:3]

    # Calculate the center of the box
    x_center, y_center = int((box[0] + box[2]) / 2), int((box[1] + box[3]) / 2)
    # Get the text size
    text_size = cv2.getTextSize(str(label), cv2.FONT_HERSHEY_SIMPLEX, self.sf - 0.15, self.tf)[0]
    # Calculate the required radius to fit the text with the margin
    required_radius = int(((text_size[0] ** 2 + text_size[1] ** 2) ** 0.5) / 2) + margin
    # Draw the circle with the required radius
    cv2.circle(self.im, (x_center, y_center), required_radius, color, -1)
    # Calculate the position for the text
    text_x = x_center - text_size[0] // 2
    text_y = y_center + text_size[1] // 2
    # Draw the text
    cv2.putText(
        self.im,
        str(label),
        (text_x, text_y),
        cv2.FONT_HERSHEY_SIMPLEX,
        self.sf - 0.15,
        self.get_txt_color(color, txt_color),
        self.tf,
        lineType=cv2.LINE_AA,
    )

def display_analytics(
    self,
    im0: np.ndarray,
    text: Dict[str, Any],
    txt_color: Tuple[int, int, int],
    bg_color: Tuple[int, int, int],
    margin: int,
):
    """
    Display the overall statistics for parking lots, object counter etc.

    Args:
        im0 (np.ndarray): Inference image.
        text (Dict[str, Any]): Labels dictionary.
        txt_color (Tuple[int, int, int]): Display color for text foreground.
        bg_color (Tuple[int, int, int]): Display color for text background.
        margin (int): Gap between text and rectangle for better display.
    """
    horizontal_gap = int(im0.shape[1] * 0.02)
    vertical_gap = int(im0.shape[0] * 0.01)
    text_y_offset = 0
    for label, value in text.items():
        txt = f"{label}: {value}"
        text_size = cv2.getTextSize(txt, 0, self.sf, self.tf)[0]
        if text_size[0] < 5 or text_size[1] < 5:
            text_size = (5, 5)
        text_x = im0.shape[1] - text_size[0] - margin * 2 - horizontal_gap
        text_y = text_y_offset + text_size[1] + margin * 2 + vertical_gap
        rect_x1 = text_x - margin * 2
        rect_y1 = text_y - text_size[1] - margin * 2
        rect_x2 = text_x + text_size[0] + margin * 2
        rect_y2 = text_y + margin * 2
        cv2.rectangle(im0, (rect_x1, rect_y1), (rect_x2, rect_y2), bg_color, -1)
        cv2.putText(im0, txt, (text_x, text_y), 0, self.sf, txt_color, self.tf, lineType=cv2.LINE_AA)
        text_y_offset = rect_y2

def display_objects_labels(
    self,
    im0: np.ndarray,
    text: str,
    txt_color: Tuple[int, int, int],
    bg_color: Tuple[int, int, int],
    x_center: float,
    y_center: float,
    margin: int,
):
    """
    Display the bounding boxes labels in parking management app.

    Args:
        im0 (np.ndarray): Inference image.
        text (str): Object/class name.
        txt_color (Tuple[int, int, int]): Display color for text foreground.
        bg_color (Tuple[int, int, int]): Display color for text background.
        x_center (float): The x position center point for bounding box.
        y_center (float): The y position center point for bounding box.
        margin (int): The gap between text and rectangle for better display.
    """
    text_size = cv2.getTextSize(text, 0, fontScale=self.sf, thickness=self.tf)[0]
    text_x = x_center - text_size[0] // 2
    text_y = y_center + text_size[1] // 2

    rect_x1 = text_x - margin
    rect_y1 = text_y - text_size[1] - margin
    rect_x2 = text_x + text_size[0] + margin
    rect_y2 = text_y + margin
    cv2.rectangle(
        im0,
        (int(rect_x1), int(rect_y1)),
        (int(rect_x2), int(rect_y2)),
        tuple(map(int, bg_color)),  # Ensure color values are int
        -1,
    )

    cv2.putText(
        im0,
        text,
        (int(text_x), int(text_y)),
        0,
        self.sf,
        tuple(map(int, txt_color)),  # Ensure color values are int
        self.tf,
        lineType=cv2.LINE_AA,
    )

def draw_region(
    self,
    reg_pts: Optional[List[Tuple[int, int]]] = None,
    color: Tuple[int, int, int] = (0, 255, 0),
    thickness: int = 5,
):
    """
    Draw a region or line on the image.

    Args:
        reg_pts (List[Tuple[int, int]], optional): Region points (for line 2 points, for region 4+ points).
        color (Tuple[int, int, int]): RGB color value for the region.
        thickness (int): Line thickness for drawing the region.
    """
    cv2.polylines(self.im, [np.array(reg_pts, dtype=np.int32)], isClosed=True, color=color, thickness=thickness)

    # Draw small circles at the corner points
    for point in reg_pts:
        cv2.circle(self.im, (point[0], point[1]), thickness * 2, color, -1)  # -1 fills the circle

def draw_specific_kpts(
    self,
    keypoints: List[List[float]],
    indices: Optional[List[int]] = None,
    radius: int = 2,
    conf_thresh: float = 0.25,
) -> np.ndarray:
    """
    Draw specific keypoints for gym steps counting.

    Args:
        keypoints (List[List[float]]): Keypoints data to be plotted, each in format [x, y, confidence].
        indices (List[int], optional): Keypoint indices to be plotted.
        radius (int): Keypoint radius.
        conf_thresh (float): Confidence threshold for keypoints.

    Returns:
        (np.ndarray): Image with drawn keypoints.

    Notes:
        Keypoint format: [x, y] or [x, y, confidence].
        Modifies self.im in-place.
    """
    indices = indices or [2, 5, 7]
    points = [(int(k[0]), int(k[1])) for i, k in enumerate(keypoints) if i in indices and k[2] >= conf_thresh]

    # Draw lines between consecutive points
    for start, end in zip(points[:-1], points[1:]):
        cv2.line(self.im, start, end, (0, 255, 0), 2, lineType=cv2.LINE_AA)

    # Draw circles for keypoints
    for pt in points:
        cv2.circle(self.im, pt, radius, (0, 0, 255), -1, lineType=cv2.LINE_AA)

    return self.im

@staticmethod
@lru_cache(maxsize=256)
def estimate_pose_angle(a: List[float], b: List[float], c: List[float]) -> float:
    """
    Calculate the angle between three points for workout monitoring.

    Args:
        a (List[float]): The coordinates of the first point.
        b (List[float]): The coordinates of the second point (vertex).
        c (List[float]): The coordinates of the third point.

    Returns:
        (float): The angle in degrees between the three points.
    """
    radians = math.atan2(c[1] - b[1], c[0] - b[0]) - math.atan2(a[1] - b[1], a[0] - b[0])
    angle = abs(radians * 180.0 / math.pi)
    return angle if angle <= 180.0 else (360 - angle)

def plot_angle_and_count_and_stage(
    self,
    angle_text: str,
    count_text: str,
    stage_text: str,
    center_kpt: List[int],
    color: Tuple[int, int, int] = (104, 31, 17),
    txt_color: Tuple[int, int, int] = (255, 255, 255),
):
    """
    Plot the pose angle, count value, and step stage for workout monitoring.

    Args:
        angle_text (str): Angle value for workout monitoring.
        count_text (str): Counts value for workout monitoring.
        stage_text (str): Stage decision for workout monitoring.
        center_kpt (List[int]): Centroid pose index for workout monitoring.
        color (Tuple[int, int, int]): Text background color.
        txt_color (Tuple[int, int, int]): Text foreground color.
    """
    # Format text
    angle_text, count_text, stage_text = f" {angle_text:.2f}", f"Steps : {count_text}", f" {stage_text}"

    # Draw angle, count and stage text
    angle_height = self.plot_workout_information(
        angle_text, (int(center_kpt[0]), int(center_kpt[1])), color, txt_color
    )
    count_height = self.plot_workout_information(
        count_text, (int(center_kpt[0]), int(center_kpt[1]) + angle_height + 20), color, txt_color
    )
    self.plot_workout_information(
        stage_text, (int(center_kpt[0]), int(center_kpt[1]) + angle_height + count_height + 40), color, txt_color
    )

def plot_distance_and_line(
    self,
    pixels_distance: float,
    centroids: List[Tuple[int, int]],
    line_color: Tuple[int, int, int] = (104, 31, 17),
    centroid_color: Tuple[int, int, int] = (255, 0, 255),
):
    """
    Plot the distance and line between two centroids on the frame.

    Args:
        pixels_distance (float): Pixels distance between two bbox centroids.
        centroids (List[Tuple[int, int]]): Bounding box centroids data.
        line_color (Tuple[int, int, int]): Distance line color.
        centroid_color (Tuple[int, int, int]): Bounding box centroid color.
    """
    # Get the text size
    text = f"Pixels Distance: {pixels_distance:.2f}"
    (text_width_m, text_height_m), _ = cv2.getTextSize(text, 0, self.sf, self.tf)

    # Define corners with 10-pixel margin and draw rectangle
    cv2.rectangle(self.im, (15, 25), (15 + text_width_m + 20, 25 + text_height_m + 20), line_color, -1)

    # Calculate the position for the text with a 10-pixel margin and draw text
    text_position = (25, 25 + text_height_m + 10)
    cv2.putText(
        self.im,
        text,
        text_position,
        0,
        self.sf,
        (255, 255, 255),
        self.tf,
        cv2.LINE_AA,
    )

    cv2.line(self.im, centroids[0], centroids[1], line_color, 3)
    cv2.circle(self.im, centroids[0], 6, centroid_color, -1)
    cv2.circle(self.im, centroids[1], 6, centroid_color, -1)

def plot_workout_information(
    self,
    display_text: str,
    position: Tuple[int, int],
    color: Tuple[int, int, int] = (104, 31, 17),
    txt_color: Tuple[int, int, int] = (255, 255, 255),
) -> int:
    """
    Draw workout text with a background on the image.

    Args:
        display_text (str): The text to be displayed.
        position (Tuple[int, int]): Coordinates (x, y) on the image where the text will be placed.
        color (Tuple[int, int, int]): Text background color.
        txt_color (Tuple[int, int, int]): Text foreground color.

    Returns:
        (int): The height of the text.
    """
    (text_width, text_height), _ = cv2.getTextSize(display_text, 0, fontScale=self.sf, thickness=self.tf)

    # Draw background rectangle
    cv2.rectangle(
        self.im,
        (position[0], position[1] - text_height - 5),
        (position[0] + text_width + 10, position[1] - text_height - 5 + text_height + 10 + self.tf),
        color,
        -1,
    )
    # Draw text
    cv2.putText(self.im, display_text, position, 0, self.sf, txt_color, self.tf)

    return text_height

def queue_counts_display(
    self,
    label: str,
    points: Optional[List[Tuple[int, int]]] = None,
    region_color: Tuple[int, int, int] = (255, 255, 255),
    txt_color: Tuple[int, int, int] = (0, 0, 0),
):
    """
    Display queue counts on an image centered at the points with customizable font size and colors.

    Args:
        label (str): Queue counts label.
        points (List[Tuple[int, int]], optional): Region points for center point calculation to display text.
        region_color (Tuple[int, int, int]): RGB queue region color.
        txt_color (Tuple[int, int, int]): RGB text display color.
    """
    x_values = [point[0] for point in points]
    y_values = [point[1] for point in points]
    center_x = sum(x_values) // len(points)
    center_y = sum(y_values) // len(points)

    text_size = cv2.getTextSize(label, 0, fontScale=self.sf, thickness=self.tf)[0]
    text_width = text_size[0]
    text_height = text_size[1]

    rect_width = text_width + 20
    rect_height = text_height + 20
    rect_top_left = (center_x - rect_width // 2, center_y - rect_height // 2)
    rect_bottom_right = (center_x + rect_width // 2, center_y + rect_height // 2)
    cv2.rectangle(self.im, rect_top_left, rect_bottom_right, region_color, -1)

    text_x = center_x - text_width // 2
    text_y = center_y + text_height // 2

    # Draw text
    cv2.putText(
        self.im,
        label,
        (text_x, text_y),
        0,
        fontScale=self.sf,
        color=txt_color,
        thickness=self.tf,
        lineType=cv2.LINE_AA,
    )

def sweep_annotator(
    self,
    line_x: int = 0,
    line_y: int = 0,
    label: Optional[str] = None,
    color: Tuple[int, int, int] = (221, 0, 186),
    txt_color: Tuple[int, int, int] = (255, 255, 255),
):
    """
    Draw a sweep annotation line and an optional label.

    Args:
        line_x (int): The x-coordinate of the sweep line.
        line_y (int): The y-coordinate limit of the sweep line.
        label (str, optional): Text label to be drawn in center of sweep line. If None, no label is drawn.
        color (Tuple[int, int, int]): RGB color for the line and label background.
        txt_color (Tuple[int, int, int]): RGB color for the label text.
    """
    # Draw the sweep line
    cv2.line(self.im, (line_x, 0), (line_x, line_y), color, self.tf * 2)

    # Draw label, if provided
    if label:
        (text_width, text_height), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, self.sf, self.tf)
        cv2.rectangle(
            self.im,
            (line_x - text_width // 2 - 10, line_y // 2 - text_height // 2 - 10),
            (line_x + text_width // 2 + 10, line_y // 2 + text_height // 2 + 10),
            color,
            -1,
        )
        cv2.putText(
            self.im,
            label,
            (line_x - text_width // 2, line_y // 2 + text_height // 2),
            cv2.FONT_HERSHEY_SIMPLEX,
            self.sf,
            txt_color,
            self.tf,
        )