Chengfang Lu
e72bc061c5
Major changes: - New frontend (platform/web/): Vite + React 18 + TypeScript + Tailwind - 4-module navigation: 数据送标 / 模型管理 / 车队管理 / 系统管理 - Data catalog with charts (DMS/ADAS/Lane 3-tab view) - Quality review workflow (标注质检): Good/Fine/Bad scoring with auto-advance - Audit enhancements: batch operations, rejection categories, Feishu notifications - Operation audit log (操作日志) - World model simulation studio (仿真工坊) - Dataset version management with snapshots and diff - ADAS 7-class dataset integration (138K images organized + compressed) - User management with Feishu integration and pagination - CRUD/search/filter on all pages, card layout redesign - PIL-optimized image overlay rendering - Auto-snapshot on build, in_review workflow stage - Removed embedded algorithm code (now in workspace)
87 lines
3.7 KiB
Python
87 lines
3.7 KiB
Python
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
|
|
|
from typing import Any
|
|
|
|
from ultralytics.engine.results import Results
|
|
from ultralytics.solutions.solutions import BaseSolution, SolutionResults
|
|
|
|
|
|
class InstanceSegmentation(BaseSolution):
|
|
"""A class to manage instance segmentation in images or video streams.
|
|
|
|
This class extends the BaseSolution class and provides functionality for performing instance segmentation, including
|
|
drawing segmented masks with bounding boxes and labels.
|
|
|
|
Attributes:
|
|
model (str): The segmentation model to use for inference.
|
|
line_width (int): Width of the bounding box and text lines.
|
|
names (dict[int, str]): Dictionary mapping class indices to class names.
|
|
clss (list[int]): List of detected class indices.
|
|
track_ids (list[int]): List of track IDs for detected instances.
|
|
masks (list[np.ndarray]): List of segmentation masks for detected instances.
|
|
show_conf (bool): Whether to display confidence scores.
|
|
show_labels (bool): Whether to display class labels.
|
|
show_boxes (bool): Whether to display bounding boxes.
|
|
|
|
Methods:
|
|
process: Process the input image to perform instance segmentation and annotate results.
|
|
extract_tracks: Extract tracks including bounding boxes, classes, and masks from model predictions.
|
|
|
|
Examples:
|
|
>>> segmenter = InstanceSegmentation()
|
|
>>> frame = cv2.imread("frame.jpg")
|
|
>>> results = segmenter.process(frame)
|
|
>>> print(f"Total segmented instances: {results.total_tracks}")
|
|
"""
|
|
|
|
def __init__(self, **kwargs: Any) -> None:
|
|
"""Initialize the InstanceSegmentation class for detecting and annotating segmented instances.
|
|
|
|
Args:
|
|
**kwargs (Any): Keyword arguments passed to the BaseSolution parent class including:
|
|
- model (str): Model name or path, defaults to "yolo26n-seg.pt".
|
|
"""
|
|
kwargs["model"] = kwargs.get("model", "yolo26n-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) -> SolutionResults:
|
|
"""Perform instance segmentation on the input image and annotate the results.
|
|
|
|
Args:
|
|
im0 (np.ndarray): The input image for segmentation.
|
|
|
|
Returns:
|
|
(SolutionResults): Object containing the annotated image and total number of tracked instances.
|
|
|
|
Examples:
|
|
>>> segmenter = InstanceSegmentation()
|
|
>>> frame = cv2.imread("image.jpg")
|
|
>>> summary = segmenter.process(frame)
|
|
>>> print(summary)
|
|
"""
|
|
self.extract_tracks(im0) # Extract tracks (bounding boxes, classes, and masks)
|
|
self.masks = getattr(self.tracks, "masks", None)
|
|
|
|
# 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:
|
|
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",
|
|
)
|
|
|
|
self.display_output(plot_im) # Display the annotated output using the base class function
|
|
|
|
# Return SolutionResults
|
|
return SolutionResults(plot_im=plot_im, total_tracks=len(self.track_ids))
|