eye-cv 1.0.0__py3-none-any.whl

eye/dataset/formats/coco.py

@@ -0,0 +1,258 @@
+import os
+from datetime import datetime
+from pathlib import Path
+from typing import TYPE_CHECKING, Dict, List, Tuple
+
+import numpy as np
+import numpy.typing as npt
+
+from eye.dataset.utils import (
+    approximate_mask_with_polygons,
+    map_detections_class_id,
+    mask_to_rle,
+    rle_to_mask,
+)
+from eye.detection.core import Detections
+from eye.detection.utils import (
+    contains_holes,
+    contains_multiple_segments,
+    polygon_to_mask,
+)
+from eye.utils.file import read_json_file, save_json_file
+
+if TYPE_CHECKING:
+    from eye.dataset.core import DetectionDataset
+
+
+def coco_categories_to_classes(coco_categories: List[dict]) -> List[str]:
+    return [
+        category["name"]
+        for category in sorted(coco_categories, key=lambda category: category["id"])
+    ]
+
+
+def build_coco_class_index_mapping(
+    coco_categories: List[dict], target_classes: List[str]
+) -> Dict[int, int]:
+    source_class_to_index = {
+        category["name"]: category["id"] for category in coco_categories
+    }
+    return {
+        source_class_to_index[target_class_name]: target_class_index
+        for target_class_index, target_class_name in enumerate(target_classes)
+    }
+
+
+def classes_to_coco_categories(classes: List[str]) -> List[dict]:
+    return [
+        {
+            "id": class_id,
+            "name": class_name,
+            "supercategory": "common-objects",
+        }
+        for class_id, class_name in enumerate(classes)
+    ]
+
+
+def group_coco_annotations_by_image_id(
+    coco_annotations: List[dict],
+) -> Dict[int, List[dict]]:
+    annotations = {}
+    for annotation in coco_annotations:
+        image_id = annotation["image_id"]
+        if image_id not in annotations:
+            annotations[image_id] = []
+        annotations[image_id].append(annotation)
+    return annotations
+
+
+def coco_annotations_to_masks(
+    image_annotations: List[dict], resolution_wh: Tuple[int, int]
+) -> npt.NDArray[np.bool_]:
+    return np.array(
+        [
+            rle_to_mask(
+                rle=np.array(image_annotation["segmentation"]["counts"]),
+                resolution_wh=resolution_wh,
+            )
+            if image_annotation["iscrowd"]
+            else polygon_to_mask(
+                polygon=np.reshape(
+                    np.asarray(image_annotation["segmentation"], dtype=np.int32),
+                    (-1, 2),
+                ),
+                resolution_wh=resolution_wh,
+            )
+            for image_annotation in image_annotations
+        ],
+        dtype=bool,
+    )
+
+
+def coco_annotations_to_detections(
+    image_annotations: List[dict], resolution_wh: Tuple[int, int], with_masks: bool
+) -> Detections:
+    if not image_annotations:
+        return Detections.empty()
+
+    class_ids = [
+        image_annotation["category_id"] for image_annotation in image_annotations
+    ]
+    xyxy = [image_annotation["bbox"] for image_annotation in image_annotations]
+    xyxy = np.asarray(xyxy)
+    xyxy[:, 2:4] += xyxy[:, 0:2]
+
+    if with_masks:
+        mask = coco_annotations_to_masks(
+            image_annotations=image_annotations, resolution_wh=resolution_wh
+        )
+        return Detections(
+            class_id=np.asarray(class_ids, dtype=int), xyxy=xyxy, mask=mask
+        )
+
+    return Detections(xyxy=xyxy, class_id=np.asarray(class_ids, dtype=int))
+
+
+def detections_to_coco_annotations(
+    detections: Detections,
+    image_id: int,
+    annotation_id: int,
+    min_image_area_percentage: float = 0.0,
+    max_image_area_percentage: float = 1.0,
+    approximation_percentage: float = 0.75,
+) -> Tuple[List[Dict], int]:
+    coco_annotations = []
+    for xyxy, mask, _, class_id, _, _ in detections:
+        box_width, box_height = xyxy[2] - xyxy[0], xyxy[3] - xyxy[1]
+        segmentation = []
+        iscrowd = 0
+        if mask is not None:
+            iscrowd = contains_holes(mask=mask) or contains_multiple_segments(mask=mask)
+
+            if iscrowd:
+                segmentation = {
+                    "counts": mask_to_rle(mask=mask),
+                    "size": list(mask.shape[:2]),
+                }
+            else:
+                segmentation = [
+                    list(
+                        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,
+                        )[0].flatten()
+                    )
+                ]
+        coco_annotation = {
+            "id": annotation_id,
+            "image_id": image_id,
+            "category_id": int(class_id),
+            "bbox": [xyxy[0], xyxy[1], box_width, box_height],
+            "area": box_width * box_height,
+            "segmentation": segmentation,
+            "iscrowd": iscrowd,
+        }
+        coco_annotations.append(coco_annotation)
+        annotation_id += 1
+    return coco_annotations, annotation_id
+
+
+def load_coco_annotations(
+    images_directory_path: str,
+    annotations_path: str,
+    force_masks: bool = False,
+) -> Tuple[List[str], List[str], Dict[str, Detections]]:
+    coco_data = read_json_file(file_path=annotations_path)
+    classes = coco_categories_to_classes(coco_categories=coco_data["categories"])
+    class_index_mapping = build_coco_class_index_mapping(
+        coco_categories=coco_data["categories"], target_classes=classes
+    )
+    coco_images = coco_data["images"]
+    coco_annotations_groups = group_coco_annotations_by_image_id(
+        coco_annotations=coco_data["annotations"]
+    )
+
+    images = []
+    annotations = {}
+
+    for coco_image in coco_images:
+        image_name, image_width, image_height = (
+            coco_image["file_name"],
+            coco_image["width"],
+            coco_image["height"],
+        )
+        image_annotations = coco_annotations_groups.get(coco_image["id"], [])
+        image_path = os.path.join(images_directory_path, image_name)
+
+        annotation = coco_annotations_to_detections(
+            image_annotations=image_annotations,
+            resolution_wh=(image_width, image_height),
+            with_masks=force_masks,
+        )
+        annotation = map_detections_class_id(
+            source_to_target_mapping=class_index_mapping,
+            detections=annotation,
+        )
+
+        images.append(image_path)
+        annotations[image_path] = annotation
+
+    return classes, images, annotations
+
+
+def save_coco_annotations(
+    dataset: "DetectionDataset",
+    annotation_path: str,
+    min_image_area_percentage: float = 0.0,
+    max_image_area_percentage: float = 1.0,
+    approximation_percentage: float = 0.75,
+) -> None:
+    Path(annotation_path).parent.mkdir(parents=True, exist_ok=True)
+    licenses = [
+        {
+            "id": 1,
+            "url": "https://creativecommons.org/licenses/by/4.0/",
+            "name": "CC BY 4.0",
+        }
+    ]
+
+    coco_annotations = []
+    coco_images = []
+    coco_categories = classes_to_coco_categories(classes=dataset.classes)
+
+    image_id, annotation_id = 1, 1
+    for image_path, image, annotation in dataset:
+        image_height, image_width, _ = image.shape
+        image_name = f"{Path(image_path).stem}{Path(image_path).suffix}"
+        coco_image = {
+            "id": image_id,
+            "license": 1,
+            "file_name": image_name,
+            "height": image_height,
+            "width": image_width,
+            "date_captured": datetime.now().strftime("%m/%d/%Y,%H:%M:%S"),
+        }
+
+        coco_images.append(coco_image)
+        coco_annotation, annotation_id = detections_to_coco_annotations(
+            detections=annotation,
+            image_id=image_id,
+            annotation_id=annotation_id,
+            min_image_area_percentage=min_image_area_percentage,
+            max_image_area_percentage=max_image_area_percentage,
+            approximation_percentage=approximation_percentage,
+        )
+
+        coco_annotations.extend(coco_annotation)
+        image_id += 1
+
+    annotation_dict = {
+        "info": {},
+        "licenses": licenses,
+        "categories": coco_categories,
+        "images": coco_images,
+        "annotations": coco_annotations,
+    }
+    save_json_file(annotation_dict, file_path=annotation_path)

