supervisely 6.73.276__py3-none-any.whl → 6.73.278__py3-none-any.whl

supervisely/convert/image/pascal_voc/pascal_voc_helper.py

@@ -1,23 +1,42 @@
 import os
-from typing import List, Tuple
+import shutil
+from pathlib import Path
+from typing import Callable, Dict, List, Optional, OrderedDict, Tuple, Union
 
 import numpy as np
-from supervisely.convert.image.image_helper import validate_image_bounds
+from PIL import Image
+from tqdm import tqdm
+
 from supervisely import (
     Annotation,
+    Dataset,
     Label,
     ObjClass,
     ObjClassCollection,
+    Project,
     ProjectMeta,
     generate_free_name,
     logger,
 )
+from supervisely.convert.image.image_helper import validate_image_bounds
 from supervisely.geometry.bitmap import Bitmap
+from supervisely.geometry.polygon import Polygon
 from supervisely.geometry.rectangle import Rectangle
+from supervisely.imaging.color import generate_rgb
 from supervisely.imaging.image import read
+from supervisely.io.fs import file_exists, get_file_ext, get_file_name
+from supervisely.io.json import load_json_file
+from supervisely.task.progress import tqdm_sly
 
 MASKS_EXTENSION = ".png"
 
+# Export
+SUPPORTED_GEOMETRY_TYPES = {Bitmap, Polygon, Rectangle}
+VALID_IMG_EXT = {".jpe", ".jpeg", ".jpg"}
+TRAIN_TAG_NAME = "train"
+VAL_TAG_NAME = "val"
+TRAINVAL_TAG_NAME = "trainval"
+
 default_classes_colors = {
     "neutral": (224, 224, 192),
     "aeroplane": (128, 0, 0),
@@ -60,15 +79,23 @@ def get_col2coord(img: np.ndarray) -> dict:
 
 def read_colors(colors_file: str) -> Tuple[ObjClassCollection, dict]:
     if os.path.isfile(colors_file):
-        logger.info("Will try to read segmentation colors from provided file.")
-        in_lines = filter(None, map(str.strip, open(colors_file, "r").readlines()))
-        in_splitted = (x.split() for x in in_lines)
-        # Format: {name: (R, G, B)}, values [0; 255]
-        cls2col = {}
-        for x in in_splitted:
-            if len(x) != 4:
-                raise ValueError("Invalid format of colors file.")
-            cls2col[x[0]] = (int(x[1]), int(x[2]), int(x[3]))
+        try:
+            logger.info("Will try to read segmentation colors from provided file.")
+            with open(colors_file, "r") as file:
+                cls2col = {}
+                for line in file:
+                    parts = line.strip().split()
+                    if len(parts) < 4:
+                        raise ValueError("Invalid format of colors file.")
+                    class_name = " ".join(parts[:-3])
+                    colors = tuple(map(int, parts[-3:]))
+                    cls2col[class_name] = colors
+        except Exception as e:
+            logger.warning(
+                "Failed to read segmentation colors from provided file. "
+                "Will use default PascalVOC color mapping."
+            )
+            cls2col = default_classes_colors
     else:
         logger.info("Will use default PascalVOC color mapping.")
         cls2col = default_classes_colors
@@ -87,6 +114,7 @@ def read_colors(colors_file: str) -> Tuple[ObjClassCollection, dict]:
     color2class_name = {v: k for k, v in cls2col.items()}
     return obj_classes, color2class_name
 
+
 def get_ann(
     item,
     color2class_name: dict,
@@ -230,3 +258,381 @@ def update_meta_from_xml(
             bbox_classes_map[original_class_name] = class_name
 
     return meta
+
+
+def sly_ann_to_pascal_voc(ann: Annotation, image_name: str) -> Tuple[dict]:
+    """
+    Convert Supervisely annotation to Pascal VOC format annotation.
+
+    :param ann: Supervisely annotation.
+    :type ann: :class:`Annotation<supervisely.annotation.annotation.Annotation>`
+    :param image_name: Image name.
+    :type image_name: :class:`str`
+    :return: Tuple with xml tree and instance and class masks in PIL.Image format.
+    :rtype: :class:`Tuple`
+
+    :Usage example:
+
+     .. code-block:: python
+
+        import supervisely as sly
+        from supervisely.convert.image.pascal_voc.pascal_voc_helper import sly_ann_to_pascal_voc
+
+        ann = sly.Annotation.from_json(ann_json, meta)
+        xml_tree, instance_mask, class_mask = sly_ann_to_pascal_voc(ann, image_name)
+    """
+
+    def from_ann_to_instance_mask(ann: Annotation, contour_thickness: int = 3):
+        mask = np.zeros((ann.img_size[0], ann.img_size[1], 3), dtype=np.uint8)
+        for label in ann.labels:
+            if label.obj_class.geometry_type == Rectangle:
+                continue
+
+            if label.obj_class.name == "neutral":
+                label.geometry.draw(mask, default_classes_colors["neutral"])
+                continue
+
+            label.geometry.draw_contour(mask, default_classes_colors["neutral"], contour_thickness)
+            label.geometry.draw(mask, label.obj_class.color)
+
+        res_mask = Image.fromarray(mask)
+        res_mask = res_mask.convert("P", palette=Image.ADAPTIVE)  # pylint: disable=no-member
+        return res_mask
+
+    def from_ann_to_class_mask(ann: Annotation, contour_thickness: int = 3):
+        exist_colors = [[0, 0, 0], default_classes_colors["neutral"]]
+        mask = np.zeros((ann.img_size[0], ann.img_size[1], 3), dtype=np.uint8)
+        for label in ann.labels:
+            if label.obj_class.geometry_type == Rectangle:
+                continue
+
+            if label.obj_class.name == "neutral":
+                label.geometry.draw(mask, default_classes_colors["neutral"])
+                continue
+
+            new_color = generate_rgb(exist_colors)
+            exist_colors.append(new_color)
+            label.geometry.draw_contour(mask, default_classes_colors["neutral"], contour_thickness)
+            label.geometry.draw(mask, new_color)
+
+        res_mask = Image.fromarray(mask)
+        res_mask = res_mask.convert("P", palette=Image.ADAPTIVE)  # pylint: disable=no-member
+        return res_mask
+
+    def from_ann_to_xml(ann: Annotation, image_name: str):
+        import lxml.etree as ET  # pylint: disable=import-error
+
+        xml_root = ET.Element("annotation")
+
+        ET.SubElement(xml_root, "folder").text = f"VOC"
+        ET.SubElement(xml_root, "filename").text = image_name
+
+        xml_root_source = ET.SubElement(xml_root, "source")
+        ET.SubElement(xml_root_source, "database").text = ""
+
+        ET.SubElement(xml_root_source, "annotation").text = "PASCAL VOC"
+        ET.SubElement(xml_root_source, "image").text = ""
+
+        xml_root_size = ET.SubElement(xml_root, "size")
+        ET.SubElement(xml_root_size, "width").text = str(ann.img_size[1])
+        ET.SubElement(xml_root_size, "height").text = str(ann.img_size[0])
+        ET.SubElement(xml_root_size, "depth").text = "3"
+
+        ET.SubElement(xml_root, "segmented").text = "1" if len(ann.labels) > 0 else "0"
+
+        for label in ann.labels:
+            if label.obj_class.name == "neutral":
+                continue
+
+            bitmap_to_bbox = label.geometry.to_bbox()
+
+            xml_ann_obj = ET.SubElement(xml_root, "object")
+            ET.SubElement(xml_ann_obj, "name").text = label.obj_class.name
+            ET.SubElement(xml_ann_obj, "pose").text = "Unspecified"
+            ET.SubElement(xml_ann_obj, "truncated").text = "0"
+            ET.SubElement(xml_ann_obj, "difficult").text = "0"
+
+            xml_ann_obj_bndbox = ET.SubElement(xml_ann_obj, "bndbox")
+            ET.SubElement(xml_ann_obj_bndbox, "xmin").text = str(bitmap_to_bbox.left)
+            ET.SubElement(xml_ann_obj_bndbox, "ymin").text = str(bitmap_to_bbox.top)
+            ET.SubElement(xml_ann_obj_bndbox, "xmax").text = str(bitmap_to_bbox.right)
+            ET.SubElement(xml_ann_obj_bndbox, "ymax").text = str(bitmap_to_bbox.bottom)
+
+        tree = ET.ElementTree(xml_root)
+        return tree
+
+    pascal_ann = from_ann_to_xml(ann, image_name)
+    instance_mask = from_ann_to_instance_mask(ann)
+    class_mask = from_ann_to_class_mask(ann)
+    return pascal_ann, instance_mask, class_mask
+
+
+def sly_ds_to_pascal_voc(
+    dataset: Dataset,
+    meta: ProjectMeta,
+    dest_dir: Optional[str] = None,
+    train_val_split_coef: float = 0.8,
+    log_progress: bool = False,
+    progress_cb: Optional[Union[tqdm, Callable]] = None,
+) -> Tuple[Dict, Optional[Dict]]:
+    """
+    Convert Supervisely dataset to Pascal VOC format.
+
+    :param meta: Project meta information.
+    :type meta: :class:`ProjectMeta<supervisely.project.project_meta.ProjectMeta>`
+    :param dest_dir: Destination directory.
+    :type dest_dir: :class:`str`, optional
+    :param train_val_split_coef: Coefficient for splitting images into train and validation sets.
+    :type train_val_split_coef: :class:`float`, optional
+    :param log_progress: If True, log progress.
+    :type log_progress: :class:`str`, optional
+    :param progress_cb: Progress callback.
+    :type progress_cb: :class:`Callable`, optional
+    :return: None
+    :rtype: NoneType
+
+    :Usage example:
+
+     .. code-block:: python
+
+        import supervisely as sly
+        from supervisely.convert.image.pascal_voc.pascal_voc_helper import sly_ds_to_pascal_voc
+
+        project_path = "/home/admin/work/supervisely/projects/lemons_annotated"
+        project = sly.Project(project_path, sly.OpenMode.READ)
+
+        for ds in project.datasets:
+            dest_dir = "/home/admin/work/supervisely/projects/lemons_annotated_pascal_voc"
+            sly_ds_to_pascal_voc(ds, project.meta, dest_dir=dest_dir)
+    """
+    import lxml.etree as ET  # pylint: disable=import-error
+
+    def write_main_set(
+        is_trainval: int,
+        images_stats: dict,
+        meta: ProjectMeta,
+        result_main_sets_dir: str,
+        result_segmentation_sets_dir: str,
+    ):
+        res_files = ["trainval.txt", "train.txt", "val.txt"]
+        for file in os.listdir(result_segmentation_sets_dir):
+            if file in res_files:
+                shutil.copyfile(
+                    os.path.join(result_segmentation_sets_dir, file),
+                    os.path.join(result_main_sets_dir, file),
+                )
+
+        train_imgs = [i for i in images_stats if i["dataset"] == TRAIN_TAG_NAME]
+        val_imgs = [i for i in images_stats if i["dataset"] == VAL_TAG_NAME]
+
+        write_objs = [
+            {"suffix": "trainval", "imgs": images_stats},
+            {"suffix": "train", "imgs": train_imgs},
+            {"suffix": "val", "imgs": val_imgs},
+        ]
+
+        if is_trainval == 1:
+            trainval_imgs = [
+                i for i in images_stats if i["dataset"] == TRAIN_TAG_NAME + VAL_TAG_NAME
+            ]
+            write_objs[0] = {"suffix": "trainval", "imgs": trainval_imgs}
+
+        for obj_cls in meta.obj_classes:
+            if obj_cls.geometry_type not in SUPPORTED_GEOMETRY_TYPES:
+                continue
+            if obj_cls.name == "neutral":
+                continue
+            for o in write_objs:
+                with open(
+                    os.path.join(result_main_sets_dir, f'{obj_cls.name}_{o["suffix"]}.txt'), "a"
+                ) as f:
+                    for img_stats in o["imgs"]:
+                        v = "1" if obj_cls.name in img_stats["classes"] else "-1"
+                        f.write(f'{img_stats["name"]} {v}\n')
+
+    def write_segm_set(is_trainval: int, images_stats: dict, result_imgsets_dir: str):
+        with open(os.path.join(result_imgsets_dir, "trainval.txt"), "a") as f:
+            if is_trainval == 1:
+                f.writelines(
+                    i["name"] + "\n"
+                    for i in images_stats
+                    if i["dataset"] == TRAIN_TAG_NAME + VAL_TAG_NAME
+                )
+            else:
+                f.writelines(i["name"] + "\n" for i in images_stats)
+        with open(os.path.join(result_imgsets_dir, "train.txt"), "a") as f:
+            f.writelines(i["name"] + "\n" for i in images_stats if i["dataset"] == TRAIN_TAG_NAME)
+        with open(os.path.join(result_imgsets_dir, "val.txt"), "a") as f:
+            f.writelines(i["name"] + "\n" for i in images_stats if i["dataset"] == VAL_TAG_NAME)
+
+    if progress_cb is not None:
+        log_progress = False
+
+    if log_progress:
+        progress_cb = tqdm_sly(
+            desc=f"Converting dataset '{dataset.short_name}' to Pascal VOC format",
+            total=len(dataset),
+        )
+
+    logger.info(f"Processing dataset: '{dataset.name}'")
+
+    # Prepare Pascal VOC root directory
+    if dest_dir is None:
+        dest_dir = str(Path(dataset.path).parent / "pascal_voc")
+
+    pascal_root_path = os.path.join(dest_dir, "VOCdevkit", "VOC")
+    result_images_dir = os.path.join(pascal_root_path, "JPEGImages")
+    result_ann_dir = os.path.join(pascal_root_path, "Annotations")
+    result_obj_dir = os.path.join(pascal_root_path, "SegmentationObject")
+    result_class_dir = os.path.join(pascal_root_path, "SegmentationClass")
+    result_image_sets_dir = os.path.join(pascal_root_path, "ImageSets")
+    result_segmentation_sets_dir = os.path.join(result_image_sets_dir, "Segmentation")
+    result_main_sets_dir = os.path.join(result_image_sets_dir, "Main")
+    result_colors_file_path = os.path.join(pascal_root_path, "colors.txt")
+
+    # Create directories if not exist
+    os.makedirs(result_images_dir, exist_ok=True)
+    os.makedirs(result_ann_dir, exist_ok=True)
+    os.makedirs(result_obj_dir, exist_ok=True)
+    os.makedirs(result_class_dir, exist_ok=True)
+    os.makedirs(result_image_sets_dir, exist_ok=True)
+    os.makedirs(result_segmentation_sets_dir, exist_ok=True)
+    os.makedirs(result_main_sets_dir, exist_ok=True)
+
+    # Create colors.txt file
+    if not file_exists(result_colors_file_path):
+        with open(result_colors_file_path, "w") as f:
+            f.write(
+                f"neutral {default_classes_colors['neutral'][0]} {default_classes_colors['neutral'][1]} {default_classes_colors['neutral'][2]}\n"
+            )
+
+    image_stats = []
+    classes_colors = {}
+    for item_name, img_path, ann_path in dataset.items():
+        # Assign unique name to avoid conflicts
+        unique_name = f"{dataset.name}_{get_file_name(item_name)}"
+
+        # Load annotation
+        ann = Annotation.from_json(load_json_file(ann_path), meta)
+        pascal_ann, instance_mask, class_mask = sly_ann_to_pascal_voc(ann, unique_name)
+
+        # Write ann
+        ann_path = os.path.join(result_ann_dir, f"{unique_name}.xml")
+        ET.indent(pascal_ann, space="    ")
+        pascal_ann.write(ann_path, pretty_print=True)
+
+        # Save instance mask
+        instance_mask_path = os.path.join(
+            result_obj_dir, f"{unique_name}_instance{MASKS_EXTENSION}"
+        )
+        instance_mask.save(instance_mask_path)
+
+        # Save class mask
+        class_mask_path = os.path.join(result_class_dir, f"{unique_name}_class{MASKS_EXTENSION}")
+        class_mask.save(class_mask_path)
+
+        # Save original image
+        img_ext = get_file_ext(img_path)
+        if img_ext not in VALID_IMG_EXT:
+            jpg_name = f"{unique_name}.jpg"
+            jpg_image_path = os.path.join(result_images_dir, jpg_name)
+            img = Image.open(img_path)
+            img.save(jpg_image_path, "JPEG")
+        else:
+            jpg_name = f"{unique_name}{img_ext}"
+            jpg_image_path = os.path.join(result_images_dir, jpg_name)
+            shutil.copyfile(img_path, jpg_image_path)
+
+        # Update stats
+        cur_img_stats = {"classes": set(), "dataset": None, "name": jpg_name}
+        image_stats.append(cur_img_stats)
+
+        # Get classes colors
+        for label in ann.labels:
+            cur_img_stats["classes"].add(label.obj_class.name)
+            classes_colors[label.obj_class.name] = tuple(label.obj_class.color)
+
+        if log_progress:
+            progress_cb.update(1)
+
+    # Update colors.txt file
+    classes_colors = OrderedDict((sorted(classes_colors.items(), key=lambda t: t[0])))
+    with open(result_colors_file_path, "a") as cc:
+        for k in classes_colors.keys():
+            if k == "neutral":
+                continue
+            cc.write(f"{k} {classes_colors[k][0]} {classes_colors[k][1]} {classes_colors[k][2]}\n")
+
+    # Create splits
+    imgs_to_split = [i for i in image_stats if i["dataset"] is None]
+    train_len = int(len(imgs_to_split) * train_val_split_coef)
+
+    for img_stat in imgs_to_split[:train_len]:
+        img_stat["dataset"] = TRAIN_TAG_NAME
+    for img_stat in imgs_to_split[train_len:]:
+        img_stat["dataset"] = VAL_TAG_NAME
+
+    is_trainval = 0
+    write_segm_set(is_trainval, image_stats, result_segmentation_sets_dir)
+    write_main_set(
+        is_trainval, image_stats, meta, result_main_sets_dir, result_segmentation_sets_dir
+    )
+
+
+def sly_project_to_pascal_voc(
+    project: Union[Project, str],
+    dest_dir: Optional[str] = None,
+    train_val_split_coef: float = 0.8,
+    log_progress: bool = True,
+    progress_cb: Optional[Union[tqdm, Callable]] = None,
+) -> None:
+    """
+    Convert Supervisely project to Pascal VOC format.
+
+    :param dest_dir: Destination directory.
+    :type dest_dir: :class:`str`, optional
+    :param train_val_split_coef: Coefficient for splitting images into train and validation sets.
+    :type train_val_split_coef: :class:`float`, optional
+    :param log_progress: Show uploading progress bar.
+    :type log_progress: :class:`bool`
+    :param progress_cb: Function for tracking conversion progress (for all items in the project).
+    :type progress_cb: callable, optional
+    :return: None
+    :rtype: NoneType
+
+    :Usage example:
+
+    .. code-block:: python
+
+        import supervisely as sly
+
+        # Local folder with Project
+        project_directory = "/home/admin/work/supervisely/source/project"
+
+        # Convert Project to Pascal VOC format
+        sly.Project(project_directory).to_pascal_voc(log_progress=True)
+    """
+    if dest_dir is None:
+        dest_dir = project.directory
+
+    Path(dest_dir).mkdir(parents=True, exist_ok=True)
+
+    if progress_cb is not None:
+        log_progress = False
+
+    if log_progress:
+        progress_cb = tqdm_sly(
+            desc="Converting Supervisely project to Pascal VOC format", total=project.total_items
+        )
+
+    for dataset in project.datasets:
+        dataset: Dataset
+        dataset.to_pascal_voc(
+            meta=project.meta,
+            dest_dir=dest_dir,
+            train_val_split_coef=train_val_split_coef,
+            log_progress=log_progress,
+            progress_cb=progress_cb,
+        )
+        logger.info(f"Dataset '{dataset.short_name}' has been converted to Pascal VOC format.")
+    logger.info(f"Project '{project.name}' has been converted to Pascal VOC format.")