eye-cv 1.0.0__py3-none-any.whl

eye/dataset/formats/yolo.py

@@ -0,0 +1,272 @@
+import os
+from pathlib import Path
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
+
+import cv2
+import numpy as np
+
+from eye.config import ORIENTED_BOX_COORDINATES
+from eye.dataset.utils import approximate_mask_with_polygons
+from eye.detection.core import Detections
+from eye.detection.utils import polygon_to_mask, polygon_to_xyxy
+from eye.utils.file import (
+    list_files_with_extensions,
+    read_txt_file,
+    read_yaml_file,
+    save_text_file,
+    save_yaml_file,
+)
+
+if TYPE_CHECKING:
+    from eye.dataset.core import DetectionDataset
+
+
+def _parse_box(values: List[str]) -> np.ndarray:
+    x_center, y_center, width, height = values
+    return np.array(
+        [
+            float(x_center) - float(width) / 2,
+            float(y_center) - float(height) / 2,
+            float(x_center) + float(width) / 2,
+            float(y_center) + float(height) / 2,
+        ],
+        dtype=np.float32,
+    )
+
+
+def _box_to_polygon(box: np.ndarray) -> np.ndarray:
+    return np.array(
+        [[box[0], box[1]], [box[2], box[1]], [box[2], box[3]], [box[0], box[3]]]
+    )
+
+
+def _parse_polygon(values: List[str]) -> np.ndarray:
+    return np.array(values, dtype=np.float32).reshape(-1, 2)
+
+
+def _polygons_to_masks(
+    polygons: List[np.ndarray], resolution_wh: Tuple[int, int]
+) -> np.ndarray:
+    return np.array(
+        [
+            polygon_to_mask(polygon=polygon, resolution_wh=resolution_wh)
+            for polygon in polygons
+        ],
+        dtype=bool,
+    )
+
+
+def _with_mask(lines: List[str]) -> bool:
+    return any([len(line.split()) > 5 for line in lines])
+
+
+def _extract_class_names(file_path: str) -> List[str]:
+    data = read_yaml_file(file_path=file_path)
+    names = data["names"]
+    if isinstance(names, dict):
+        names = [names[key] for key in sorted(names.keys())]
+    return names
+
+
+def _image_name_to_annotation_name(image_name: str) -> str:
+    base_name, _ = os.path.splitext(image_name)
+    return base_name + ".txt"
+
+
+def yolo_annotations_to_detections(
+    lines: List[str],
+    resolution_wh: Tuple[int, int],
+    with_masks: bool,
+    is_obb: bool = False,
+) -> Detections:
+    if len(lines) == 0:
+        return Detections.empty()
+
+    class_id, relative_xyxy, relative_polygon, relative_xyxyxyxy = [], [], [], []
+    w, h = resolution_wh
+    for line in lines:
+        values = line.split()
+        class_id.append(int(values[0]))
+        if len(values) == 5:
+            box = _parse_box(values=values[1:])
+            relative_xyxy.append(box)
+            if with_masks:
+                relative_polygon.append(_box_to_polygon(box=box))
+        elif len(values) > 5:
+            polygon = _parse_polygon(values=values[1:])
+            relative_xyxy.append(polygon_to_xyxy(polygon=polygon))
+            if is_obb:
+                relative_xyxyxyxy.append(np.array(values[1:]))
+            if with_masks:
+                relative_polygon.append(polygon)
+
+    class_id = np.array(class_id, dtype=int)
+    relative_xyxy = np.array(relative_xyxy, dtype=np.float32)
+    xyxy = relative_xyxy * np.array([w, h, w, h], dtype=np.float32)
+    data = {}
+
+    if is_obb:
+        relative_xyxyxyxy = np.array(relative_xyxyxyxy, dtype=np.float32)
+        xyxyxyxy = relative_xyxyxyxy.reshape(-1, 4, 2)
+        xyxyxyxy *= np.array([w, h], dtype=np.float32)
+        data[ORIENTED_BOX_COORDINATES] = xyxyxyxy
+
+    if not with_masks:
+        return Detections(class_id=class_id, xyxy=xyxy, data=data)
+
+    polygons = [
+        (polygon * np.array(resolution_wh)).astype(int) for polygon in relative_polygon
+    ]
+    mask = _polygons_to_masks(polygons=polygons, resolution_wh=resolution_wh)
+    return Detections(class_id=class_id, xyxy=xyxy, data=data, mask=mask)
+
+
+def load_yolo_annotations(
+    images_directory_path: str,
+    annotations_directory_path: str,
+    data_yaml_path: str,
+    force_masks: bool = False,
+    is_obb: bool = False,
+) -> Tuple[List[str], List[str], Dict[str, Detections]]:
+    """
+    Loads YOLO annotations and returns class names, images,
+        and their corresponding detections.
+
+    Args:
+        images_directory_path (str): The path to the directory containing the images.
+        annotations_directory_path (str): The path to the directory
+            containing the YOLO annotation files.
+        data_yaml_path (str): The path to the data
+            YAML file containing class information.
+        force_masks (bool): If True, forces masks to be loaded
+            for all annotations, regardless of whether they are present.
+        is_obb (bool): If True, loads the annotations in OBB format.
+            OBB annotations are defined as `[class_id, x, y, x, y, x, y, x, y]`,
+            where pairs of [x, y] are box corners.
+
+    Returns:
+        Tuple[List[str], List[str], Dict[str, Detections]]:
+            A tuple containing a list of class names, a dictionary with
+            image names as keys and images as values, and a dictionary
+            with image names as keys and corresponding Detections instances as values.
+    """
+    image_paths = [
+        str(path)
+        for path in list_files_with_extensions(
+            directory=images_directory_path, extensions=["jpg", "jpeg", "png"]
+        )
+    ]
+
+    classes = _extract_class_names(file_path=data_yaml_path)
+    annotations = {}
+
+    for image_path in image_paths:
+        image_stem = Path(image_path).stem
+        annotation_path = os.path.join(annotations_directory_path, f"{image_stem}.txt")
+        if not os.path.exists(annotation_path):
+            annotations[image_path] = Detections.empty()
+            continue
+
+        image = cv2.imread(image_path)
+        lines = read_txt_file(file_path=annotation_path, skip_empty=True)
+        h, w, _ = image.shape
+        resolution_wh = (w, h)
+
+        with_masks = _with_mask(lines=lines)
+        with_masks = force_masks if force_masks else with_masks
+        annotation = yolo_annotations_to_detections(
+            lines=lines,
+            resolution_wh=resolution_wh,
+            with_masks=with_masks,
+            is_obb=is_obb,
+        )
+        annotations[image_path] = annotation
+    return classes, image_paths, annotations
+
+
+def object_to_yolo(
+    xyxy: np.ndarray,
+    class_id: int,
+    image_shape: Tuple[int, int, int],
+    polygon: Optional[np.ndarray] = None,
+) -> str:
+    h, w, _ = image_shape
+    if polygon is None:
+        xyxy_relative = xyxy / np.array([w, h, w, h], dtype=np.float32)
+        x_min, y_min, x_max, y_max = xyxy_relative
+        x_center = (x_min + x_max) / 2
+        y_center = (y_min + y_max) / 2
+        width = x_max - x_min
+        height = y_max - y_min
+        return f"{int(class_id)} {x_center:.5f} {y_center:.5f} {width:.5f} {height:.5f}"
+    else:
+        polygon_relative = polygon / np.array([w, h], dtype=np.float32)
+        polygon_relative = polygon_relative.reshape(-1)
+        polygon_parsed = " ".join([f"{value:.5f}" for value in polygon_relative])
+        return f"{int(class_id)} {polygon_parsed}"
+
+
+def detections_to_yolo_annotations(
+    detections: Detections,
+    image_shape: Tuple[int, int, int],
+    min_image_area_percentage: float = 0.0,
+    max_image_area_percentage: float = 1.0,
+    approximation_percentage: float = 0.75,
+) -> List[str]:
+    annotation = []
+    for xyxy, mask, _, class_id, _, _ in detections:
+        if class_id is None:
+            raise ValueError("Class ID is required for YOLO annotations.")
+
+        if mask is not None:
+            polygons = approximate_mask_with_polygons(
+                mask=mask,
+                min_image_area_percentage=min_image_area_percentage,
+                max_image_area_percentage=max_image_area_percentage,
+                approximation_percentage=approximation_percentage,
+            )
+            for polygon in polygons:
+                xyxy = polygon_to_xyxy(polygon=polygon)
+                next_object = object_to_yolo(
+                    xyxy=xyxy,
+                    class_id=class_id,
+                    image_shape=image_shape,
+                    polygon=polygon,
+                )
+                annotation.append(next_object)
+        else:
+            next_object = object_to_yolo(
+                xyxy=xyxy, class_id=class_id, image_shape=image_shape
+            )
+            annotation.append(next_object)
+    return annotation
+
+
+def save_yolo_annotations(
+    dataset: "DetectionDataset",
+    annotations_directory_path: str,
+    min_image_area_percentage: float = 0.0,
+    max_image_area_percentage: float = 1.0,
+    approximation_percentage: float = 0.75,
+) -> None:
+    Path(annotations_directory_path).mkdir(parents=True, exist_ok=True)
+    for image_path, image, annotation in dataset:
+        image_name = Path(image_path).name
+        yolo_annotations_name = _image_name_to_annotation_name(image_name=image_name)
+        yolo_annotations_path = os.path.join(
+            annotations_directory_path, yolo_annotations_name
+        )
+        lines = detections_to_yolo_annotations(
+            detections=annotation,
+            image_shape=image.shape,  # type: ignore
+            min_image_area_percentage=min_image_area_percentage,
+            max_image_area_percentage=max_image_area_percentage,
+            approximation_percentage=approximation_percentage,
+        )
+        save_text_file(lines=lines, file_path=yolo_annotations_path)
+
+
+def save_data_yaml(data_yaml_path: str, classes: List[str]) -> None:
+    data = {"nc": len(classes), "names": classes}
+    Path(data_yaml_path).parent.mkdir(parents=True, exist_ok=True)
+    save_yaml_file(data=data, file_path=data_yaml_path)

