HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

Update distance-calculation solution (#16907)

Browse source 
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
This commit is contained in:
Muhammad Rizwan Munawar 2024-10-14 19:44:36 +05:00 • committed by GitHub
parent 1052cf41f8
commit 8e3846d377
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 41 additions and 106 deletions
Download patch file Download diff file Expand all files Collapse all files

31
docs/en/guides/distance-calculation.md
View file

@ -43,12 +43,9 @@ Measuring the gap between two objects is known as distance calculation within a
```python
import cv2
from ultralytics import YOLO, solutions
from ultralytics import solutions
model = YOLO("yolo11n.pt")
names = model.model.names
cap = cv2.VideoCapture("path/to/video/file.mp4")
cap = cv2.VideoCapture("Path/to/video/file.mp4")
assert cap.isOpened(), "Error reading video file"
w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))
@ -56,16 +53,14 @@ Measuring the gap between two objects is known as distance calculation within a
video_writer = cv2.VideoWriter("distance_calculation.avi", cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
# Init distance-calculation obj
dist_obj = solutions.DistanceCalculation(names=names, view_img=True)
distance = solutions.DistanceCalculation(model="yolo11n.pt", show=True)
while cap.isOpened():
success, im0 = cap.read()
if not success:
print("Video frame is empty or video processing has been successfully completed.")
break
tracks = model.track(im0, persist=True, show=False)
im0 = dist_obj.start_process(im0, tracks)
im0 = distance.calculate(im0)
video_writer.write(im0)
cap.release()
@ -85,12 +80,10 @@ Measuring the gap between two objects is known as distance calculation within a
### Arguments `DistanceCalculation()`
| `Name` | `Type` | `Default` | Description |
| ---------------- | ------- | --------------- | --------------------------------------------------------- |
| `names` | `dict` | `None` | Dictionary of classes names. |
| `view_img` | `bool` | `False` | Flag to indicate if the video stream should be displayed. |
| `line_thickness` | `int` | `2` | Thickness of the lines drawn on the image. |
| `line_color` | `tuple` | `(255, 255, 0)` | Color of the lines drawn on the image (BGR format). |
| `centroid_color` | `tuple` | `(255, 0, 255)` | Color of the centroids drawn (BGR format). |
| ------------ | ------ | --------- | ---------------------------------------------------- |
| `model` | `str` | `None` | Path to Ultralytics YOLO Model File |
| `line_width` | `int` | `2` | Line thickness for bounding boxes. |
| `show` | `bool` | `False` | Flag to control whether to display the video stream. |
### Arguments `model.track`
@ -122,10 +115,8 @@ To delete points drawn during distance calculation with Ultralytics YOLO11, you
The key arguments for initializing the `DistanceCalculation` class in Ultralytics YOLO11 include:
- `names`: Dictionary mapping class indices to class names.
- `view_img`: Flag to indicate if the video stream should be displayed.
- `line_thickness`: Thickness of the lines drawn on the image.
- `line_color`: Color of the lines drawn on the image (BGR format).
- `centroid_color`: Color of the centroids (BGR format).
- `model`: Model file path.
- `show`: Flag to indicate if the video stream should be displayed.
- `line_width`: Thickness of bounding box and the lines drawn on the image.
For an exhaustive list and default values, see the [arguments of DistanceCalculation](#arguments-distancecalculation).

1
docs/en/guides/heatmaps.md
View file

@ -222,6 +222,7 @@ A heatmap generated with [Ultralytics YOLO11](https://github.com/ultralytics/ult
| Name | Type | Default | Description |
| ------------ | ------ | ------------------ | ----------------------------------------------------------------- |
| `model` | `str` | `None` | Path to Ultralytics YOLO Model File |
| `colormap` | `int` | `cv2.COLORMAP_JET` | Colormap to use for the heatmap. |
| `show` | `bool` | `False` | Whether to display the image with the heatmap overlay. |
| `show_in` | `bool` | `True` | Whether to display the count of objects entering the region. |

1
docs/en/guides/workouts-monitoring.md
View file

@ -106,6 +106,7 @@ Monitoring workouts through pose estimation with [Ultralytics YOLO11](https://gi
| `show` | `bool` | `False` | Flag to display the image. |
| `up_angle` | `float` | `145.0` | Angle threshold for the 'up' pose. |
| `down_angle` | `float` | `90.0` | Angle threshold for the 'down' pose. |
| `model` | `str` | `None` | Path to Ultralytics YOLO Pose Model File |
### Arguments `model.predict`

87
ultralytics/solutions/distance_calculation.py
View file

@ -4,55 +4,21 @@ import math
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, colors
class DistanceCalculation:
class DistanceCalculation(BaseSolution):
"""A class to calculate distance between two objects in a real-time video stream based on their tracks."""
def __init__(
self,
names,
view_img=False,
line_thickness=2,
line_color=(255, 0, 255),
centroid_color=(104, 31, 17),
):
"""
Initializes the DistanceCalculation class with the given parameters.
Args:
names (dict): Dictionary of classes names.
view_img (bool, optional): Flag to indicate if the video stream should be displayed. Defaults to False.
line_thickness (int, optional): Thickness of the lines drawn on the image. Defaults to 2.
line_color (tuple, optional): Color of the lines drawn on the image (BGR format). Defaults to (255, 255, 0).
centroid_color (tuple, optional): Color of the centroids drawn (BGR format). Defaults to (255, 0, 255).
"""
# Visual & image information
self.im0 = None
self.annotator = None
self.view_img = view_img
self.line_color = line_color
self.centroid_color = centroid_color
# Prediction & tracking information
self.names = names
self.boxes = None
self.line_thickness = line_thickness
self.trk_ids = None
# Distance calculation information
self.centroids = []
def __init__(self, **kwargs):
"""Initializes the DistanceCalculation class with the given parameters."""
super().__init__(**kwargs)
# Mouse event information
self.left_mouse_count = 0
self.selected_boxes = {}
# Check if environment supports imshow
self.env_check = check_imshow(warn=True)
self.window_name = "Ultralytics Solutions"
def mouse_event_for_distance(self, event, x, y, flags, param):
"""
Handles mouse events to select regions in a real-time video stream.
@ -67,7 +33,7 @@ class DistanceCalculation:
if event == cv2.EVENT_LBUTTONDOWN:
self.left_mouse_count += 1
if self.left_mouse_count <= 2:
for box, track_id in zip(self.boxes, self.trk_ids):
for box, track_id in zip(self.boxes, self.track_ids):
if box[0] < x < box[2] and box[1] < y < box[3] and track_id not in self.selected_boxes:
self.selected_boxes[track_id] = box
@ -75,30 +41,21 @@ class DistanceCalculation:
self.selected_boxes = {}
self.left_mouse_count = 0
def start_process(self, im0, tracks):
def calculate(self, im0):
"""
Processes the video frame and calculates the distance between two bounding boxes.
Args:
im0 (ndarray): The image frame.
tracks (list): List of tracks obtained from the object tracking process.
Returns:
(ndarray): The processed image frame.
"""
self.im0 = im0
if tracks[0].boxes.id is None:
if self.view_img:
self.display_frames()
return im0
self.annotator = Annotator(im0, line_width=self.line_width) # Initialize annotator
self.extract_tracks(im0) # Extract tracks
self.boxes = tracks[0].boxes.xyxy.cpu()
clss = tracks[0].boxes.cls.cpu().tolist()
self.trk_ids = tracks[0].boxes.id.int().cpu().tolist()
self.annotator = Annotator(self.im0, line_width=self.line_thickness)
for box, cls, track_id in zip(self.boxes, clss, self.trk_ids):
# Iterate over bounding boxes, track ids and classes index
for box, track_id, cls in zip(self.boxes, self.track_ids, self.clss):
self.annotator.box_label(box, color=colors(int(cls), True), label=self.names[int(cls)])
if len(self.selected_boxes) == 2:
@ -115,25 +72,11 @@ class DistanceCalculation:
pixels_distance = math.sqrt(
(self.centroids[0][0] - self.centroids[1][0]) ** 2 + (self.centroids[0][1] - self.centroids[1][1]) ** 2
)
self.annotator.plot_distance_and_line(pixels_distance, self.centroids, self.line_color, self.centroid_color)
self.annotator.plot_distance_and_line(pixels_distance, self.centroids)
self.centroids = []
if self.view_img and self.env_check:
self.display_frames()
self.display_output(im0) # display output with base class function
cv2.setMouseCallback("Ultralytics Solutions", self.mouse_event_for_distance)
return im0
def display_frames(self):
"""Displays the current frame with annotations."""
cv2.namedWindow(self.window_name)
cv2.setMouseCallback(self.window_name, self.mouse_event_for_distance)
cv2.imshow(self.window_name, self.im0)
if cv2.waitKey(1) & 0xFF == ord("q"):
return
if __name__ == "__main__":
names = {0: "person", 1: "car"} # example class names
distance_calculation = DistanceCalculation(names)
return im0 # return output image for more usage

4
ultralytics/solutions/object_counter.py
View file

@ -112,13 +112,13 @@ class ObjectCounter(BaseSolution):
# Iterate over bounding boxes, track ids and classes index
for box, track_id, cls in zip(self.boxes, self.track_ids, self.clss):
# Draw bounding box and counting region
self.annotator.box_label(box, label=self.names[cls], color=colors(track_id, True))
self.annotator.box_label(box, label=self.names[cls], color=colors(cls, True))
self.store_tracking_history(track_id, box) # Store track history
self.store_classwise_counts(cls) # store classwise counts in dict
# Draw tracks of objects
self.annotator.draw_centroid_and_tracks(
self.track_line, color=colors(int(track_id), True), track_thickness=self.line_width
self.track_line, color=colors(int(cls), True), track_thickness=self.line_width
)
# store previous position of track for object counting

21
ultralytics/utils/plotting.py
View file

@ -804,31 +804,30 @@ class Annotator:
self.im, label, (int(mask[0][0]) - text_size[0] // 2, int(mask[0][1])), 0, self.sf, txt_color, self.tf
)
def plot_distance_and_line(self, pixels_distance, centroids, line_color, centroid_color):
def plot_distance_and_line(
self, pixels_distance, centroids, line_color=(104, 31, 17), centroid_color=(255, 0, 255)
):
"""
Plot the distance and line on frame.
Args:
pixels_distance (float): Pixels distance between two bbox centroids.
centroids (list): Bounding box centroids data.
line_color (tuple): RGB distance line color.
centroid_color (tuple): RGB bounding box centroid color.
line_color (tuple, optional): Distance line color.
centroid_color (tuple, optional): Bounding box centroid color.
"""
# Get the text size
(text_width_m, text_height_m), _ = cv2.getTextSize(
f"Pixels Distance: {pixels_distance:.2f}", 0, self.sf, self.tf
)
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
top_left = (15, 25)
bottom_right = (15 + text_width_m + 20, 25 + text_height_m + 20)
cv2.rectangle(self.im, top_left, bottom_right, centroid_color, -1)
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 = (top_left[0] + 10, top_left[1] + text_height_m + 10)
text_position = (25, 25 + text_height_m + 10)
cv2.putText(
self.im,
f"Pixels Distance: {pixels_distance:.2f}",
text,
text_position,
0,
self.sf,