Segmentation problem

[ana111todorova]

Search before asking

• I have searched the Supervision issues and found no similar feature requests.

Question

Dear @SkalskiP

I am trying to adapt your code for velocity estimation on cars, so that besides detection can display segmentation also.
I chaned the model into model = YOLO('yolov8s-seg.pt'), but there is no change in the result.

Bellow is the code, that I just want to add segmentation on it so that I can use it for my purposes:

import argparse
from collections import defaultdict, deque

import cv2
import numpy as np
from ultralytics import YOLO
from tqdm.notebook import tqdm

import supervision as sv

SOURCE = np.array([[1252, 787], [2298, 803], [5039, 2159], [-550, 2159]])           #coordinates of the tracking region)

TARGET_WIDTH = 25                                                                   #physical dimentions of the targeted region)
TARGET_HEIGHT = 250


TARGET = np.array(
    [
        [0, 0],
        [TARGET_WIDTH - 1, 0],
        [TARGET_WIDTH - 1, TARGET_HEIGHT - 1],                                       #targeted region in coordinates
        [0, TARGET_HEIGHT - 1],
    ]
)

LINE_START = sv.Point(50, 1500)
LINE_END = sv.Point(3840-50, 1500)



class ViewTransformer:
    def __init__(self, source: np.ndarray, target: np.ndarray) -> None:
        source = source.astype(np.float32)
        target = target.astype(np.float32)
        self.m = cv2.getPerspectiveTransform(source, target)

    def transform_points(self, points: np.ndarray) -> np.ndarray:                      #transform points from source to the target for the tracking objects
        if points.size == 0:
            return points

        reshaped_points = points.reshape(-1, 1, 2).astype(np.float32)
        transformed_points = cv2.perspectiveTransform(reshaped_points, self.m)
        return transformed_points.reshape(-1, 2)


def parse_arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Vehicle Speed Estimation using Ultralytics and Supervision"
    )
    parser.add_argument(
        "--source_video_path",
        required=True,
        help="Path to the source video file",
        type=str,
    )
    parser.add_argument(
        "--target_video_path",
        required=True,
        help="Path to the target video file (output)",                                 #function to upload the source video and the putput !!!! Done from command window!!!!
        type=str,
    )
    parser.add_argument(
        "--confidence_threshold",
        default=0.3,
        help="Confidence threshold for the model",
        type=float,
    )
    parser.add_argument(
        "--iou_threshold", default=0.7, help="IOU threshold for the model", type=float
    )

    return parser.parse_args()