eye/dataset/utils.py

@@ -0,0 +1,259 @@
+import copy
+import os
+import random
+import shutil
+from pathlib import Path
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, TypeVar, Union
+
+import cv2
+import numpy as np
+import numpy.typing as npt
+
+from eye.detection.core import Detections
+from eye.detection.utils import (
+    approximate_polygon,
+    filter_polygons_by_area,
+    mask_to_polygons,
+)
+
+if TYPE_CHECKING:
+    from eye.dataset.core import DetectionDataset
+
+T = TypeVar("T")
+
+
+def approximate_mask_with_polygons(
+    mask: np.ndarray,
+    min_image_area_percentage: float = 0.0,
+    max_image_area_percentage: float = 1.0,
+    approximation_percentage: float = 0.75,
+) -> List[np.ndarray]:
+    height, width = mask.shape
+    image_area = height * width
+    minimum_detection_area = min_image_area_percentage * image_area
+    maximum_detection_area = max_image_area_percentage * image_area
+
+    polygons = mask_to_polygons(mask=mask)
+    if len(polygons) == 1:
+        polygons = filter_polygons_by_area(
+            polygons=polygons, min_area=None, max_area=maximum_detection_area
+        )
+    else:
+        polygons = filter_polygons_by_area(
+            polygons=polygons,
+            min_area=minimum_detection_area,
+            max_area=maximum_detection_area,
+        )
+    return [
+        approximate_polygon(polygon=polygon, percentage=approximation_percentage)
+        for polygon in polygons
+    ]
+
+
+def merge_class_lists(class_lists: List[List[str]]) -> List[str]:
+    unique_classes = set()
+
+    for class_list in class_lists:
+        for class_name in class_list:
+            unique_classes.add(class_name)
+
+    return sorted(list(unique_classes))
+
+
+def build_class_index_mapping(
+    source_classes: List[str], target_classes: List[str]
+) -> Dict[int, int]:
+    """Returns the index map of source classes -> target classes."""
+    index_mapping = {}
+
+    for i, class_name in enumerate(source_classes):
+        if class_name not in target_classes:
+            raise ValueError(
+                f"Class {class_name} not found in target classes. "
+                "source_classes must be a subset of target_classes."
+            )
+        corresponding_index = target_classes.index(class_name)
+        index_mapping[i] = corresponding_index
+
+    return index_mapping
+
+
+def map_detections_class_id(
+    source_to_target_mapping: Dict[int, int], detections: Detections
+) -> Detections:
+    if detections.class_id is None:
+        raise ValueError("Detections must have class_id attribute.")
+    if set(np.unique(detections.class_id)) - set(source_to_target_mapping.keys()):
+        raise ValueError(
+            "Detections class_id must be a subset of source_to_target_mapping keys."
+        )
+
+    detections_copy = copy.deepcopy(detections)
+
+    if len(detections) > 0:
+        detections_copy.class_id = np.vectorize(source_to_target_mapping.get)(
+            detections_copy.class_id
+        )
+
+    return detections_copy
+
+
+def save_dataset_images(
+    dataset: "DetectionDataset", images_directory_path: str
+) -> None:
+    Path(images_directory_path).mkdir(parents=True, exist_ok=True)
+    for image_path in dataset.image_paths:
+        final_path = os.path.join(images_directory_path, Path(image_path).name)
+        if image_path in dataset._images_in_memory:
+            image = dataset._images_in_memory[image_path]
+            cv2.imwrite(final_path, image)
+        else:
+            shutil.copyfile(image_path, final_path)
+
+
+def train_test_split(
+    data: List[T],
+    train_ratio: float = 0.8,
+    random_state: Optional[int] = None,
+    shuffle: bool = True,
+) -> Tuple[List[T], List[T]]:
+    """
+    Splits the data into two parts using the provided train_ratio.
+
+    Args:
+        data (List[T]): The data to split.
+        train_ratio (float): The ratio of the training set to the entire dataset.
+        random_state (Optional[int]): The seed for the random number generator.
+        shuffle (bool): Whether to shuffle the data before splitting.
+
+    Returns:
+        Tuple[List[T], List[T]]: The split data.
+    """
+    if random_state is not None:
+        random.seed(random_state)
+
+    if shuffle:
+        random.shuffle(data)
+
+    split_index = int(len(data) * train_ratio)
+    return data[:split_index], data[split_index:]
+
+
+def rle_to_mask(
+    rle: Union[npt.NDArray[np.int_], List[int]], resolution_wh: Tuple[int, int]
+) -> npt.NDArray[np.bool_]:
+    """
+    Converts run-length encoding (RLE) to a binary mask.
+
+    Args:
+        rle (Union[npt.NDArray[np.int_], List[int]]): The 1D RLE array, the format
+            used in the COCO dataset (column-wise encoding, values of an array with
+            even indices represent the number of pixels assigned as background,
+            values of an array with odd indices represent the number of pixels
+            assigned as foreground object).
+        resolution_wh (Tuple[int, int]): The width (w) and height (h)
+            of the desired binary mask.
+
+    Returns:
+        The generated 2D Boolean mask of shape `(h, w)`, where the foreground object is
+            marked with `True`'s and the rest is filled with `False`'s.
+
+    Raises:
+        AssertionError: If the sum of pixels encoded in RLE differs from the
+            number of pixels in the expected mask (computed based on resolution_wh).
+
+    Examples:
+        ```python
+        import eye as sv
+
+        sv.rle_to_mask([5, 2, 2, 2, 5], (4, 4))
+        # array([
+        #     [False, False, False, False],
+        #     [False, True,  True,  False],
+        #     [False, True,  True,  False],
+        #     [False, False, False, False],
+        # ])
+        ```
+    """
+    if isinstance(rle, list):
+        rle = np.array(rle, dtype=int)
+
+    width, height = resolution_wh
+
+    assert width * height == np.sum(rle), (
+        "the sum of the number of pixels in the RLE must be the same "
+        "as the number of pixels in the expected mask"
+    )
+
+    zero_one_values = np.zeros(shape=(rle.size, 1), dtype=np.uint8)
+    zero_one_values[1::2] = 1
+
+    decoded_rle = np.repeat(zero_one_values, rle, axis=0)
+    decoded_rle = np.append(
+        decoded_rle, np.zeros(width * height - len(decoded_rle), dtype=np.uint8)
+    )
+    return decoded_rle.reshape((height, width), order="F")
+
+
+def mask_to_rle(mask: npt.NDArray[np.bool_]) -> List[int]:
+    """
+    Converts a binary mask into a run-length encoding (RLE).
+
+    Args:
+        mask (npt.NDArray[np.bool_]): 2D binary mask where `True` indicates foreground
+            object and `False` indicates background.
+
+    Returns:
+        The run-length encoded mask. Values of a list with even indices
+            represent the number of pixels assigned as background (`False`), values
+            of a list with odd indices represent the number of pixels assigned
+            as foreground object (`True`).
+
+    Raises:
+        AssertionError: If input mask is not 2D or is empty.
+
+    Examples:
+        ```python
+        import numpy as np
+        import eye as sv
+
+        mask = np.array([
+            [True, True, True, True],
+            [True, True, True, True],
+            [True, True, True, True],
+            [True, True, True, True],
+        ])
+        sv.mask_to_rle(mask)
+        # [0, 16]
+
+        mask = np.array([
+            [False, False, False, False],
+            [False, True,  True,  False],
+            [False, True,  True,  False],
+            [False, False, False, False],
+        ])
+        sv.mask_to_rle(mask)
+        # [5, 2, 2, 2, 5]
+        ```
+
+    ![mask_to_rle](https://media.roboflow.com/eye-docs/mask-to-rle.png){ align=center width="800" }
+    """  # noqa E501 // docs
+    assert mask.ndim == 2, "Input mask must be 2D"
+    assert mask.size != 0, "Input mask cannot be empty"
+
+    on_value_change_indices = np.where(
+        mask.ravel(order="F") != np.roll(mask.ravel(order="F"), 1)
+    )[0]
+
+    on_value_change_indices = np.append(on_value_change_indices, mask.size)
+    # need to add 0 at the beginning when the same value is in the first and
+    # last element of the flattened mask
+    if on_value_change_indices[0] != 0:
+        on_value_change_indices = np.insert(on_value_change_indices, 0, 0)
+
+    rle = np.diff(on_value_change_indices)
+
+    if mask[0][0] == 1:
+        rle = np.insert(rle, 0, 0)
+
+    return list(rle)

