vidformer 0.9.0__py3-none-any.whl → 0.10.0__py3-none-any.whl

vidformer/supervision/__init__.py

@@ -0,0 +1,529 @@
+"""
+vidformer.supervision is the [supervision](https://supervision.roboflow.com/) frontend for [vidformer](https://github.com/ixlab/vidformer).
+"""
+
+import vidformer.cv2 as vf_cv2
+
+import supervision as _sv
+import numpy as np
+from supervision import Color, ColorPalette, ColorLookup, Detections
+from supervision.annotators.utils import resolve_color, resolve_text_background_xyxy
+from supervision.detection.utils import spread_out_boxes
+from supervision.config import CLASS_NAME_DATA_FIELD
+from math import sqrt
+from supervision.geometry.core import Position
+
+CV2_FONT = vf_cv2.FONT_HERSHEY_SIMPLEX
+
+
+class BoxAnnotator:
+    def __init__(
+        self,
+        color=ColorPalette.DEFAULT,
+        thickness=2,
+        color_lookup=ColorLookup.CLASS,
+    ):
+        self.color = color
+        self.thickness = thickness
+        self.color_lookup = color_lookup
+
+    def annotate(
+        self,
+        scene: vf_cv2.Frame,
+        detections: Detections,
+        custom_color_lookup=None,
+    ):
+        for detection_idx in range(len(detections)):
+            x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
+            color = resolve_color(
+                color=self.color,
+                detections=detections,
+                detection_idx=detection_idx,
+                color_lookup=(
+                    self.color_lookup
+                    if custom_color_lookup is None
+                    else custom_color_lookup
+                ),
+            )
+            vf_cv2.rectangle(
+                img=scene,
+                pt1=(x1, y1),
+                pt2=(x2, y2),
+                color=color.as_rgb(),
+                thickness=self.thickness,
+            )
+        return scene
+
+
+class RoundBoxAnnotator:
+    def __init__(
+        self,
+        color=ColorPalette.DEFAULT,
+        thickness: int = 2,
+        color_lookup: ColorLookup = ColorLookup.CLASS,
+        roundness: float = 0.6,
+    ):
+        self.color = color
+        self.thickness = thickness
+        self.color_lookup = color_lookup
+        if not 0 < roundness <= 1.0:
+            raise ValueError("roundness attribute must be float between (0, 1.0]")
+        self.roundness = roundness
+
+    def annotate(
+        self,
+        scene: vf_cv2.Frame,
+        detections: _sv.Detections,
+        custom_color_lookup=None,
+    ):
+        for detection_idx in range(len(detections)):
+            x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
+            color = resolve_color(
+                color=self.color,
+                detections=detections,
+                detection_idx=detection_idx,
+                color_lookup=(
+                    self.color_lookup
+                    if custom_color_lookup is None
+                    else custom_color_lookup
+                ),
+            )
+            radius = (
+                int((x2 - x1) // 2 * self.roundness)
+                if abs(x1 - x2) < abs(y1 - y2)
+                else int((y2 - y1) // 2 * self.roundness)
+            )
+            circle_coordinates = [
+                ((x1 + radius), (y1 + radius)),
+                ((x2 - radius), (y1 + radius)),
+                ((x2 - radius), (y2 - radius)),
+                ((x1 + radius), (y2 - radius)),
+            ]
+            line_coordinates = [
+                ((x1 + radius, y1), (x2 - radius, y1)),
+                ((x2, y1 + radius), (x2, y2 - radius)),
+                ((x1 + radius, y2), (x2 - radius, y2)),
+                ((x1, y1 + radius), (x1, y2 - radius)),
+            ]
+            start_angles = (180, 270, 0, 90)
+            end_angles = (270, 360, 90, 180)
+            for center_coordinates, line, start_angle, end_angle in zip(
+                circle_coordinates, line_coordinates, start_angles, end_angles
+            ):
+                vf_cv2.ellipse(
+                    img=scene,
+                    center=center_coordinates,
+                    axes=(radius, radius),
+                    angle=0,
+                    startAngle=start_angle,
+                    endAngle=end_angle,
+                    color=color.as_rgb(),
+                    thickness=self.thickness,
+                )
+                vf_cv2.line(
+                    img=scene,
+                    pt1=line[0],
+                    pt2=line[1],
+                    color=color.as_rgb(),
+                    thickness=self.thickness,
+                )
+        return scene
+
+
+class BoxCornerAnnotator:
+    def __init__(
+        self,
+        color=ColorPalette.DEFAULT,
+        thickness=4,
+        corner_length=15,
+        color_lookup=ColorLookup.CLASS,
+    ):
+        self.color = color
+        self.thickness: int = thickness
+        self.corner_length: int = corner_length
+        self.color_lookup: ColorLookup = color_lookup
+
+    def annotate(
+        self,
+        scene: vf_cv2.Frame,
+        detections: Detections,
+        custom_color_lookup=None,
+    ):
+        for detection_idx in range(len(detections)):
+            x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
+            color = resolve_color(
+                color=self.color,
+                detections=detections,
+                detection_idx=detection_idx,
+                color_lookup=(
+                    self.color_lookup
+                    if custom_color_lookup is None
+                    else custom_color_lookup
+                ),
+            )
+            corners = [(x1, y1), (x2, y1), (x1, y2), (x2, y2)]
+            for x, y in corners:
+                x_end = x + self.corner_length if x == x1 else x - self.corner_length
+                vf_cv2.line(
+                    scene, (x, y), (x_end, y), color.as_rgb(), thickness=self.thickness
+                )
+
+                y_end = y + self.corner_length if y == y1 else y - self.corner_length
+                vf_cv2.line(
+                    scene, (x, y), (x, y_end), color.as_rgb(), thickness=self.thickness
+                )
+        return scene
+
+
+class ColorAnnotator:
+    def __init__(
+        self,
+        color=ColorPalette.DEFAULT,
+        opacity: float = 0.5,
+        color_lookup: ColorLookup = ColorLookup.CLASS,
+    ):
+        self.color = color
+        self.color_lookup: ColorLookup = color_lookup
+        self.opacity = opacity
+
+    def annotate(
+        self,
+        scene: vf_cv2.Frame,
+        detections: Detections,
+        custom_color_lookup=None,
+    ):
+        scene_with_boxes = scene.copy()
+        for detection_idx in range(len(detections)):
+            x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
+            color = resolve_color(
+                color=self.color,
+                detections=detections,
+                detection_idx=detection_idx,
+                color_lookup=(
+                    self.color_lookup
+                    if custom_color_lookup is None
+                    else custom_color_lookup
+                ),
+            )
+            vf_cv2.rectangle(
+                img=scene_with_boxes,
+                pt1=(x1, y1),
+                pt2=(x2, y2),
+                color=color.as_rgb(),
+                thickness=-1,
+            )
+
+        vf_cv2.addWeighted(
+            scene_with_boxes, self.opacity, scene, 1 - self.opacity, gamma=0, dst=scene
+        )
+        return scene
+
+
+class CircleAnnotator:
+    def __init__(
+        self,
+        color=ColorPalette.DEFAULT,
+        thickness: int = 2,
+        color_lookup: ColorLookup = ColorLookup.CLASS,
+    ):
+        self.color = color
+        self.thickness: int = thickness
+        self.color_lookup: ColorLookup = color_lookup
+
+    def annotate(
+        self,
+        scene: vf_cv2.Frame,
+        detections: Detections,
+        custom_color_lookup=None,
+    ):
+        for detection_idx in range(len(detections)):
+            x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
+            center = ((x1 + x2) // 2, (y1 + y2) // 2)
+            distance = sqrt((x1 - center[0]) ** 2 + (y1 - center[1]) ** 2)
+            color = resolve_color(
+                color=self.color,
+                detections=detections,
+                detection_idx=detection_idx,
+                color_lookup=(
+                    self.color_lookup
+                    if custom_color_lookup is None
+                    else custom_color_lookup
+                ),
+            )
+            vf_cv2.circle(
+                img=scene,
+                center=center,
+                radius=int(distance),
+                color=color.as_rgb(),
+                thickness=self.thickness,
+            )
+
+        return scene
+
+
+class DotAnnotator:
+    def __init__(
+        self,
+        color=ColorPalette.DEFAULT,
+        radius: int = 4,
+        position: Position = Position.CENTER,
+        color_lookup: ColorLookup = ColorLookup.CLASS,
+        outline_thickness: int = 0,
+        outline_color=Color.BLACK,
+    ):
+
+        self.color = color
+        self.radius: int = radius
+        self.position: Position = position
+        self.color_lookup: ColorLookup = color_lookup
+        self.outline_thickness = outline_thickness
+        self.outline_color = outline_color
+
+    def annotate(
+        self,
+        scene: vf_cv2.Frame,
+        detections: Detections,
+        custom_color_lookup=None,
+    ):
+        xy = detections.get_anchors_coordinates(anchor=self.position)
+        for detection_idx in range(len(detections)):
+            color = resolve_color(
+                color=self.color,
+                detections=detections,
+                detection_idx=detection_idx,
+                color_lookup=(
+                    self.color_lookup
+                    if custom_color_lookup is None
+                    else custom_color_lookup
+                ),
+            )
+            center = (int(xy[detection_idx, 0]), int(xy[detection_idx, 1]))
+
+            vf_cv2.circle(scene, center, self.radius, color.as_rgb(), -1)
+            if self.outline_thickness:
+                outline_color = resolve_color(
+                    color=self.outline_color,
+                    detections=detections,
+                    detection_idx=detection_idx,
+                    color_lookup=(
+                        self.color_lookup
+                        if custom_color_lookup is None
+                        else custom_color_lookup
+                    ),
+                )
+                vf_cv2.circle(
+                    scene,
+                    center,
+                    self.radius,
+                    outline_color.as_rgb(),
+                    self.outline_thickness,
+                )
+        return scene
+
+
+class LabelAnnotator:
+    def __init__(
+        self,
+        color=ColorPalette.DEFAULT,
+        text_color=Color.WHITE,
+        text_scale: float = 0.5,
+        text_thickness: int = 1,
+        text_padding: int = 10,
+        text_position: Position = Position.TOP_LEFT,
+        color_lookup: ColorLookup = ColorLookup.CLASS,
+        border_radius: int = 0,
+        smart_position: bool = False,
+    ):
+        self.border_radius: int = border_radius
+        self.color = color
+        self.text_color = text_color
+        self.text_scale: float = text_scale
+        self.text_thickness: int = text_thickness
+        self.text_padding: int = text_padding
+        self.text_anchor: Position = text_position
+        self.color_lookup: ColorLookup = color_lookup
+        self.smart_position = smart_position
+
+    def annotate(
+        self,
+        scene,
+        detections: Detections,
+        labels,
+        custom_color_lookup=None,
+    ):
+        self._validate_labels(labels, detections)
+
+        labels = self._get_labels_text(detections, labels)
+        label_properties = self._get_label_properties(detections, labels)
+
+        if self.smart_position:
+            xyxy = label_properties[:, :4]
+            xyxy = spread_out_boxes(xyxy)
+            label_properties[:, :4] = xyxy
+
+        self._draw_labels(
+            scene=scene,
+            labels=labels,
+            label_properties=label_properties,
+            detections=detections,
+            custom_color_lookup=custom_color_lookup,
+        )
+
+        return scene
+
+    def _validate_labels(self, labels, detections: Detections):
+        if labels is not None and len(labels) != len(detections):
+            raise ValueError(
+                f"The number of labels ({len(labels)}) does not match the "
+                f"number of detections ({len(detections)}). Each detection "
+                f"should have exactly 1 label."
+            )
+
+    def _get_label_properties(
+        self,
+        detections: Detections,
+        labels,
+    ):
+        label_properties = []
+        anchors_coordinates = detections.get_anchors_coordinates(
+            anchor=self.text_anchor
+        ).astype(int)
+
+        for label, center_coords in zip(labels, anchors_coordinates):
+            (text_w, text_h) = vf_cv2.getTextSize(
+                text=label,
+                fontFace=CV2_FONT,
+                fontScale=self.text_scale,
+                thickness=self.text_thickness,
+            )[0]
+
+            width_padded = text_w + 2 * self.text_padding
+            height_padded = text_h + 2 * self.text_padding
+
+            text_background_xyxy = resolve_text_background_xyxy(
+                center_coordinates=tuple(center_coords),
+                text_wh=(width_padded, height_padded),
+                position=self.text_anchor,
+            )
+
+            label_properties.append(
+                [
+                    *text_background_xyxy,
+                    text_h,
+                ]
+            )
+
+        return np.array(label_properties).reshape(-1, 5)
+
+    @staticmethod
+    def _get_labels_text(detections: Detections, custom_labels):
+        if custom_labels is not None:
+            return custom_labels
+
+        labels = []
+        for idx in range(len(detections)):
+            if CLASS_NAME_DATA_FIELD in detections.data:
+                labels.append(detections.data[CLASS_NAME_DATA_FIELD][idx])
+            elif detections.class_id is not None:
+                labels.append(str(detections.class_id[idx]))
+            else:
+                labels.append(str(idx))
+        return labels
+
+    def _draw_labels(
+        self,
+        scene,
+        labels,
+        label_properties,
+        detections,
+        custom_color_lookup,
+    ) -> None:
+        assert len(labels) == len(label_properties) == len(detections), (
+            f"Number of label properties ({len(label_properties)}), "
+            f"labels ({len(labels)}) and detections ({len(detections)}) "
+            "do not match."
+        )
+
+        color_lookup = (
+            custom_color_lookup
+            if custom_color_lookup is not None
+            else self.color_lookup
+        )
+
+        for idx, label_property in enumerate(label_properties):
+            background_color = resolve_color(
+                color=self.color,
+                detections=detections,
+                detection_idx=idx,
+                color_lookup=color_lookup,
+            )
+            text_color = resolve_color(
+                color=self.text_color,
+                detections=detections,
+                detection_idx=idx,
+                color_lookup=color_lookup,
+            )
+
+            box_xyxy = label_property[:4]
+            text_height_padded = label_property[4]
+            self.draw_rounded_rectangle(
+                scene=scene,
+                xyxy=box_xyxy,
+                color=background_color.as_rgb(),
+                border_radius=self.border_radius,
+            )
+
+            text_x = box_xyxy[0] + self.text_padding
+            text_y = box_xyxy[1] + self.text_padding + text_height_padded
+            vf_cv2.putText(
+                img=scene,
+                text=labels[idx],
+                org=(text_x, text_y),
+                fontFace=CV2_FONT,
+                fontScale=self.text_scale,
+                color=text_color.as_rgb(),
+                thickness=self.text_thickness,
+                lineType=vf_cv2.LINE_AA,
+            )
+
+    @staticmethod
+    def draw_rounded_rectangle(
+        scene: np.ndarray,
+        xyxy,
+        color,
+        border_radius: int,
+    ) -> np.ndarray:
+        x1, y1, x2, y2 = xyxy
+        width = x2 - x1
+        height = y2 - y1
+
+        border_radius = min(border_radius, min(width, height) // 2)
+
+        rectangle_coordinates = [
+            ((x1 + border_radius, y1), (x2 - border_radius, y2)),
+            ((x1, y1 + border_radius), (x2, y2 - border_radius)),
+        ]
+        circle_centers = [
+            (x1 + border_radius, y1 + border_radius),
+            (x2 - border_radius, y1 + border_radius),
+            (x1 + border_radius, y2 - border_radius),
+            (x2 - border_radius, y2 - border_radius),
+        ]
+
+        for coordinates in rectangle_coordinates:
+            vf_cv2.rectangle(
+                img=scene,
+                pt1=coordinates[0],
+                pt2=coordinates[1],
+                color=color,
+                thickness=-1,
+            )
+        for center in circle_centers:
+            vf_cv2.circle(
+                img=scene,
+                center=center,
+                radius=border_radius,
+                color=color,
+                thickness=-1,
+            )
+        return scene

{vidformer-0.9.0.dist-info → vidformer-0.10.0.dist-info}/METADATA

@@ -1,7 +1,7 @@
 Metadata-Version: 2.3
 Name: vidformer
-Version: 0.9.0
-Summary: A Python library for creating and viewing videos with vidformer.
+Version: 0.10.0
+Summary: vidformer-py is a Python 🐍 interface for [vidformer](https://github.com/ixlab/vidformer).
 Author-email: Dominik Winecki <dominikwinecki@gmail.com>
 Requires-Python: >=3.8
 Description-Content-Type: text/markdown
@@ -19,13 +19,15 @@ Project-URL: Issues, https://github.com/ixlab/vidformer/issues
 [![PyPI version](https://img.shields.io/pypi/v/vidformer.svg)](https://pypi.org/project/vidformer/)
 [![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://github.com/ixlab/vidformer/blob/main/LICENSE)
 
-vidformer-py is a Python 🐍 interface for [vidformer](https://github.com/ixlab/vidformer).
+vidformer-py is a Python 🐍 frontend for [vidformer](https://github.com/ixlab/vidformer).
+It has an API compatability layer with OpenCV cv2, as well as some [supervision](https://github.com/roboflow/supervision) annotators.
 Our [getting started guide](https://ixlab.github.io/vidformer/getting-started.html) explains how to use it.
 
 **Quick links:**
 * [📦 PyPI](https://pypi.org/project/vidformer/)
-* [📘 Documentation - vidformer-py](https://ixlab.github.io/vidformer/vidformer-py/)
-* [📘 Documentation - vidformer.cv2](https://ixlab.github.io/vidformer/vidformer-py-cv2/)
+* [📘 Documentation - vidformer-py](https://ixlab.github.io/vidformer/vidformer-py/pdoc/)
+* [📘 Documentation - vidformer.cv2](https://ixlab.github.io/vidformer/vidformer-py/pdoc/vidformer/cv2.html)
+* [📘 Documentation - vidformer.supervision](https://ixlab.github.io/vidformer/vidformer-py/pdoc/vidformer/supervision.html)
 * [🧑‍💻 Source Code](https://github.com/ixlab/vidformer/tree/main/vidformer-py/)
 
 **Publish:**

vidformer-0.10.0.dist-info/RECORD

@@ -0,0 +1,6 @@
+vidformer/__init__.py,sha256=qpWcttHsW1wnaGx5__qJqaT-m5VF7yMiHCxdm10Fjek,44916
+vidformer/cv2/__init__.py,sha256=DGm5NB4FGCHxPVez-yO748DjocKruxn4QBqqThgskWI,25555
+vidformer/supervision/__init__.py,sha256=T2QJ3gKtUSoKOlxAf06TG4fD9IgIDuBpiVOBKVk3qAw,17150
+vidformer-0.10.0.dist-info/WHEEL,sha256=CpUCUxeHQbRN5UGRQHYRJorO5Af-Qy_fHMctcQ8DSGI,82
+vidformer-0.10.0.dist-info/METADATA,sha256=f_aUIFbQUoVJuYTKtqexwyN5uzWRoS2Fvd2dHZ_EGbo,1800
+vidformer-0.10.0.dist-info/RECORD,,