if __name__ == "__main__":
    args = parse_arguments()

    video_info = sv.VideoInfo.from_video_path(video_path=args.source_video_path)      #Uploads the video and uses YOLO
    model = YOLO('yolov8s-seg.pt')
    #model = YOLO("yolov8n-seg.pt")

    byte_track = sv.ByteTrack(
        frame_rate=video_info.fps, track_thresh=args.confidence_threshold             #tracking objects in the video frames
    )

    thickness = sv.calculate_dynamic_line_thickness(                                  #annotating bounding boxes and traces on the frames
        resolution_wh=video_info.resolution_wh
    )
    text_scale = sv.calculate_dynamic_text_scale(resolution_wh=video_info.resolution_wh)
    bounding_box_annotator = sv.BoundingBoxAnnotator(thickness=thickness)
    label_annotator = sv.LabelAnnotator(                                              #labeling with number in the bottom center
        text_scale=text_scale,
        text_thickness=thickness,
        text_position=sv.Position.BOTTOM_CENTER,                                     #label colour not specified and selects a colour based on the label id
    )
    trace_annotator = sv.TraceAnnotator(
        thickness=thickness,
        trace_length=video_info.fps * 2,                                               #tracer display bottom center
        position=sv.Position.BOTTOM_CENTER,color_lookup=sv.ColorLookup.TRACK           #tracer colous changes based on the track number 
    )

    frame_generator = sv.get_video_frames_generator(source_path=args.source_video_path)    #појма немам 

    polygon_zone = sv.PolygonZone(        
        polygon=SOURCE, frame_resolution_wh=video_info.resolution_wh                       #tracer display bottom center
    )
    view_transformer = ViewTransformer(source=SOURCE, target=TARGET)                       #појма немам 

    coordinates = defaultdict(lambda: deque(maxlen=video_info.fps))                        # dictionary corresponds to a tracker ID, and the associated value is a deque (double-ended queue)
                                                                                           # with a maximum length of video_info.fps, which is likely the frames per second of the video.

    line_counter = sv.LineZone(start=LINE_START, end=LINE_END)
    line_annotator = sv.LineZoneAnnotator(thickness=thickness)
    box_annotator = sv.BoxAnnotator(
        thickness=thickness,
        text_thickness=thickness,
        text_scale=text_scale
    )   
     


    with sv.VideoSink(args.target_video_path, video_info) as sink:
        for frame in frame_generator:
            result = model(frame)[0]
            detections = sv.Detections.from_ultralytics(result)
            detections = detections[detections.confidence > args.confidence_threshold]
            detections = detections[polygon_zone.trigger(detections)]                       #detections chack up 
            detections = detections.with_nms(threshold=args.iou_threshold)
            detections = byte_track.update_with_detections(detections=detections)

            points = detections.get_anchors_coordinates(
                anchor=sv.Position.BOTTOM_CENTER
            )
            points = view_transformer.transform_points(points=points).astype(int)

            for tracker_id, [_, y] in zip(detections.tracker_id, points):                   #storing y coordinates in dictenary 
                coordinates[tracker_id].append(y)

            labels = []
            for tracker_id in detections.tracker_id:
                if len(coordinates[tracker_id]) < video_info.fps / 2:
                    labels.append(f"#{tracker_id}")
                else:
                    coordinate_start = coordinates[tracker_id][-1]                           #speed estimation
                    coordinate_end = coordinates[tracker_id][0]
                    distance = abs(coordinate_start - coordinate_end)
                    time = len(coordinates[tracker_id]) / video_info.fps
                    speed = distance / time * 3.6
                    labels.append(f"#{tracker_id} {int(speed)} km/h")

            annotated_frame = frame.copy()
            annotated_frame=sv.draw_polygon(annotated_frame,polygon=SOURCE,color=sv.Color.red())   #draw the poligone red for the detection zone 
            annotated_frame = trace_annotator.annotate(
                scene=annotated_frame, detections=detections
            )
            annotated_frame = bounding_box_annotator.annotate(
                scene=annotated_frame, detections=detections
            )
            annotated_frame = label_annotator.annotate(
                scene=annotated_frame, detections=detections, labels=labels
            )

            line_counter.trigger(detections=detections)
            line_annotator.annotate(frame=annotated_frame, line_counter=line_counter)

                    

            sink.write_frame(annotated_frame)
            cv2.imshow("frame", annotated_frame)
            if cv2.waitKey(1) & 0xFF == ord("q"):                                # displays images, and q is to terminate the loop  CHEERS
                break
        cv2.destroyAllWindows()

Additional

I will be so grateful for your help.

Also, is there a possibility to estimate the area of the segmentation, in real dimensions?

Thank you in advance

[SkalskiP]

Hi @ana111todorova 👋🏻 Let me first convert this issue into discussion and move it into the Q&A section.

[SkalskiP]

Hi @ana111todorova ! 👋🏻 Thanks for your interest in supervision. Swapping the model is not enough. You need to add MaskAnnotator as well to visualize your segmentation.

[ana111todorova]

**Hii,

okey I added this:**

if name == "main":
mask_annotator = sv.MaskAnnotator(opacity=0.9)

with sv.VideoSink(args.target_video_path, video_info) as sink:
for frame in frame_generator:
result = model(frame)[0]
detections = sv.Detections.from_ultralytics(result)

annotated_frame = mask_annotator.annotate(
scene=annotated_frame, detections=detections)

sink.write_frame(annotated_frame)
cv2.imshow("frame", annotated_frame)
if cv2.waitKey(1) & 0xFF == ord("q"):
break
cv2.destroyAllWindows()

However, I am still not able to visualize the segmentation....

What else am I missing?

Thank you in advance