mirror of
https://github.com/ultralytics/ultralytics.git
synced 2025-09-15 15:48:41 +08:00
Improve instance-segmentation mask plotting (#20025)
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com> Co-authored-by: UltralyticsAssistant <web@ultralytics.com> Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Muhammad Rizwan Munawar
GitHub
parent
b9ade810c6
commit
a1d3c94a9d
@ -393,8 +393,8 @@ Ultralytics includes an `Annotator` class for annotating various data types. It'
|
||||
import cv2
|
||||
|
||||
from ultralytics import YOLO
|
||||
from ultralytics.engine.results import Results
|
||||
from ultralytics.solutions.solutions import SolutionAnnotator
|
||||
from ultralytics.utils.plotting import colors
|
||||
|
||||
# User defined video path and model file
|
||||
cap = cv2.VideoCapture("path/to/video.mp4")
|
||||
@ -416,7 +416,8 @@ Ultralytics includes an `Annotator` class for annotating various data types. It'
|
||||
window_name = "Ultralytics Sweep Annotator"
|
||||
|
||||
|
||||
def drag_line(event, x, y, flags, param): # Mouse callback for dragging line.
|
||||
def drag_line(event, x, _, flags, param):
|
||||
"""Mouse callback function to enable dragging a vertical sweep line across the video frame."""
|
||||
global line_x, dragging
|
||||
if event == cv2.EVENT_LBUTTONDOWN or (flags & cv2.EVENT_FLAG_LBUTTON):
|
||||
line_x = max(0, min(x, w))
|
||||
@ -429,44 +430,47 @@ Ultralytics includes an `Annotator` class for annotating various data types. It'
|
||||
break
|
||||
f = f + 1 # Increment frame count.
|
||||
count = 0 # Re-initialize count variable on every frame for precise counts.
|
||||
annotator = SolutionAnnotator(im0)
|
||||
results = model.track(im0, persist=True) # Track objects using track method.
|
||||
results = model.track(im0, persist=True)[0]
|
||||
|
||||
if f == 1:
|
||||
cv2.namedWindow(window_name)
|
||||
cv2.setMouseCallback(window_name, drag_line)
|
||||
|
||||
if results[0].boxes.id is not None:
|
||||
if results[0].masks is not None:
|
||||
masks = results[0].masks.xy
|
||||
track_ids = results[0].boxes.id.int().cpu().tolist()
|
||||
clss = results[0].boxes.cls.cpu().tolist()
|
||||
boxes = results[0].boxes.xyxy.cpu()
|
||||
if results.boxes.id is not None:
|
||||
if results.masks is not None:
|
||||
masks = results.masks
|
||||
boxes = results.boxes
|
||||
|
||||
track_ids = results.boxes.id.int().cpu().tolist()
|
||||
clss = results.boxes.cls.cpu().tolist()
|
||||
|
||||
for mask, box, cls, t_id in zip(masks or [None] * len(boxes), boxes, clss, track_ids):
|
||||
color = colors(t_id, True) # Assign different color to each tracked object.
|
||||
if mask is not None and mask.size > 0:
|
||||
# If you want to overlay the masks
|
||||
# mask[:, 0] = np.clip(mask[:, 0], line_x, w)
|
||||
# mask_img = cv2.fillPoly(im0.copy(), [mask.astype(int)], color)
|
||||
# cv2.addWeighted(mask_img, 0.5, im0, 0.5, 0, im0)
|
||||
box_data = box.xyxy.cpu().tolist()[0][0]
|
||||
if box_data > line_x:
|
||||
count += 1
|
||||
results = Results(
|
||||
im0, path=None, names=classes, boxes=box.data, masks=None if mask is None else mask.data
|
||||
)
|
||||
im0 = results.plot(
|
||||
boxes=True, # display bounding box
|
||||
conf=False, # hide confidence score
|
||||
labels=True, # display labels
|
||||
color_mode="instance",
|
||||
)
|
||||
|
||||
if box[0] > line_x:
|
||||
count += 1
|
||||
annotator.seg_bbox(mask=mask, mask_color=color, label=str(classes[cls]))
|
||||
else:
|
||||
if box[0] > line_x:
|
||||
count += 1
|
||||
annotator.box_label(box=box, color=color, label=str(classes[cls]))
|
||||
# Generate draggable sweep line
|
||||
annotator = SolutionAnnotator(im0)
|
||||
annotator.sweep_annotator(line_x=line_x, line_y=h, label=f"COUNT:{count}")
|
||||
|
||||
annotator.sweep_annotator(line_x=line_x, line_y=h, label=f"COUNT:{count}") # Display the sweep
|
||||
cv2.imshow(window_name, im0)
|
||||
video_writer.write(im0)
|
||||
if cv2.waitKey(1) & 0xFF == ord("q"):
|
||||
break
|
||||
|
||||
cap.release() # Release the video capture.
|
||||
video_writer.release() # Release the video writer.
|
||||
cv2.destroyAllWindows() # Destroy all opened windows.
|
||||
# Release the resources
|
||||
cap.release()
|
||||
video_writer.release()
|
||||
cv2.destroyAllWindows()
|
||||
```
|
||||
|
||||
#### Horizontal Bounding Boxes
|
||||
|
||||
@ -1,7 +1,7 @@
|
||||
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
||||
|
||||
from ultralytics.solutions.solutions import BaseSolution, SolutionAnnotator, SolutionResults
|
||||
from ultralytics.utils.plotting import colors
|
||||
from ultralytics.engine.results import Results
|
||||
from ultralytics.solutions.solutions import BaseSolution, SolutionResults
|
||||
|
||||
|
||||
class InstanceSegmentation(BaseSolution):
|
||||
@ -41,6 +41,10 @@ class InstanceSegmentation(BaseSolution):
|
||||
kwargs["model"] = kwargs.get("model", "yolo11n-seg.pt")
|
||||
super().__init__(**kwargs)
|
||||
|
||||
self.show_conf = self.CFG.get("show_conf", True)
|
||||
self.show_labels = self.CFG.get("show_labels", True)
|
||||
self.show_boxes = self.CFG.get("show_boxes", True)
|
||||
|
||||
def process(self, im0):
|
||||
"""
|
||||
Perform instance segmentation on the input image and annotate the results.
|
||||
@ -58,17 +62,21 @@ class InstanceSegmentation(BaseSolution):
|
||||
>>> print(summary)
|
||||
"""
|
||||
self.extract_tracks(im0) # Extract tracks (bounding boxes, classes, and masks)
|
||||
annotator = SolutionAnnotator(im0, self.line_width)
|
||||
|
||||
# Iterate over detected classes, track IDs, and segmentation masks
|
||||
if self.masks is None:
|
||||
self.LOGGER.warning("⚠️ No masks detected! Ensure you're using a supported Ultralytics segmentation model.")
|
||||
plot_im = im0
|
||||
else:
|
||||
for cls, t_id, mask in zip(self.clss, self.track_ids, self.masks):
|
||||
# Annotate the image with segmentation mask, mask color, and label
|
||||
annotator.segmentation_mask(mask=mask, mask_color=colors(t_id, True), label=self.names[cls])
|
||||
results = Results(im0, path=None, names=self.names, boxes=self.track_data.data, masks=self.masks.data)
|
||||
plot_im = results.plot(
|
||||
line_width=self.line_width,
|
||||
boxes=self.show_boxes,
|
||||
conf=self.show_conf,
|
||||
labels=self.show_labels,
|
||||
color_mode="instance",
|
||||
)
|
||||
|
||||
plot_im = annotator.result()
|
||||
self.display_output(plot_im) # Display the annotated output using the base class function
|
||||
|
||||
# Return SolutionResults
|
||||
|
||||
@ -122,7 +122,7 @@ class BaseSolution:
|
||||
self.track_data = self.tracks[0].obb or self.tracks[0].boxes # Extract tracks for OBB or object detection
|
||||
|
||||
self.masks = (
|
||||
self.tracks[0].masks.xy if hasattr(self.tracks[0], "masks") and self.tracks[0].masks is not None else None
|
||||
self.tracks[0].masks if hasattr(self.tracks[0], "masks") and self.tracks[0].masks is not None else None
|
||||
)
|
||||
|
||||
if self.track_data and self.track_data.id is not None:
|
||||
@ -225,7 +225,6 @@ class SolutionAnnotator(Annotator):
|
||||
plot_angle_and_count_and_stage: Visualizes angle, step count, and stage for workout monitoring.
|
||||
plot_distance_and_line: Displays the distance between centroids and connects them with a line.
|
||||
display_objects_labels: Annotates bounding boxes with object class labels.
|
||||
segmentation_mask: Draws mask for segmented objects and optionally labels them.
|
||||
sweep_annotator: Visualizes a vertical sweep line and optional label.
|
||||
visioneye: Maps and connects object centroids to a visual "eye" point.
|
||||
circle_label: Draws a circular label within a bounding box.
|
||||
@ -519,50 +518,6 @@ class SolutionAnnotator(Annotator):
|
||||
lineType=cv2.LINE_AA,
|
||||
)
|
||||
|
||||
def segmentation_mask(self, mask, mask_color=(255, 0, 255), label=None, alpha=0.5):
|
||||
"""
|
||||
Draw an optimized segmentation mask with smooth corners, highlighted edge, and dynamic text box size.
|
||||
|
||||
Args:
|
||||
mask (np.ndarray): A 2D array of shape (N, 2) containing the object mask.
|
||||
mask_color (Tuple[int, int, int]): RGB color for the mask.
|
||||
label (str, optional): Text label for the object.
|
||||
alpha (float): Transparency level (0 = fully transparent, 1 = fully opaque).
|
||||
"""
|
||||
if mask.size == 0:
|
||||
return
|
||||
|
||||
overlay = self.im.copy()
|
||||
mask = np.int32([mask])
|
||||
|
||||
# Approximate polygon for smooth corners with epsilon
|
||||
refined_mask = cv2.approxPolyDP(mask, 0.002 * cv2.arcLength(mask, True), True)
|
||||
|
||||
# Apply a highlighter effect by drawing a thick outer shadow
|
||||
cv2.polylines(overlay, [refined_mask], isClosed=True, color=mask_color, thickness=self.lw * 3)
|
||||
cv2.fillPoly(overlay, [refined_mask], mask_color) # draw mask with primary color
|
||||
|
||||
# Apply an inner glow effect for extra clarity
|
||||
cv2.polylines(overlay, [refined_mask], isClosed=True, color=mask_color, thickness=self.lw)
|
||||
|
||||
self.im = cv2.addWeighted(overlay, alpha, self.im, 1 - alpha, 0) # blend overlay with the original image
|
||||
|
||||
# Draw label if provided
|
||||
if label:
|
||||
text_size, _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, self.sf, self.tf)
|
||||
text_x, text_y = refined_mask[0][0][0], refined_mask[0][0][1]
|
||||
rect_start, rect_end = (text_x - 5, text_y - text_size[1] - 5), (text_x + text_size[0] + 5, text_y + 5)
|
||||
cv2.rectangle(self.im, rect_start, rect_end, mask_color, -1)
|
||||
cv2.putText(
|
||||
self.im,
|
||||
label,
|
||||
(text_x, text_y),
|
||||
cv2.FONT_HERSHEY_SIMPLEX,
|
||||
self.sf,
|
||||
self.get_txt_color(mask_color),
|
||||
self.tf,
|
||||
)
|
||||
|
||||
def sweep_annotator(self, line_x=0, line_y=0, label=None, color=(221, 0, 186), txt_color=(255, 255, 255)):
|
||||
"""
|
||||
Draw a sweep annotation line and an optional label.
|
||||
|
||||
Loading…
Reference in New Issue
Block a user