eye/dataset/formats/pascal_voc.py

@@ -0,0 +1,279 @@
+import os
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+from xml.etree.ElementTree import Element, SubElement
+
+import cv2
+import numpy as np
+from defusedxml.ElementTree import parse, tostring
+from defusedxml.minidom import parseString
+
+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
+
+
+def object_to_pascal_voc(
+    xyxy: np.ndarray, name: str, polygon: Optional[np.ndarray] = None
+) -> Element:
+    root = Element("object")
+
+    object_name = SubElement(root, "name")
+    object_name.text = name
+
+    # https://github.com/roboflow/eye/issues/144
+    xyxy += 1
+
+    bndbox = SubElement(root, "bndbox")
+    xmin = SubElement(bndbox, "xmin")
+    xmin.text = str(int(xyxy[0]))
+    ymin = SubElement(bndbox, "ymin")
+    ymin.text = str(int(xyxy[1]))
+    xmax = SubElement(bndbox, "xmax")
+    xmax.text = str(int(xyxy[2]))
+    ymax = SubElement(bndbox, "ymax")
+    ymax.text = str(int(xyxy[3]))
+
+    if polygon is not None:
+        # https://github.com/roboflow/eye/issues/144
+        polygon += 1
+        object_polygon = SubElement(root, "polygon")
+        for index, point in enumerate(polygon, start=1):
+            x_coordinate, y_coordinate = point
+            x = SubElement(object_polygon, f"x{index}")
+            x.text = str(x_coordinate)
+            y = SubElement(object_polygon, f"y{index}")
+            y.text = str(y_coordinate)
+
+    return root
+
+
+def detections_to_pascal_voc(
+    detections: Detections,
+    classes: List[str],
+    filename: str,
+    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,
+) -> str:
+    """
+    Converts Detections object to Pascal VOC XML format.
+
+    Args:
+        detections (Detections): A Detections object containing bounding boxes,
+            class ids, and other relevant information.
+        classes (List[str]): A list of class names corresponding to the
+            class ids in the Detections object.
+        filename (str): The name of the image file associated with the detections.
+        image_shape (Tuple[int, int, int]): The shape of the image
+            file associated with the detections.
+        min_image_area_percentage (float): Minimum detection area
+            relative to area of image associated with it.
+        max_image_area_percentage (float): Maximum detection area
+            relative to area of image associated with it.
+        approximation_percentage (float): The percentage of
+            polygon points to be removed from the input polygon, in the range [0, 1).
+    Returns:
+        str: An XML string in Pascal VOC format representing the detections.
+    """
+    height, width, depth = image_shape
+
+    # Create root element
+    annotation = Element("annotation")
+
+    # Add folder element
+    folder = SubElement(annotation, "folder")
+    folder.text = "VOC"
+
+    # Add filename element
+    file_name = SubElement(annotation, "filename")
+    file_name.text = filename
+
+    # Add source element
+    source = SubElement(annotation, "source")
+    database = SubElement(source, "database")
+    database.text = "roboflow.ai"
+
+    # Add size element
+    size = SubElement(annotation, "size")
+    w = SubElement(size, "width")
+    w.text = str(width)
+    h = SubElement(size, "height")
+    h.text = str(height)
+    d = SubElement(size, "depth")
+    d.text = str(depth)
+
+    # Add segmented element
+    segmented = SubElement(annotation, "segmented")
+    segmented.text = "0"
+
+    # Add object elements
+    for xyxy, mask, _, class_id, _, _ in detections:
+        name = classes[class_id]
+        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_pascal_voc(
+                    xyxy=xyxy, name=name, polygon=polygon
+                )
+                annotation.append(next_object)
+        else:
+            next_object = object_to_pascal_voc(xyxy=xyxy, name=name)
+            annotation.append(next_object)
+
+    # Generate XML string
+    xml_string = parseString(tostring(annotation)).toprettyxml(indent="  ")
+    return xml_string
+
+
+def load_pascal_voc_annotations(
+    images_directory_path: str,
+    annotations_directory_path: str,
+    force_masks: bool = False,
+) -> Tuple[List[str], List[str], Dict[str, Detections]]:
+    """
+    Loads PASCAL VOC XML annotations and returns the image name,
+        a Detections instance, and a list of class names.
+
+    Args:
+        images_directory_path (str): The path to the directory containing the images.
+        annotations_directory_path (str): The path to the directory containing the
+            PASCAL VOC annotation files.
+        force_masks (bool): If True, forces masks to be loaded for all
+            annotations, regardless of whether they are present.
+
+    Returns:
+        Tuple[List[str], List[str], Dict[str, Detections]]: A tuple with a list
+            of class names, a list of paths to images, and a dictionary with image
+            paths 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: List[str] = []
+    annotations = {}
+
+    for image_path in image_paths:
+        image_stem = Path(image_path).stem
+        annotation_path = os.path.join(annotations_directory_path, f"{image_stem}.xml")
+        if not os.path.exists(annotation_path):
+            annotations[image_path] = Detections.empty()
+            continue
+
+        tree = parse(annotation_path)
+        root = tree.getroot()
+
+        image = cv2.imread(image_path)
+        resolution_wh = (image.shape[1], image.shape[0])
+        annotation, classes = detections_from_xml_obj(
+            root, classes, resolution_wh, force_masks
+        )
+        annotations[image_path] = annotation
+
+    return classes, image_paths, annotations
+
+
+def detections_from_xml_obj(
+    root: Element, classes: List[str], resolution_wh, force_masks: bool = False
+) -> Tuple[Detections, List[str]]:
+    """
+    Converts an XML object in Pascal VOC format to a Detections object.
+    Expected XML format:
+    <annotation>
+        ...
+        <object>
+            <name>dog</name>
+            <bndbox>
+                <xmin>48</xmin>
+                <ymin>240</ymin>
+                <xmax>195</xmax>
+                <ymax>371</ymax>
+            </bndbox>
+            <polygon>
+                <x1>48</x1>
+                <y1>240</y1>
+                <x2>195</x2>
+                <y2>240</y2>
+                <x3>195</x3>
+                <y3>371</y3>
+                <x4>48</x4>
+                <y4>371</y4>
+            </polygon>
+        </object>
+    </annotation>
+
+    Returns:
+        Tuple[Detections, List[str]]: A tuple containing a Detections object and an
+            updated list of class names, extended with the class names
+            from the XML object.
+    """
+    xyxy = []
+    class_names = []
+    masks = []
+    with_masks = False
+    extended_classes = classes[:]
+    for obj in root.findall("object"):
+        class_name = obj.find("name").text
+        class_names.append(class_name)
+
+        bbox = obj.find("bndbox")
+        x1 = int(bbox.find("xmin").text)
+        y1 = int(bbox.find("ymin").text)
+        x2 = int(bbox.find("xmax").text)
+        y2 = int(bbox.find("ymax").text)
+
+        xyxy.append([x1, y1, x2, y2])
+
+        with_masks = obj.find("polygon") is not None
+        with_masks = force_masks if force_masks else with_masks
+
+        for polygon in obj.findall("polygon"):
+            polygon = parse_polygon_points(polygon)
+            # https://github.com/roboflow/eye/issues/144
+            polygon -= 1
+
+            mask_from_polygon = polygon_to_mask(
+                polygon=polygon,
+                resolution_wh=resolution_wh,
+            )
+            masks.append(mask_from_polygon)
+
+    xyxy = np.array(xyxy) if len(xyxy) > 0 else np.empty((0, 4))
+
+    # https://github.com/roboflow/eye/issues/144
+    xyxy -= 1
+
+    for k in set(class_names):
+        if k not in extended_classes:
+            extended_classes.append(k)
+    class_id = np.array(
+        [extended_classes.index(class_name) for class_name in class_names]
+    )
+
+    annotation = Detections(
+        xyxy=xyxy.astype(np.float32),
+        mask=np.array(masks).astype(bool) if with_masks else None,
+        class_id=class_id,
+    )
+
+    return annotation, extended_classes
+
+
+def parse_polygon_points(polygon: Element) -> np.ndarray:
+    coordinates = [int(coord.text) for coord in polygon.findall(".//*")]
+    return np.array(
+        [(coordinates[i], coordinates[i + 1]) for i in range(0, len(coordinates), 2)]
+    )