supervisely 6.73.375__py3-none-any.whl → 6.73.377__py3-none-any.whl

supervisely/video/sampling.py

@@ -0,0 +1,550 @@
+from typing import Dict, List, Union
+
+import cv2
+import numpy as np
+
+from supervisely._utils import batched_iter
+from supervisely.annotation.annotation import Annotation
+from supervisely.annotation.label import Label
+from supervisely.annotation.tag import Tag
+from supervisely.annotation.tag_collection import TagCollection
+from supervisely.annotation.tag_meta import TagApplicableTo, TagMeta, TagValueType
+from supervisely.api.api import Api
+from supervisely.api.dataset_api import DatasetInfo
+from supervisely.api.image_api import ImageInfo
+from supervisely.api.project_api import ProjectInfo
+from supervisely.api.video.video_api import VideoInfo
+from supervisely.project.project_meta import ProjectMeta
+from supervisely.task.progress import tqdm_sly
+from supervisely.video.video import VideoFrameReader
+from supervisely.video_annotation.frame import Frame
+from supervisely.video_annotation.key_id_map import KeyIdMap
+from supervisely.video_annotation.video_annotation import VideoAnnotation
+
+VIDEO_OBJECT_TAG_META = TagMeta(
+    "object_id",
+    value_type=TagValueType.ANY_NUMBER,
+    applicable_to=TagApplicableTo.OBJECTS_ONLY,
+)
+AUTO_TRACKED_TAG_META = TagMeta(
+    "auto-tracked", TagValueType.NONE, applicable_to=TagApplicableTo.OBJECTS_ONLY
+)
+
+
+class ApiContext:
+    def __init__(self):
+        self.project_info: Dict[int, ProjectInfo] = {}
+        self.project_meta: Dict[int, ProjectMeta] = {}
+        self.dataset_info: Dict[int, DatasetInfo] = {}
+        self.video_info: Dict[int, VideoInfo] = {}
+        self.children_datasets: Dict[int, List[DatasetInfo]] = {}
+        self.children_items: Dict[int, List[Union[ImageInfo, VideoInfo]]] = {}
+
+    def __getattr__(self, item: str) -> Dict:
+        if not hasattr(self, item):
+            new_dict = {}
+            setattr(self, item, new_dict)
+        return getattr(self, item)
+
+
+class SamplingSettings:
+    ONLY_ANNOTATED = "annotated"
+    START = "start"
+    END = "end"
+    STEP = "step"
+    RESIZE = "resize"
+    COPY_ANNOTATIONS = "copy_annotations"
+
+
+def _get_frame_indices(frames_count, start, end, step, only_annotated, video_annotation):
+    frame_indices = list(range(start, end if end is not None else frames_count))
+    if only_annotated:
+        annotated_frames = set()
+        for frame in video_annotation.frames:
+            frame: Frame
+            if frame.figures:
+                annotated_frames.add(frame.index)
+        frame_indices = [idx for idx in frame_indices if idx in annotated_frames]
+    frame_indices = [frame_indices[i] for i in range(0, len(frame_indices), step)]
+    return frame_indices
+
+
+def _frame_to_annotation(frame: Frame, video_annotation: VideoAnnotation) -> Annotation:
+
+    labels = []
+    for figure in frame.figures:
+        tags = []
+        video_object = figure.parent_object
+        obj_class = video_object.obj_class
+        geometry = figure.geometry
+        for tag in video_object.tags:
+            if tag.frame_range is None or tag.frame_range[0] <= frame.index <= tag.frame_range[1]:
+                tags.append(Tag(tag.meta, tag.value, labeler_login=tag.labeler_login))
+        tags.append(Tag(VIDEO_OBJECT_TAG_META, video_object.class_id))
+        if figure.track_id is not None:
+            tags.append(Tag(AUTO_TRACKED_TAG_META, None, labeler_login=video_object.labeler_login))
+        label = Label(geometry, obj_class, TagCollection(tags))
+        labels.append(label)
+    img_tags = []
+    for tag in video_annotation.tags:
+        if tag.frame_range is None or tag.frame_range[0] <= frame.index <= tag.frame_range[1]:
+            img_tags.append(Tag(tag.meta, tag.value, labeler_login=tag.labeler_login))
+    return Annotation(video_annotation.img_size, labels=labels, img_tags=TagCollection(img_tags))
+
+
+def _upload_annotations(api: Api, image_ids, frame_indices, video_annotation: VideoAnnotation):
+    anns = []
+    for image_id, frame_index in zip(image_ids, frame_indices):
+        frame = video_annotation.frames.get(frame_index, None)
+        if frame is not None:
+            anns.append(_frame_to_annotation(frame, video_annotation))
+    api.annotation.upload_anns(image_ids, anns=anns)
+
+
+def _upload_frames(
+    api: Api,
+    frames: List[np.ndarray],
+    video_name: str,
+    video_frames_count: int,
+    indices: List[int],
+    dataset_id: int,
+    sample_info: Dict = None,
+    context: ApiContext = None,
+    copy_annotations: bool = False,
+    video_annotation: VideoAnnotation = None,
+) -> List[int]:
+    if sample_info is None:
+        sample_info = {}
+    if context is not None:
+        if dataset_id not in context.children_items:
+            context.children_items[dataset_id] = api.image.get_list(dataset_id)
+        existing_images = context.children_items[dataset_id]
+    else:
+        existing_images = api.image.get_list(dataset_id)
+
+    name_to_info = {image.name: image for image in existing_images}
+
+    image_ids = [None] * len(frames)
+    to_upload = []
+    for i, index in enumerate(indices):
+        digits = len(str(video_frames_count))
+        image_name = f"{video_name}_frame_{str(index).zfill(digits)}.png"
+        if image_name in name_to_info:
+            image_ids[i] = name_to_info[image_name].id
+        else:
+            to_upload.append((image_name, i))
+
+    if to_upload:
+        frames = [frames[i] for _, i in to_upload]
+        names = [name for name, _ in to_upload]
+        metas = [{**sample_info, "frame_index": i} for _, i in to_upload]
+        uploaded = api.image.upload_nps(
+            dataset_id=dataset_id,
+            names=names,
+            imgs=frames,
+            metas=metas,
+        )
+
+        for image_info, (_, i) in zip(uploaded, to_upload):
+            image_ids[i] = image_info.id
+
+    if copy_annotations:
+        _upload_annotations(api, image_ids, indices, video_annotation)
+
+    return image_ids
+
+
+def sample_video(
+    api: Api,
+    video_id: int,
+    dst_dataset_info: DatasetInfo,
+    settings: Dict,
+    sample_info: Dict = None,
+    context: ApiContext = None,
+    progress: tqdm_sly = None,
+):
+    dst_parent_info = dst_dataset_info
+    only_annotated = settings.get(SamplingSettings.ONLY_ANNOTATED, False)
+    start_frame = settings.get(SamplingSettings.START, 0)
+    end_frame = settings.get(SamplingSettings.END, None)
+    step = settings.get(SamplingSettings.STEP, 1)
+    resize = settings.get(SamplingSettings.RESIZE, None)
+    copy_annotations = settings.get(SamplingSettings.COPY_ANNOTATIONS, False)
+
+    if context is None:
+        context = ApiContext()
+    if video_id not in context.video_info:
+        context.video_info[video_id] = api.video.get_info_by_id(video_id)
+    video_info = context.video_info[video_id]
+
+    project_id = video_info.project_id
+    if project_id not in context.project_meta:
+        context.project_meta[project_id] = ProjectMeta.from_json(api.project.get_meta(project_id))
+    project_meta = context.project_meta[project_id]
+
+    video_annotation = VideoAnnotation.from_json(
+        api.video.annotation.download(video_info.id), project_meta, key_id_map=KeyIdMap()
+    )
+
+    progress_cb = None
+    if progress is not None:
+        size = int(video_info.file_meta["size"])
+        progress.reset(size)
+        progress.unit = "B"
+        progress.unit_scale = True
+        progress.unit_divisor = 1024
+        progress.message = f"Downloading {video_info.name} [{video_info.id}]"
+        progress.desc = progress.message
+        progress.refresh()
+        progress_cb = progress.update
+
+    video_path = f"/tmp/{video_info.name}"
+    api.video.download_path(video_info.id, video_path, progress_cb=progress_cb)
+
+    frame_indices = _get_frame_indices(
+        video_info.frames_count, start_frame, end_frame, step, only_annotated, video_annotation
+    )
+
+    dst_dataset_info = _get_or_create_dst_dataset(
+        api=api,
+        src_info=video_info,
+        dst_parent_info=dst_parent_info,
+        sample_info=sample_info,
+        context=context,
+    )
+
+    if progress is not None:
+        progress.reset(len(frame_indices))
+        progress.unit = "it"
+        progress.unit_scale = False
+        progress.unit_divisor = 1000
+        progress.message = f"Processing {video_info.name} [{video_info.id}]"
+        progress.desc = progress.message
+        progress.miniters = 1
+        progress.refresh()
+
+    with VideoFrameReader(video_path, frame_indices) as reader:
+        for batch in batched_iter(zip(reader, frame_indices), 10):
+            frames, indices = zip(*batch)
+            for frame in frames:
+                if resize:
+                    cv2.resize(frame, [*resize, frame.shape[2]], interpolation=cv2.INTER_LINEAR)
+
+            image_ids = _upload_frames(
+                api=api,
+                frames=frames,
+                video_name=video_info.name,
+                video_frames_count=video_info.frames_count,
+                indices=indices,
+                dataset_id=dst_dataset_info.id,
+                sample_info=sample_info,
+                context=context,
+                copy_annotations=copy_annotations,
+                video_annotation=video_annotation,
+            )
+
+            if progress is not None:
+                progress.update(len(image_ids))
+
+
+def _get_or_create_dst_dataset(
+    api: Api,
+    src_info: Union[DatasetInfo, VideoInfo],
+    dst_parent_info: Union[ProjectInfo, DatasetInfo],
+    sample_info: Dict = None,
+    context: ApiContext = None,
+) -> DatasetInfo:
+    if isinstance(dst_parent_info, ProjectInfo):
+        parent = None
+        project_id = dst_parent_info.id
+    else:
+        parent = dst_parent_info.id
+        project_id = dst_parent_info.project_id
+
+    if context is not None:
+        if dst_parent_info.id not in context.children_datasets:
+            context.children_datasets[dst_parent_info.id] = api.dataset.get_list(
+                project_id, parent_id=parent
+            )
+        dst_datasets = context.children_datasets[dst_parent_info.id]
+    else:
+        dst_datasets = api.dataset.get_list(project_id, parent_id=parent)
+
+    for dst_dataset in dst_datasets:
+        if dst_dataset.name == src_info.name:
+            return dst_dataset
+
+    if isinstance(src_info, DatasetInfo):
+        src_dataset_info = src_info
+        description = (
+            f"Sample dataset made from project #{src_info.parent_id}, dataset #{src_info.id}"
+        )
+    else:
+        if context is not None:
+            if src_info.dataset_id not in context.dataset_info:
+                context.dataset_info[src_info.dataset_id] = api.dataset.get_info_by_id(
+                    src_info.dataset_id
+                )
+            src_dataset_info = context.dataset_info[src_info.dataset_id]
+        else:
+            src_dataset_info = api.dataset.get_info_by_id(src_info.dataset_id)
+        description = f"Sample dataset made from project #{src_info.project_id}, dataset #{src_info.dataset_id}, video #{src_info.id}"
+
+    dst_dataset = api.dataset.create(
+        project_id,
+        name=src_info.name,
+        description=description,
+        parent_id=parent,
+    )
+    if sample_info is None:
+        if context is not None:
+            if src_info.project_id not in context.project_info:
+                context.project_info[src_info.project_id] = api.project.get_info_by_id(
+                    src_info.project_id
+                )
+            src_project_info = context.project_info[src_info.project_id]
+        else:
+            src_project_info = api.project.get_info_by_id(src_info.project_id)
+        sample_info = {
+            "is_sample": True,
+            "video_project_id": src_info.project_id,
+            "video_project_name": src_project_info.name,
+        }
+    sample_info.update(
+        {
+            "video_dataset_id": src_dataset_info.id,
+            "video_dataset_name": src_dataset_info.name,
+        }
+    )
+    if isinstance(src_info, VideoInfo):
+        sample_info.update(
+            {
+                "video_id": src_info.id,
+                "video_name": src_info.name,
+            }
+        )
+    api.dataset.update_custom_data(
+        dst_dataset.id,
+        custom_data=sample_info,
+    )
+    if context is not None:
+        context.dataset_info[dst_dataset.id] = dst_dataset
+    return dst_dataset
+
+
+def sample_video_dataset(
+    api: Api,
+    src_dataset_id: int,
+    dst_parent_info: Union[ProjectInfo, DatasetInfo],
+    settings: Dict,
+    sample_info: Dict = None,
+    context: ApiContext = None,
+    datasets_ids_whitelist: List[int] = None,
+    items_progress_cb: tqdm_sly = None,
+    video_progress: tqdm_sly = None,
+):
+    if context is None:
+        context = ApiContext()
+
+    if not (
+        datasets_ids_whitelist is None
+        or src_dataset_id in datasets_ids_whitelist  # this dataset should be sampled
+        or _has_children_datasets(
+            api, src_dataset_id, context=context, children_ids=datasets_ids_whitelist
+        )  # has children datasets that should be sampled
+    ):
+        return None
+
+    src_dataset_info = context.dataset_info.get(src_dataset_id, None)
+    if src_dataset_info is None:
+        src_dataset_info = api.dataset.get_info_by_id(src_dataset_id)
+        context.dataset_info[src_dataset_id] = src_dataset_info
+
+    dst_dataset = _get_or_create_dst_dataset(
+        api=api,
+        src_info=src_dataset_info,
+        dst_parent_info=dst_parent_info,
+        sample_info=sample_info,
+        context=context,
+    )
+
+    if datasets_ids_whitelist is None or src_dataset_id in datasets_ids_whitelist:
+        video_infos = api.video.get_list(src_dataset_id)
+        for video_info in video_infos:
+            sample_video(
+                api=api,
+                video_id=video_info.id,
+                dst_dataset_info=dst_dataset,
+                settings=settings,
+                sample_info=sample_info.copy(),
+                context=context,
+                progress=video_progress,
+            )
+            if items_progress_cb is not None:
+                items_progress_cb()
+
+    if src_dataset_id not in context.children_datasets:
+        if src_dataset_id not in context.dataset_info:
+            context.dataset_info[src_dataset_id] = api.dataset.get_info_by_id(src_dataset_id)
+        src_dataset_info = context.dataset_info[src_dataset_id]
+        context.children_datasets[src_dataset_id] = api.dataset.get_list(
+            src_dataset_info.project_id, parent_id=src_dataset_info.id
+        )
+    datasets = context.children_datasets[src_dataset_id]
+    for dataset in datasets:
+        sample_video_dataset(
+            api=api,
+            src_dataset_id=dataset.id,
+            dst_parent_info=dst_dataset,
+            settings=settings,
+            datasets_ids_whitelist=datasets_ids_whitelist,
+            sample_info=sample_info,
+            context=context,
+            items_progress_cb=items_progress_cb,
+            video_progress=video_progress,
+        )
+    return dst_dataset
+
+
+def _update_meta(
+    api: Api, src_project_meta: ProjectMeta, dst_project_id: int, context: ApiContext = None
+):
+    if context is None:
+        context = ApiContext()
+    if dst_project_id not in context.project_meta:
+        context.project_meta[dst_project_id] = ProjectMeta.from_json(
+            api.project.get_meta(dst_project_id)
+        )
+    dst_project_meta = context.project_meta[dst_project_id]
+    dst_project_meta.merge(src_project_meta)
+    if dst_project_meta.get_tag_meta(VIDEO_OBJECT_TAG_META.name) is None:
+        dst_project_meta.add_tag_meta(VIDEO_OBJECT_TAG_META)
+    if dst_project_meta.get_tag_meta(AUTO_TRACKED_TAG_META.name) is None:
+        dst_project_meta.add_tag_meta(AUTO_TRACKED_TAG_META)
+
+    if dst_project_meta != src_project_meta:
+        api.project.update_meta(dst_project_id, dst_project_meta.to_json())
+        context.project_meta[dst_project_id] = dst_project_meta
+
+
+def _get_or_create_dst_project(
+    api: Api,
+    src_project_id: int,
+    dst_project_id: Union[int, None] = None,
+    sample_info: Dict = None,
+    context: ApiContext = None,
+) -> ProjectInfo:
+    if dst_project_id is None:
+        # get source project info
+        if context is None:
+            src_project_info = api.project.get_info_by_id(src_project_id)
+        else:
+            if src_project_id not in context.project_info:
+                context.project_info[src_project_id] = api.project.get_info_by_id(src_project_id)
+            src_project_info = context.project_info[src_project_id]
+        # create new project
+        if sample_info is None:
+            sample_info = {}
+        sample_info.update(
+            {
+                "is_sample": True,
+                "video_project_id": src_project_id,
+                "video_project_name": src_project_info.name,
+            }
+        )
+        dst_project = api.project.create(
+            src_project_info.workspace_id,
+            f"{src_project_info.name}(images)",
+            description=f"Sample project made from project #{src_project_info.id}",
+            change_name_if_conflict=True,
+        )
+        api.project.update_custom_data(dst_project.id, sample_info)
+    else:
+        # use existing project
+        dst_project = api.project.get_info_by_id(dst_project_id)
+    if context is not None:
+        context.project_info[dst_project.id] = dst_project
+    if src_project_id not in context.project_meta:
+        context.project_meta[src_project_id] = ProjectMeta.from_json(
+            api.project.get_meta(src_project_id)
+        )
+    src_project_meta = context.project_meta[src_project_id]
+    _update_meta(api, src_project_meta, dst_project.id, context=context)
+    return dst_project
+
+
+def _has_children_datasets(
+    api: Api, dataset_id: int, context: ApiContext = None, children_ids: List[int] = None
+) -> bool:
+    if context is None:
+        context = ApiContext()
+
+    if dataset_id not in context.dataset_info:
+        context.dataset_info[dataset_id] = api.dataset.get_info_by_id(dataset_id)
+    dataset_info = context.dataset_info[dataset_id]
+    if dataset_id not in context.children_datasets:
+        context.children_datasets[dataset_id] = api.dataset.get_list(
+            project_id=dataset_info.project_id, parent_id=dataset_id
+        )
+
+    if children_ids is None:
+        return len(context.children_datasets[dataset_id]) > 0
+    for child_dataset in context.children_datasets[dataset_id]:
+        if child_dataset.id in children_ids:
+            return True
+        if _has_children_datasets(
+            api, child_dataset.id, context=context, children_ids=children_ids
+        ):
+            return True
+    return False
+
+
+def sample_video_project(
+    api: Api,
+    project_id: int,
+    settings: Dict,
+    dst_project_id: Union[int, None] = None,
+    datasets_ids: List[int] = None,
+    context: ApiContext = None,
+    items_progress_cb: tqdm_sly = None,
+    video_progress: tqdm_sly = None,
+):
+    if context is None:
+        context = ApiContext()
+
+    if project_id not in context.project_info:
+        context.project_info[project_id] = api.project.get_info_by_id(project_id)
+
+    src_project_info = context.project_info[project_id]
+
+    sample_info = {
+        "is_sample": True,
+        "video_project_id": src_project_info.id,
+        "video_project_name": src_project_info.name,
+    }
+    dst_project_info = _get_or_create_dst_project(
+        api, project_id, dst_project_id, sample_info, context
+    )
+
+    # non recursive. Nested datasets are handled by sample_video_dataset
+    if project_id not in context.children_datasets:
+        datasets = api.dataset.get_list(project_id)
+        context.children_datasets[project_id] = datasets
+
+    dataset_infos = context.children_datasets[project_id]
+    for dataset_info in dataset_infos:
+        context.dataset_info[dataset_info.id] = dataset_info
+    for dataset_info in dataset_infos:
+        sample_video_dataset(
+            api=api,
+            src_dataset_id=dataset_info.id,
+            dst_parent_info=dst_project_info,
+            settings=settings,
+            sample_info=sample_info,
+            context=context,
+            datasets_ids_whitelist=datasets_ids,
+            items_progress_cb=items_progress_cb,
+            video_progress=video_progress,
+        )
+
+    return dst_project_info

{supervisely-6.73.375.dist-info → supervisely-6.73.377.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: supervisely
-Version: 6.73.375
+Version: 6.73.377
 Summary: Supervisely Python SDK.
 Home-page: https://github.com/supervisely/supervisely
 Author: Supervisely