supervisely 6.73.277__py3-none-any.whl → 6.73.279__py3-none-any.whl

supervisely/convert/image/coco/coco_helper.py

@@ -1,12 +1,27 @@
 import os
+import shutil
 import sys
 import uuid
 from copy import deepcopy
-from typing import List
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union
 
 import cv2
 import numpy as np
 
+from supervisely._utils import generate_free_name
+from supervisely.api.image_api import ImageApi
+from supervisely.imaging.image import read as sly_image
+from supervisely.io.fs import get_file_name_with_ext
+from supervisely.io.json import dump_json_file, load_json_file
+from supervisely.project.project import Dataset, OpenMode, Project
+from supervisely.project.project_meta import ProjectMeta
+from supervisely.task.progress import tqdm_sly
+
+COCO_INSTANCES_FILE = "coco_instances.json"
+COCO_CAPTIONS_FILE = "coco_captions.json"
+
 
 class HiddenCocoPrints:
     def __enter__(self):
@@ -321,5 +336,482 @@ def create_custom_geometry_config(num_keypoints=None, cat_labels=None, cat_edges
         for edge in cat_edges:
             template.add_edge(src=cat_labels[edge[0] - 1], dst=cat_labels[edge[1] - 1])
     else:
-        logger.warn("Edges can not be mapped without skeleton, please check your annotation")
+        logger.warning("Edges can not be mapped without skeleton, please check your annotation")
     return template
+
+
+def _get_graph_info(idx, obj_class):
+    data = {"supercategory": obj_class.name, "id": idx, "name": obj_class.name}
+    kp = {i: n["label"] for i, n in obj_class.geometry_config["nodes"].items()}
+    keys = {k: j for j, k in enumerate(list(kp.keys()), 1)}
+    edges = obj_class.geometry_config["edges"]
+    sk = [[keys[e["src"]], keys[e["dst"]]] for e in edges]
+    data["keypoints"] = list(kp.values())
+    data["skeleton"] = sk
+    return data
+
+
+def get_categories_from_meta(meta: ProjectMeta):
+    cat = lambda idx, c: {"supercategory": c.name, "id": idx, "name": c.name}
+    return [
+        cat(idx, c) if c.geometry_type != GraphNodes else _get_graph_info(idx, c)
+        for idx, c in enumerate(meta.obj_classes)
+    ]
+
+
+def sly_ann_to_coco(
+    ann: Annotation,
+    coco_image_id: int,
+    class_mapping: Dict[str, int],
+    coco_ann: Optional[Union[Dict, List]] = None,
+    last_label_id: Optional[int] = None,
+    coco_captions: Optional[Union[Dict, List]] = None,
+    last_caption_id: Optional[int] = None,
+) -> Tuple[List, List]:
+    """
+    Convert Supervisely annotation to COCO format annotation ("annotations" field).
+
+    :param coco_image_id: Image id in COCO format.
+    :type coco_image_id: int
+    :param class_mapping: Dictionary that maps class names to class ids.
+    :type class_mapping: Dict[str, int]
+    :param coco_ann: COCO annotation in dictionary or list format to append new annotations.
+    :type coco_ann: Union[Dict, List], optional
+    :param last_label_id: Last label id in COCO format to continue counting.
+    :type last_label_id: int, optional
+    :param coco_captions: COCO captions in dictionary or list format to append new captions.
+    :type coco_captions: Union[Dict, List], optional
+    :return: Tuple with list of COCO objects and list of COCO captions.
+    :rtype: :class:`tuple`
+
+
+    :Usage example:
+
+     .. code-block:: python
+
+        import supervisely as sly
+        from supervisely.convert.image.coco.coco_helper import sly_ann_to_coco
+
+
+        coco_instances = dict(
+            info=dict(
+                description="COCO dataset converted from Supervisely",
+                url="None",
+                version=str(1.0),
+                year=2025,
+                contributor="Supervisely",
+                date_created="2025-01-01 00:00:00",
+            ),
+            licenses=[dict(url="None", id=0, name="None")],
+            images=[],
+            annotations=[],
+            categories=get_categories_from_meta(meta),  # [{"supercategory": "lemon", "id": 0, "name": "lemon"}, ...]
+        )
+
+        ann = sly.Annotation.from_json(ann_json, meta)
+        image_id = 11
+        label_id = 222
+        class_mapping = {obj_cls.name: idx for idx, obj_cls in enumerate(meta.obj_classes)}
+
+        curr_coco_ann, _ = sly_ann_to_coco(ann, image_id, class_mapping, coco_instances, label_id)
+        # or
+        # curr_coco_ann, _ = sly_ann_to_coco(ann, image_id, class_mapping, label_id=label_id)
+        # coco_instances["annotations"].extend(curr_coco_ann)
+
+        label_id += len(curr_coco_ann)
+        image_id += 1
+    """
+
+    coco_obj_template = lambda x, y, z: {
+        "id": x,
+        "segmentation": [],
+        "area": 0,
+        "iscrowd": 0,
+        "image_id": y,
+        "bbox": [],
+        "category_id": z,
+    }
+    label_id = last_label_id + 1 if last_label_id is not None else 1
+    if isinstance(coco_ann, dict):
+        coco_ann = coco_ann["annotations"]
+    if isinstance(coco_ann, list):
+        label_id = len(coco_ann) + 1
+
+    last_caption_id = last_caption_id + 1 if last_caption_id is not None else 1
+    if isinstance(coco_captions, dict):
+        coco_captions = coco_captions["annotations"]
+    if isinstance(coco_captions, list):
+        last_caption_id = len(coco_captions) + 1
+
+    def _update_inst_results(label_id, coco_ann, coco_obj, res):
+        label_id += 1
+        if isinstance(coco_ann, list):
+            coco_ann.append(coco_obj)
+        res.append(coco_obj)
+        return label_id
+
+    def _get_common_bbox(labels, sly_bbox=False, approx=False):
+        bboxes = [l.geometry.to_bbox() for l in labels]
+        x = min([bbox.left for bbox in bboxes])
+        y = min([bbox.top for bbox in bboxes])
+        max_x = max([bbox.right for bbox in bboxes])
+        max_y = max([bbox.bottom for bbox in bboxes])
+        if approx:
+            x, y, max_x, max_y = x - 10, y - 10, max_x + 10, max_y + 10
+        if sly_bbox:
+            return Rectangle(top=y, left=x, bottom=max_y, right=max_x)
+        return [x, y, max_x - x, max_y - y]
+
+    def _create_keypoints_obj(label, cat_id, label_id, coco_image_id):
+        nodes_dict = label.obj_class.geometry_config["nodes"]
+        keypoint_uuid_labels = {i: d["label"] for i, d in nodes_dict.items()}
+        keypoints = []
+        for key in keypoint_uuid_labels.keys():
+            if key not in label.geometry.nodes:
+                keypoints.extend([0, 0, 0])
+            else:
+                loc = label.geometry.nodes[key].location
+                keypoints.extend([loc.col, loc.row, 2])
+        coco_obj = coco_obj_template(label_id, coco_image_id, cat_id)
+        coco_obj["keypoints"] = keypoints
+        coco_obj["num_keypoints"] = len(keypoint_uuid_labels)
+        x, y = keypoints[0::3], keypoints[1::3]
+        x0, x1, y0, y1 = (np.min(x), np.max(x), np.min(y), np.max(y))
+        x0, x1, y0, y1 = int(x0), int(x1), int(y0), int(y1)
+        coco_obj["area"] = int((x1 - x0) * (y1 - y0))
+        coco_obj["bbox"] = [x0, y0, x1 - x0, y1 - y0]
+        return coco_obj
+
+    def _update_caption_results(caption_id, coco_captions, caption, res):
+        caption_id += 1
+        if isinstance(coco_captions, list):
+            coco_captions.append(caption)
+        res.append(caption)
+        return caption_id
+
+    res_inst = []  # result list of COCO objects
+
+    for binding_key, labels in ann.get_bindings().items():
+        if binding_key is None:
+            polygons = [l for l in labels if l.obj_class.geometry_type == Polygon]
+            masks = [l for l in labels if l.obj_class.geometry_type == Bitmap]
+            bboxes = [l for l in labels if l.obj_class.geometry_type == Rectangle]
+            graphs = [l for l in labels if l.obj_class.geometry_type == GraphNodes]
+            if len(masks) > 0:
+                for l in masks:
+                    polygon_cls = l.obj_class.clone(geometry_type=Polygon)
+                    polygons.extend(l.convert(polygon_cls))
+            for label in polygons + bboxes:
+                cat_id = class_mapping[label.obj_class.name]
+                coco_obj = coco_obj_template(label_id, coco_image_id, cat_id)
+                coco_obj["bbox"] = _get_common_bbox([label])
+                coco_obj["area"] = label.geometry.area
+                if label.obj_class.geometry_type == Polygon:
+                    poly = label.geometry.to_json()["points"]["exterior"]
+                    poly = np.array(poly).flatten().astype(float).tolist()
+                    coco_obj["segmentation"] = [poly]
+
+                label_id = _update_inst_results(label_id, coco_ann, coco_obj, res_inst)
+
+            for label in graphs:
+                cat_id = class_mapping[label.obj_class.name]
+                new_obj = _create_keypoints_obj(label, cat_id, label_id, coco_image_id)
+                label_id = _update_inst_results(label_id, coco_ann, new_obj, res_inst)
+
+            continue
+
+        bboxes = [l for l in labels if l.obj_class.geometry_type == Rectangle]
+        polygons = [l for l in labels if l.obj_class.geometry_type == Polygon]
+        masks = [l for l in labels if l.obj_class.geometry_type == Bitmap]
+        graphs = [l for l in labels if l.obj_class.geometry_type == GraphNodes]
+
+        if len(masks) > 0:  # convert Bitmap to Polygon
+            for l in masks:
+                polygon_cls = l.obj_class.clone(geometry_type=Polygon)
+                polygons.extend(l.convert(polygon_cls))
+
+        matched_bbox = False
+        if len(polygons) > 0:  # process polygons
+            cat_id = class_mapping[polygons[0].obj_class.name]
+            coco_obj = coco_obj_template(label_id, coco_image_id, cat_id)
+            if len(bboxes) > 0:
+                found = _get_common_bbox(bboxes, sly_bbox=True, approx=True)
+                new = _get_common_bbox(polygons, sly_bbox=True)
+                matched_bbox = found.contains(new)
+
+            polys = [l.geometry.to_json()["points"]["exterior"] for l in polygons]
+            polys = [np.array(p).flatten().astype(float).tolist() for p in polys]
+            coco_obj["segmentation"] = polys
+            coco_obj["area"] = sum([l.geometry.area for l in polygons])
+            coco_obj["bbox"] = _get_common_bbox(bboxes if matched_bbox else polygons)
+            label_id = _update_inst_results(label_id, coco_ann, coco_obj, res_inst)
+
+        if len(graphs) > 0:
+            if len(graphs) > 1:
+                logger.warning(
+                    "Multiple Keypoints in one binding key are not supported. "
+                    "Only the first graph will be converted."
+                )
+            cat_id = class_mapping[graphs[0].obj_class.name]
+            coco_obj = _create_keypoints_obj(graphs[0], cat_id, label_id, coco_image_id)
+            label_id = _update_inst_results(label_id, coco_ann, coco_obj, res_inst)
+
+        if len(bboxes) > 0 and not matched_bbox:  # process bboxes separately
+            for label in bboxes:
+                cat_id = class_mapping[label.obj_class.name]
+                coco_obj = coco_obj_template(label_id, coco_image_id, cat_id)
+                coco_obj["bbox"] = _get_common_bbox([label])
+                coco_obj["area"] = label.geometry.area
+
+                label_id = _update_inst_results(label_id, coco_ann, coco_obj, res_inst)
+
+    is_caption = lambda t: t.meta.name == "caption" and t.meta.value_type == TagValueType.ANY_STRING
+    caption_tags = [tag for tag in ann.img_tags if is_caption(tag)]
+
+    res_captions = []  # result list of COCO captions
+    for tag in caption_tags:
+        caption = {"image_id": coco_image_id, "id": caption_id, "caption": tag.value}
+        caption_id = _update_caption_results(caption_id, coco_captions, caption, res_captions)
+
+    return res_inst, res_captions
+
+
+def has_caption_tag(meta: ProjectMeta) -> bool:
+    tag = meta.get_tag_meta("caption")
+    return tag is not None and tag.value_type == TagValueType.ANY_STRING
+
+
+def create_coco_ann_template(meta: ProjectMeta) -> Dict[str, Any]:
+    now = datetime.now()
+    coco_ann = dict(
+        info=dict(
+            description="COCO dataset converted from Supervisely",
+            url="None",
+            version=str(1.0),
+            year=now.year,
+            contributor="Supervisely",
+            date_created=now.strftime("%Y-%m-%d %H:%M:%S"),
+        ),
+        licenses=[dict(url="None", id=0, name="None")],
+        images=[],
+        annotations=[],
+        categories=[],
+    )
+    coco_ann["categories"] = get_categories_from_meta(meta)
+    return coco_ann
+
+
+def sly_ds_to_coco(
+    dataset: Dataset,
+    meta: ProjectMeta,
+    return_type: Literal["path", "dict"] = "path",
+    dest_dir: Optional[str] = None,
+    copy_images: bool = False,
+    with_captions: bool = False,
+    log_progress: bool = False,
+    progress_cb: Optional[Callable] = None,
+) -> Union[str, Tuple[str, str], Dict, Tuple[Dict, Dict]]:
+    """
+    Convert Supervisely dataset to COCO format.
+
+    Note: Depending on the `return_type` and `with_captions` parameters, the function returns different values.
+
+    :param dataset: Supervisely dataset.
+    :type dataset: :class:`Dataset<supervisely.project.dataset.Dataset>`
+    :param meta: Project meta information.
+    :type meta: :class:`ProjectMeta<supervisely.project.project_meta.ProjectMeta>`
+    :param return_type: Type of return value.
+                        If 'path', returns paths to COCO dataset files.
+                        If 'dict', returns COCO dataset dictionaries.
+    :param dest_dir: Destination path to save COCO dataset.
+    :type dest_dir: :class:`str`, optional
+    :param copy_images: If True, copies images to the destination directory.
+    :type copy_images: :class:`bool`, optional
+    :param with_captions: If True, returns COCO captions.
+    :type with_captions: :class:`bool`, optional
+    :param log_progress: If True, logs the progress of the conversion.
+    :type log_progress: :class:`bool`, optional
+    :param progress_cb: Callback function to track the progress of the conversion.
+    :type progress_cb: :class:`Callable`, optional
+    :return:
+            If return_type is 'path', returns paths to COCO dataset file or file (instances or instances and captions).
+            If return_type is 'dict', returns COCO dataset dictionary or dictionaries (instances or instances and captions).
+    :rtype: :class:`tuple`
+
+    :Usage example:
+
+     .. code-block:: python
+
+        import supervisely as sly
+        from supervisely.convert.image.coco.coco_helper import sly_ds_to_coco
+
+        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/ds1"
+            coco_json, coco_captions, coco_json_path, coco_captions_path = sly_ds_to_coco(ds, project.meta, save=True, dest_dir=dest_dir)
+    """
+    dest_dir = Path(dataset.path).parent / "coco" if dest_dir is None else Path(dest_dir)
+    save_json = return_type == "path"
+    if save_json is True:
+        annotations_dir = dest_dir / "annotations"
+        annotations_dir.mkdir(parents=True, exist_ok=True)
+    if copy_images is True:
+        images_dir = dest_dir / "images"
+        images_dir.mkdir(parents=True, exist_ok=True)
+
+    if progress_cb is not None:
+        log_progress = False
+
+    if log_progress:
+        progress_cb = tqdm_sly(
+            desc=f"Converting dataset '{dataset.short_name}' to COCO format",
+            total=len(dataset),
+        ).update
+
+    coco_ann = create_coco_ann_template(meta)
+    coco_captions = create_coco_ann_template(meta) if with_captions else None
+
+    image_coco = lambda info, idx: dict(
+        license="None",
+        file_name=info.name,
+        url="None",
+        height=info.height,
+        width=info.width,
+        date_captured=info.created_at,
+        id=idx,
+        sly_id=info.id,
+    )
+
+    class_mapping = {cls.name: idx for idx, cls in enumerate(meta.obj_classes)}
+    label_id = 0
+    caption_id = 0
+    for image_idx, name in enumerate(dataset.get_items_names(), 1):
+        img_path = dataset.get_img_path(name)
+        img_name = get_file_name_with_ext(img_path)
+        ann_path = dataset.get_ann_path(name)
+        img_info_path = dataset.get_img_info_path(name)
+
+        if copy_images:
+            dst_img_path = images_dir / img_name
+            shutil.copy(img_path, dst_img_path)
+
+        if os.path.exists(img_info_path):
+            image_info_json = load_json_file(img_info_path)
+        else:
+            img = sly_image(img_path, remove_alpha_channel=False)
+            now = datetime.now()
+            image_info_json = {
+                "id": None,
+                "name": img_name,
+                "height": img.shape[0],
+                "width": img.shape[1],
+                "created_at": now.strftime("%Y-%m-%d %H:%M:%S"),
+            }
+        image_info = ImageApi._convert_json_info(ImageApi(None), image_info_json)
+
+        coco_ann["images"].append(image_coco(image_info, image_idx))
+        if with_captions is True:
+            coco_captions["images"].append(image_coco(image_info, image_idx))  # pylint: disable=unsubscriptable-object
+
+        ann = Annotation.load_json_file(ann_path, meta)
+        if ann.img_size is None or ann.img_size == (0, 0) or ann.img_size == (None, None):
+            img = sly_image(img_path)
+            ann = ann.clone(img_size=[img.shape[0], img.shape[1]])
+        insts, captions = ann.to_coco(
+            image_idx, class_mapping, coco_ann, label_id, coco_captions, caption_id
+        )
+        label_id += len(insts)
+        caption_id += len(captions)
+
+        if progress_cb is not None:
+            progress_cb(1)
+
+    ann_path = None
+    captions_path = None
+    if save_json is True:
+        logger.info("Saving COCO annotations to disk...")
+        ann_path = str(annotations_dir / COCO_INSTANCES_FILE)
+        dump_json_file(coco_ann, ann_path)
+        logger.info(f"Saved COCO instances to '{ann_path}'")
+
+        if with_captions is True:
+            captions_path = str(annotations_dir / COCO_CAPTIONS_FILE)
+            dump_json_file(coco_captions, captions_path)
+            logger.info(f"Saved COCO captions to '{captions_path}'")
+
+            return ann_path, captions_path
+        return ann_path
+    if with_captions:
+        return coco_ann, coco_captions
+    return coco_ann
+
+
+def sly_project_to_coco(
+    project: Union[Project, str],
+    dest_dir: Optional[str] = None,
+    copy_images: bool = False,
+    with_captions: bool = False,
+    log_progress: bool = True,
+    progress_cb: Optional[Callable] = None,
+) -> None:
+    """
+    Convert Supervisely project to COCO format.
+
+    :param project: Supervisely project.
+    :type project: :class:`Project<supervisely.project.project.Project>` or :class:`str`
+    :param dest_dir: Destination directory.
+    :type dest_dir: :class:`str`, optional
+    :param copy_images: Copy images to destination directory.
+    :type copy_images: :class:`bool`, optional
+    :param with_captions: Return COCO captions.
+    :type with_captions: :class:`bool`, 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 COCO format
+        sly.Project(project_directory).to_coco(log_progress=True)
+    """
+    if isinstance(project, str):
+        project = Project(project, mode=OpenMode.READ)
+
+    dest_dir = Path(dest_dir) if dest_dir is not None else Path(project.directory).parent / "coco"
+    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 COCO format", total=project.total_items
+        ).update
+
+    used_ds_names = set()
+    for ds in project.datasets:
+        ds: Dataset
+        coco_dir = generate_free_name(used_ds_names, ds.short_name, extend_used_names=True)
+        ds.to_coco(
+            meta=project.meta,
+            return_type="path",
+            dest_dir=dest_dir / coco_dir,
+            copy_images=copy_images,
+            with_captions=with_captions,
+            log_progress=log_progress,
+            progress_cb=progress_cb,
+        )
+        logger.info(f"Dataset '{ds.short_name}' has been converted to COCO format.")
+    logger.info(f"Project '{project.name}' has been converted to COCO format.")