eye/detection/auto_convert.py

@@ -0,0 +1,155 @@
+"""Universal auto-conversion for ALL model formats - Works with supervision!"""
+
+import numpy as np
+from typing import Any, Optional, Union
+from eye.detection.core import Detections
+
+
+def auto_convert(results: Any) -> Detections:
+    """Automatically convert ANY model output to eye.Detections.
+    
+    Supports: YOLO (Ultralytics), PyTorch, TensorFlow, OpenCV, ONNX,
+             TensorRT, MMDetection, Detectron2, PaddlePaddle, and more!
+    
+    Args:
+        results: Model output in any supported format
+        
+    Returns:
+        eye.Detections object ready to use
+        
+    Example:
+        >>> results = model(image)
+        >>> detections = eye.auto_convert(results)
+        >>> annotated = annotator.annotate(image, detections)
+    """
+    # YOLO (Ultralytics) - Most common in this project
+    if hasattr(results, 'boxes'):
+        boxes = results.boxes
+        if hasattr(boxes, 'xyxy'):
+            xyxy = boxes.xyxy.cpu().numpy() if hasattr(boxes.xyxy, 'cpu') else boxes.xyxy
+            confidence = boxes.conf.cpu().numpy() if hasattr(boxes.conf, 'cpu') else boxes.conf
+            class_id = boxes.cls.cpu().numpy().astype(int) if hasattr(boxes.cls, 'cpu') else boxes.cls.astype(int)
+            
+            # Handle optional tracking IDs
+            tracker_id = None
+            if hasattr(boxes, 'id') and boxes.id is not None:
+                tracker_id = boxes.id.cpu().numpy().astype(int) if hasattr(boxes.id, 'cpu') else boxes.id.astype(int)
+            
+            # Handle segmentation masks
+            mask = None
+            if hasattr(results, 'masks') and results.masks is not None:
+                try:
+                    mask = results.masks.data.cpu().numpy() if hasattr(results.masks.data, 'cpu') else results.masks.data
+                except Exception:
+                    pass
+            
+            return Detections(
+                xyxy=xyxy,
+                confidence=confidence,
+                class_id=class_id,
+                tracker_id=tracker_id,
+                mask=mask
+            )
+    
+    # YOLO list format (batch)
+    if isinstance(results, list) and len(results) > 0:
+        if hasattr(results[0], 'boxes'):
+            return auto_convert(results[0])  # Take first result
+    
+    # PyTorch (torchvision) - Dict with boxes/scores/labels
+    if isinstance(results, dict):
+        if all(k in results for k in ['boxes', 'scores', 'labels']):
+            boxes = results['boxes']
+            xyxy = boxes.cpu().numpy() if hasattr(boxes, 'cpu') else boxes
+            confidence = results['scores'].cpu().numpy() if hasattr(results['scores'], 'cpu') else results['scores']
+            class_id = results['labels'].cpu().numpy().astype(int) if hasattr(results['labels'], 'cpu') else results['labels'].astype(int)
+            
+            return Detections(
+                xyxy=xyxy,
+                confidence=confidence,
+                class_id=class_id
+            )
+    
+    # TensorFlow Object Detection API
+    if isinstance(results, (tuple, list)) and len(results) >= 3:
+        boxes, scores, classes = results[:3]
+        # TF format is [ymin, xmin, ymax, xmax] normalized - convert to xyxy
+        if hasattr(boxes, 'numpy'):
+            boxes = boxes.numpy()
+        boxes = boxes.squeeze()
+        
+        # Convert to absolute coordinates (assuming image size is available)
+        # For now, keep as is and let user handle scaling
+        xyxy = boxes[:, [1, 0, 3, 2]]  # Reorder to [xmin, ymin, xmax, ymax]
+        
+        confidence = scores.numpy() if hasattr(scores, 'numpy') else scores
+        class_id = classes.numpy().astype(int) if hasattr(classes, 'numpy') else classes.astype(int)
+        
+        return Detections(
+            xyxy=xyxy.squeeze(),
+            confidence=confidence.squeeze(),
+            class_id=class_id.squeeze()
+        )
+    
+    # OpenCV DNN format
+    if isinstance(results, np.ndarray):
+        if results.ndim == 4 and results.shape[0] == 1:
+            # OpenCV format: [1, 1, N, 7] where 7 = [image_id, class_id, conf, x1, y1, x2, y2]
+            detections_2d = results.reshape(-1, results.shape[-1])
+            
+            confidence_mask = detections_2d[:, 2] > 0  # Filter valid detections
+            valid = detections_2d[confidence_mask]
+            
+            return Detections(
+                xyxy=valid[:, 3:7],  # [x1, y1, x2, y2]
+                confidence=valid[:, 2],
+                class_id=valid[:, 1].astype(int)
+            )
+        
+        # Raw array format (ONNX/TensorRT): [N, 6+] where columns are [x, y, w, h, conf, class, ...]
+        if results.ndim == 2 and results.shape[1] >= 6:
+            # Convert xywh to xyxy
+            x = results[:, 0]
+            y = results[:, 1]
+            w = results[:, 2]
+            h = results[:, 3]
+            
+            xyxy = np.stack([
+                x - w/2,  # x1
+                y - h/2,  # y1
+                x + w/2,  # x2
+                y + h/2   # y2
+            ], axis=1)
+            
+            return Detections(
+                xyxy=xyxy,
+                confidence=results[:, 4],
+                class_id=results[:, 5].astype(int)
+            )
+    
+    # If nothing matched, raise helpful error
+    raise TypeError(
+        f"Unsupported model output format: {type(results)}. "
+        f"Supported: YOLO, PyTorch, TensorFlow, OpenCV, ONNX, TensorRT. "
+        f"Got: {results.__class__.__name__}"
+    )
+
+
+def from_yolo(results) -> Detections:
+    """Convert YOLO (Ultralytics) results to Detections."""
+    return auto_convert(results)
+
+
+def from_pytorch(results: dict) -> Detections:
+    """Convert PyTorch/torchvision results to Detections."""
+    return auto_convert(results)
+
+
+def from_tensorflow(results: tuple) -> Detections:
+    """Convert TensorFlow Object Detection API results to Detections."""
+    return auto_convert(results)
+
+
+def from_opencv(results: np.ndarray) -> Detections:
+    """Convert OpenCV DNN results to Detections."""
+    return auto_convert(results)