Counting based on different Categories of Fishes ( Instead of In and out)

[NandiniLReddy]

Search before asking

• I have searched the Roboflow Notebooks issues and found no similar bug report.

Notebook name

how-to-track-and-count-vehicles-with-yolov8-and-supervison.ipynb

Bug

I have been working on counting based on different categories of fish on a conveyor belt; I have modified your LineCounter to make it customized.
`# Custom Linecounter

from typing import Dict

import cv2
import numpy as np

from supervision.draw.color import Color
from supervision.geometry.dataclasses import Point, Rect, Vector

from supervision import Detections

class LineCounter1:
"""
Count the number of objects that cross a line.
"""

def __init__(self, start: Point, end: Point):
    """
    Initialize a LineCounter object.

    Attributes:
        start (Point): The starting point of the line.
        end (Point): The ending point of the line.

    """
    self.vector = Vector(start=start, end=end)
    self.tracker_state: Dict[str, bool] = {}
    self.fl_witch: int = 0
    self.haddock: int = 0
    self.hk_red: int = 0
    self.hk_slvr: int = 0
    self.pollock: int = 0
    self.redfish: int = 0
    self.in_count: int = 0
    self.out_count: int = 0


def trigger(self, detections: Detections):
    """
    Update the in_count and out_count for the detections that cross the line.

    Attributes:
        detections (Detections): The detections for which to update the counts.

    """
    for xyxy, confidence, class_id, tracker_id, length, gt_length in detections:
        # handle detections with no tracker_id
        if tracker_id is None:
            continue

        # we check if all four anchors of bbox are on the same side of vector
        x1, y1, x2, y2 = xyxy
        anchors = [
            Point(x=x1, y=y1),
            Point(x=x1, y=y2),
            Point(x=x2, y=y1),
            Point(x=x2, y=y2),
        ]
        triggers = [self.vector.is_in(point=anchor) for anchor in anchors]

        # detection is partially in and partially out
        if len(set(triggers)) == 2:
            continue

        tracker_state = triggers[0]
        # handle new detection
        if tracker_id not in self.tracker_state:
            self.tracker_state[tracker_id] = tracker_state
            continue

        # handle detection on the same side of the line
        if self.tracker_state.get(tracker_id) == tracker_state:
            continue

        self.tracker_state[tracker_id] = tracker_state
        if tracker_state:
            #lets put conditions to update fish count based on class_id
            if class_id == 0 or class_id == 1:
                self.fl_witch += 1
            if class_id == 2:
                self.haddock += 1
            if class_id == 3:
                self.hk_red += 1
            if class_id == 4:
                self.hk_slvr += 1
            if class_id == 5:
                self.pollock += 1
            if class_id == 6:
                self.redfish += 1
        else:
            continue

class LineCounterAnnotator1:
def init(
self,
thickness: float = 2,
color: Color = Color.white(),
text_thickness: float = 2,
text_color: Color = Color.black(),
text_scale: float = 0.5,
text_offset: float = 1.5,
text_padding: int = 10,
):
"""
Initialize the LineCounterAnnotator object with default values.

    Attributes:
        thickness (float): The thickness of the line that will be drawn.
        color (Color): The color of the line that will be drawn.
        text_thickness (float): The thickness of the text that will be drawn.
        text_color (Color): The color of the text that will be drawn.
        text_scale (float): The scale of the text that will be drawn.
        text_offset (float): The offset of the text that will be drawn.
        text_padding (int): The padding of the text that will be drawn.

    """
    self.thickness: float = thickness
    self.color: Color = color
    self.text_thickness: float = text_thickness
    self.text_color: Color = text_color
    self.text_scale: float = text_scale
    self.text_offset: float = text_offset
    self.text_padding: int = text_padding

def annotate(self, frame: np.ndarray, line_counter: LineCounter1) -> np.ndarray:
    """
    Draws the line on the frame using the line_counter provided.

    Attributes:
        frame (np.ndarray): The image on which the line will be drawn.
        line_counter (LineCounter): The line counter that will be used to draw the line.

    Returns:
        np.ndarray: The image with the line drawn on it.

    """
    cv2.line(
        frame,
        line_counter.vector.start.as_xy_int_tuple(),
        line_counter.vector.end.as_xy_int_tuple(),
        self.color.as_bgr(),
        self.thickness,
        lineType=cv2.LINE_AA,
        shift=0,
    )
    cv2.circle(
        frame,
        line_counter.vector.start.as_xy_int_tuple(),
        radius=5,
        color=self.text_color.as_bgr(),
        thickness=-1,
        lineType=cv2.LINE_AA,
    )
    cv2.circle(
        frame,
        line_counter.vector.end.as_xy_int_tuple(),
        radius=5,
        color=self.text_color.as_bgr(),
        thickness=-1,
        lineType=cv2.LINE_AA,
    )

    # Fish texts
    fl_witch_text = f"Flounder Witch: {line_counter.fl_witch}"
    haddock_text = f"Haddock: {line_counter.haddock}"
    hk_red_text = f"Hake Red: {line_counter.hk_red}"
    hk_slvr_text = f"Hake Silver: {line_counter.hk_slvr}"
    pollock_text = f"Pollock: {line_counter.pollock}"
    redfish_text = f"Red Fish: {line_counter.redfish}"

    # Getting text sizes
    (fl_witch_text_width, fl_witch_text_height), _ = cv2.getTextSize(
        fl_witch_text, cv2.FONT_HERSHEY_SIMPLEX, self.text_scale, self.text_thickness
    )
    (haddock_text_width, haddock_text_height), _ = cv2.getTextSize(
        haddock_text, cv2.FONT_HERSHEY_SIMPLEX, self.text_scale, self.text_thickness
    )
    (hk_red_text_width, hk_red_text_height), _ = cv2.getTextSize(
        hk_red_text, cv2.FONT_HERSHEY_SIMPLEX, self.text_scale, self.text_thickness
    )
    (hk_slvr_text_width, hk_slvr_text_height), _ = cv2.getTextSize(
        hk_slvr_text, cv2.FONT_HERSHEY_SIMPLEX, self.text_scale, self.text_thickness
    )
    (pollock_text_width, pollock_text_height), _ = cv2.getTextSize(
        pollock_text, cv2.FONT_HERSHEY_SIMPLEX, self.text_scale, self.text_thickness
    )
    (redfish_text_width, redfish_text_height), _ = cv2.getTextSize(
        redfish_text, cv2.FONT_HERSHEY_SIMPLEX, self.text_scale, self.text_thickness
    )

    
    # Defining x and y coordinates
    fl_witch_text_x = int(
        (line_counter.vector.end.x + line_counter.vector.start.x - fl_witch_text_width)
        / 2
    )
    fl_witch_text_y = int(
        (line_counter.vector.end.y + line_counter.vector.start.y + fl_witch_text_height)
        / 2
        + self.text_offset * fl_witch_text_height * 6
    )

    haddock_text_x = int(
        (line_counter.vector.end.x + line_counter.vector.start.x - haddock_text_width)
        / 2
    )
    haddock_text_y = int(
        (line_counter.vector.end.y + line_counter.vector.start.y + haddock_text_height)
        / 2
        + self.text_offset * haddock_text_height * 4
    )

    hk_red_text_x = int(
        (line_counter.vector.end.x + line_counter.vector.start.x - hk_red_text_width)
        / 2
    )
    hk_red_text_y = int(
        (line_counter.vector.end.y + line_counter.vector.start.y + hk_red_text_height)
        / 2
        + self.text_offset * hk_red_text_height
    )

    hk_slvr_text_x = int(
        (line_counter.vector.end.x + line_counter.vector.start.x - hk_slvr_text_width)
        / 2
    )
    hk_slvr_text_y = int(
        (line_counter.vector.end.y + line_counter.vector.start.y + hk_slvr_text_height)
        / 2
        - self.text_offset * hk_slvr_text_height
    )

    pollock_text_x = int(
        (line_counter.vector.end.x + line_counter.vector.start.x - pollock_text_width)
        / 2
    )
    pollock_text_y = int(
        (line_counter.vector.end.y + line_counter.vector.start.y + pollock_text_height)
        / 2
        - self.text_offset * pollock_text_height * 4
    )

    redfish_text_x = int(
        (line_counter.vector.end.x + line_counter.vector.start.x - redfish_text_width)
        / 2
    )
    redfish_text_y = int(
        (line_counter.vector.end.y + line_counter.vector.start.y + redfish_text_height)
        / 2
        - self.text_offset * redfish_text_height * 6
    )

    # Building rectangles
    fl_witch_text_background_rect = Rect(
        x=fl_witch_text_x,
        y=fl_witch_text_y - fl_witch_text_height,
        width=fl_witch_text_width,
        height=fl_witch_text_height,
    ).pad(padding=self.text_padding)
    haddock_text_background_rect = Rect(
        x=haddock_text_x,
        y=haddock_text_y - haddock_text_height,
        width=haddock_text_width,
        height=haddock_text_height,
    ).pad(padding=self.text_padding)
    hk_red_text_background_rect = Rect(
        x=hk_red_text_x,
        y=hk_red_text_y - hk_red_text_height,
        width=hk_red_text_width,
        height=hk_red_text_height,
    ).pad(padding=self.text_padding)
    hk_slvr_text_background_rect = Rect(
        x=hk_slvr_text_x,
        y=hk_slvr_text_y - hk_slvr_text_height,
        width=hk_slvr_text_width,
        height=hk_slvr_text_height,
    ).pad(padding=self.text_padding)
    pollock_text_background_rect = Rect(
        x=pollock_text_x,
        y=pollock_text_y - pollock_text_height,
        width=pollock_text_width,
        height=pollock_text_height,
    ).pad(padding=self.text_padding)
    redfish_text_background_rect = Rect(
        x=redfish_text_x,
        y=redfish_text_y - redfish_text_height,
        width=redfish_text_width,
        height=redfish_text_height,
    ).pad(padding=self.text_padding)

    # Displaying rectanges
    cv2.rectangle(
        frame,
        fl_witch_text_background_rect.top_left.as_xy_int_tuple(),
        fl_witch_text_background_rect.bottom_right.as_xy_int_tuple(),
        self.color.as_bgr(),
        -1,
    )
    cv2.rectangle(
        frame,
        haddock_text_background_rect.top_left.as_xy_int_tuple(),
        haddock_text_background_rect.bottom_right.as_xy_int_tuple(),
        self.color.as_bgr(),
        -1,
    )
    cv2.rectangle(
        frame,
        hk_red_text_background_rect.top_left.as_xy_int_tuple(),
        hk_red_text_background_rect.bottom_right.as_xy_int_tuple(),
        self.color.as_bgr(),
        -1,
    )
    cv2.rectangle(
        frame,
        hk_slvr_text_background_rect.top_left.as_xy_int_tuple(),
        hk_slvr_text_background_rect.bottom_right.as_xy_int_tuple(),
        self.color.as_bgr(),
        -1,
    )
    cv2.rectangle(
        frame,
        pollock_text_background_rect.top_left.as_xy_int_tuple(),
        pollock_text_background_rect.bottom_right.as_xy_int_tuple(),
        self.color.as_bgr(),
        -1,
    )
    cv2.rectangle(
        frame,
        redfish_text_background_rect.top_left.as_xy_int_tuple(),
        redfish_text_background_rect.bottom_right.as_xy_int_tuple(),
        self.color.as_bgr(),
        -1,
    )

    # Displaying Texts
    cv2.putText(
        frame,
        fl_witch_text,
        (fl_witch_text_x, fl_witch_text_y),
        cv2.FONT_HERSHEY_SIMPLEX,
        self.text_scale,
        Color.green().as_bgr(),
        self.text_thickness,
        cv2.LINE_AA,
    )
    cv2.putText(
        frame,
        haddock_text,
        (haddock_text_x, haddock_text_y),
        cv2.FONT_HERSHEY_SIMPLEX,
        self.text_scale,
        self.text_color.as_bgr(),
        self.text_thickness,
        cv2.LINE_AA,
    )
    cv2.putText(
        frame,
        hk_red_text,
        (hk_red_text_x, hk_red_text_y),
        cv2.FONT_HERSHEY_SIMPLEX,
        self.text_scale,
        Color.blue().as_bgr(),
        self.text_thickness,
        cv2.LINE_AA,
    )
    cv2.putText(
        frame,
        hk_slvr_text,
        (hk_slvr_text_x, hk_slvr_text_y),
        cv2.FONT_HERSHEY_SIMPLEX,
        self.text_scale,
        (0,165,255),
        self.text_thickness,
        cv2.LINE_AA,
    )
    cv2.putText(
        frame,
        pollock_text,
        (pollock_text_x, pollock_text_y),
        cv2.FONT_HERSHEY_SIMPLEX,
        self.text_scale,
        self.text_color.as_bgr(),
        self.text_thickness,
        cv2.LINE_AA,
    )
    cv2.putText(
        frame,
        redfish_text,
        (redfish_text_x, redfish_text_y),
        cv2.FONT_HERSHEY_SIMPLEX,
        self.text_scale,
        Color.red().as_bgr(),
        self.text_thickness,
        cv2.LINE_AA,
    )

`

Environment

python:3.9.0

Minimal Reproducible Example

create BYTETracker instance

byte_tracker = sv.ByteTrack(track_thresh=0.25, track_buffer=30, match_thresh=0.8, frame_rate=30)

create VideoInfo instance

video_info = sv.VideoInfo.from_video_path(SOURCE_VIDEO_PATH)

create frame generator

generator = sv.get_video_frames_generator(SOURCE_VIDEO_PATH)

create LineZone instance, it is previously called LineCounter class

line_zone = sv.LineZone(start=LINE_START, end=LINE_END)

create instance of BoxAnnotator

box_annotator = sv.BoxAnnotator(thickness=4, text_thickness=4, text_scale=2)

create instance of TraceAnnotator

trace_annotator = sv.TraceAnnotator(thickness=4, trace_length=50)

create LineZoneAnnotator instance, it is previously called LineCounterAnnotator class

line_zone_annotator = sv.LineZoneAnnotator(thickness=4, text_thickness=4, text_scale=2)

define call back function to be used in video processing

def callback(frame: np.ndarray, index:int) -> np.ndarray:
# model prediction on single frame and conversion to supervision Detections
results = model(frame, verbose=False)[0]
detections = sv.Detections.from_ultralytics(results)
# only consider class id from selected_classes define above
detections = detections[np.isin(detections.class_id, selected_classes)]
# tracking detections
detections = byte_tracker.update_with_detections(detections)
labels = [
f"#{tracker_id} {model.model.names[class_id]} {confidence:0.2f}"
for confidence, class_id, tracker_id
in zip(detections.confidence, detections.class_id, detections.tracker_id)
]
annotated_frame = trace_annotator.annotate(
scene=frame.copy(),
detections=detections
)
annotated_frame=box_annotator.annotate(
scene=annotated_frame,
detections=detections,
labels=labels)

# update line counter
line_zone.trigger(detections)
# return frame with box and line annotated result
return  line_zone_annotator.annotate(annotated_frame, line_counter=line_zone)

process the whole video

sv.process_video(
source_path = SOURCE_VIDEO_PATH,
target_path = TARGET_VIDEO_PATH,
callback=callback
)

Additional

Can I use this to annotate the fish video; Although I'm not able to embed it into this?

Are you willing to submit a PR?

• Yes I'd like to help by submitting a PR!