nnInteractive 2.0.0__py3-none-any.whl

nnInteractive/supervoxel/src/sam2/training/dataset/vos_dataset.py

@@ -0,0 +1,154 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+import random
+from copy import deepcopy
+
+import numpy as np
+
+import torch
+from iopath.common.file_io import g_pathmgr
+from PIL import Image as PILImage
+from torchvision.datasets.vision import VisionDataset
+
+from training.dataset.vos_raw_dataset import VOSRawDataset
+from training.dataset.vos_sampler import VOSSampler
+from training.dataset.vos_segment_loader import JSONSegmentLoader
+
+from training.utils.data_utils import Frame, Object, VideoDatapoint
+
+MAX_RETRIES = 100
+
+
+class VOSDataset(VisionDataset):
+    def __init__(
+        self,
+        transforms,
+        training: bool,
+        video_dataset: VOSRawDataset,
+        sampler: VOSSampler,
+        multiplier: int,
+        always_target=True,
+        target_segments_available=True,
+    ):
+        self._transforms = transforms
+        self.training = training
+        self.video_dataset = video_dataset
+        self.sampler = sampler
+
+        self.repeat_factors = torch.ones(len(self.video_dataset), dtype=torch.float32)
+        self.repeat_factors *= multiplier
+        print(f"Raw dataset length = {len(self.video_dataset)}")
+
+        self.curr_epoch = 0  # Used in case data loader behavior changes across epochs
+        self.always_target = always_target
+        self.target_segments_available = target_segments_available
+
+    def _get_datapoint(self, idx):
+
+        for retry in range(MAX_RETRIES):
+            try:
+                if isinstance(idx, torch.Tensor):
+                    idx = idx.item()
+                # sample a video
+                video, segment_loader = self.video_dataset.get_video(idx)
+                # sample frames and object indices to be used in a datapoint
+                sampled_frms_and_objs = self.sampler.sample(video, segment_loader, epoch=self.curr_epoch)
+                break  # Succesfully loaded video
+            except Exception as e:
+                if self.training:
+                    logging.warning(f"Loading failed (id={idx}); Retry {retry} with exception: {e}")
+                    idx = random.randrange(0, len(self.video_dataset))
+                else:
+                    # Shouldn't fail to load a val video
+                    raise e
+
+        datapoint = self.construct(video, sampled_frms_and_objs, segment_loader)
+        for transform in self._transforms:
+            datapoint = transform(datapoint, epoch=self.curr_epoch)
+        return datapoint
+
+    def construct(self, video, sampled_frms_and_objs, segment_loader):
+        """
+        Constructs a VideoDatapoint sample to pass to transforms
+        """
+        sampled_frames = sampled_frms_and_objs.frames
+        sampled_object_ids = sampled_frms_and_objs.object_ids
+
+        images = []
+        rgb_images = load_images(sampled_frames)
+        # Iterate over the sampled frames and store their rgb data and object data (bbox, segment)
+        for frame_idx, frame in enumerate(sampled_frames):
+            w, h = rgb_images[frame_idx].size
+            images.append(
+                Frame(
+                    data=rgb_images[frame_idx],
+                    objects=[],
+                )
+            )
+            # We load the gt segments associated with the current frame
+            if isinstance(segment_loader, JSONSegmentLoader):
+                segments = segment_loader.load(frame.frame_idx, obj_ids=sampled_object_ids)
+            else:
+                segments = segment_loader.load(frame.frame_idx)
+            for obj_id in sampled_object_ids:
+                # Extract the segment
+                if obj_id in segments:
+                    assert segments[obj_id] is not None, "None targets are not supported"
+                    # segment is uint8 and remains uint8 throughout the transforms
+                    segment = segments[obj_id].to(torch.uint8)
+                else:
+                    # There is no target, we either use a zero mask target or drop this object
+                    if not self.always_target:
+                        continue
+                    segment = torch.zeros(h, w, dtype=torch.uint8)
+
+                images[frame_idx].objects.append(
+                    Object(
+                        object_id=obj_id,
+                        frame_index=frame.frame_idx,
+                        segment=segment,
+                    )
+                )
+        return VideoDatapoint(
+            frames=images,
+            video_id=video.video_id,
+            size=(h, w),
+        )
+
+    def __getitem__(self, idx):
+        return self._get_datapoint(idx)
+
+    def __len__(self):
+        return len(self.video_dataset)
+
+
+def load_images(frames):
+    all_images = []
+    cache = {}
+    for frame in frames:
+        if frame.data is None:
+            # Load the frame rgb data from file
+            path = frame.image_path
+            if path in cache:
+                all_images.append(deepcopy(all_images[cache[path]]))
+                continue
+            with g_pathmgr.open(path, "rb") as fopen:
+                all_images.append(PILImage.open(fopen).convert("RGB"))
+            cache[path] = len(all_images) - 1
+        else:
+            # The frame rgb data has already been loaded
+            # Convert it to a PILImage
+            all_images.append(tensor_2_PIL(frame.data))
+
+    return all_images
+
+
+def tensor_2_PIL(data: torch.Tensor) -> PILImage.Image:
+    data = data.cpu().numpy().transpose((1, 2, 0)) * 255.0
+    data = data.astype(np.uint8)
+    return PILImage.fromarray(data)

nnInteractive/supervoxel/src/sam2/training/dataset/vos_raw_dataset.py

@@ -0,0 +1,290 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import glob
+import logging
+import os
+from dataclasses import dataclass
+
+from typing import List, Optional
+
+import pandas as pd
+
+import torch
+
+from iopath.common.file_io import g_pathmgr
+
+from omegaconf.listconfig import ListConfig
+
+from training.dataset.vos_segment_loader import (
+    JSONSegmentLoader,
+    MultiplePNGSegmentLoader,
+    PalettisedPNGSegmentLoader,
+    SA1BSegmentLoader,
+)
+
+
+@dataclass
+class VOSFrame:
+    frame_idx: int
+    image_path: str
+    data: Optional[torch.Tensor] = None
+    is_conditioning_only: Optional[bool] = False
+
+
+@dataclass
+class VOSVideo:
+    video_name: str
+    video_id: int
+    frames: List[VOSFrame]
+
+    def __len__(self):
+        return len(self.frames)
+
+
+class VOSRawDataset:
+    def __init__(self):
+        pass
+
+    def get_video(self, idx):
+        raise NotImplementedError()
+
+
+class PNGRawDataset(VOSRawDataset):
+    def __init__(
+        self,
+        img_folder,
+        gt_folder,
+        file_list_txt=None,
+        excluded_videos_list_txt=None,
+        sample_rate=1,
+        is_palette=True,
+        single_object_mode=False,
+        truncate_video=-1,
+        frames_sampling_mult=False,
+    ):
+        self.img_folder = img_folder
+        self.gt_folder = gt_folder
+        self.sample_rate = sample_rate
+        self.is_palette = is_palette
+        self.single_object_mode = single_object_mode
+        self.truncate_video = truncate_video
+
+        # Read the subset defined in file_list_txt
+        if file_list_txt is not None:
+            with g_pathmgr.open(file_list_txt, "r") as f:
+                subset = [os.path.splitext(line.strip())[0] for line in f]
+        else:
+            subset = os.listdir(self.img_folder)
+
+        # Read and process excluded files if provided
+        if excluded_videos_list_txt is not None:
+            with g_pathmgr.open(excluded_videos_list_txt, "r") as f:
+                excluded_files = [os.path.splitext(line.strip())[0] for line in f]
+        else:
+            excluded_files = []
+
+        # Check if it's not in excluded_files
+        self.video_names = sorted([video_name for video_name in subset if video_name not in excluded_files])
+
+        if self.single_object_mode:
+            # single object mode
+            self.video_names = sorted(
+                [
+                    os.path.join(video_name, obj)
+                    for video_name in self.video_names
+                    for obj in os.listdir(os.path.join(self.gt_folder, video_name))
+                ]
+            )
+
+        if frames_sampling_mult:
+            video_names_mult = []
+            for video_name in self.video_names:
+                num_frames = len(os.listdir(os.path.join(self.img_folder, video_name)))
+                video_names_mult.extend([video_name] * num_frames)
+            self.video_names = video_names_mult
+
+    def get_video(self, idx):
+        """
+        Given a VOSVideo object, return the mask tensors.
+        """
+        video_name = self.video_names[idx]
+
+        if self.single_object_mode:
+            video_frame_root = os.path.join(self.img_folder, os.path.dirname(video_name))
+        else:
+            video_frame_root = os.path.join(self.img_folder, video_name)
+
+        video_mask_root = os.path.join(self.gt_folder, video_name)
+
+        if self.is_palette:
+            segment_loader = PalettisedPNGSegmentLoader(video_mask_root)
+        else:
+            segment_loader = MultiplePNGSegmentLoader(video_mask_root, self.single_object_mode)
+
+        all_frames = sorted(glob.glob(os.path.join(video_frame_root, "*.jpg")))
+        if self.truncate_video > 0:
+            all_frames = all_frames[: self.truncate_video]
+        frames = []
+        for _, fpath in enumerate(all_frames[:: self.sample_rate]):
+            fid = int(os.path.basename(fpath).split(".")[0])
+            frames.append(VOSFrame(fid, image_path=fpath))
+        video = VOSVideo(video_name, idx, frames)
+        return video, segment_loader
+
+    def __len__(self):
+        return len(self.video_names)
+
+
+class SA1BRawDataset(VOSRawDataset):
+    def __init__(
+        self,
+        img_folder,
+        gt_folder,
+        file_list_txt=None,
+        excluded_videos_list_txt=None,
+        num_frames=1,
+        mask_area_frac_thresh=1.1,  # no filtering by default
+        uncertain_iou=-1,  # no filtering by default
+    ):
+        self.img_folder = img_folder
+        self.gt_folder = gt_folder
+        self.num_frames = num_frames
+        self.mask_area_frac_thresh = mask_area_frac_thresh
+        self.uncertain_iou = uncertain_iou  # stability score
+
+        # Read the subset defined in file_list_txt
+        if file_list_txt is not None:
+            with g_pathmgr.open(file_list_txt, "r") as f:
+                subset = [os.path.splitext(line.strip())[0] for line in f]
+        else:
+            subset = os.listdir(self.img_folder)
+            subset = [path.split(".")[0] for path in subset if path.endswith(".jpg")]  # remove extension
+
+        # Read and process excluded files if provided
+        if excluded_videos_list_txt is not None:
+            with g_pathmgr.open(excluded_videos_list_txt, "r") as f:
+                excluded_files = [os.path.splitext(line.strip())[0] for line in f]
+        else:
+            excluded_files = []
+
+        # Check if it's not in excluded_files and it exists
+        self.video_names = [video_name for video_name in subset if video_name not in excluded_files]
+
+    def get_video(self, idx):
+        """
+        Given a VOSVideo object, return the mask tensors.
+        """
+        video_name = self.video_names[idx]
+
+        video_frame_path = os.path.join(self.img_folder, video_name + ".jpg")
+        video_mask_path = os.path.join(self.gt_folder, video_name + ".json")
+
+        segment_loader = SA1BSegmentLoader(
+            video_mask_path,
+            mask_area_frac_thresh=self.mask_area_frac_thresh,
+            video_frame_path=video_frame_path,
+            uncertain_iou=self.uncertain_iou,
+        )
+
+        frames = []
+        for frame_idx in range(self.num_frames):
+            frames.append(VOSFrame(frame_idx, image_path=video_frame_path))
+        video_name = video_name.split("_")[-1]  # filename is sa_{int}
+        # video id needs to be image_id to be able to load correct annotation file during eval
+        video = VOSVideo(video_name, int(video_name), frames)
+        return video, segment_loader
+
+    def __len__(self):
+        return len(self.video_names)
+
+
+class JSONRawDataset(VOSRawDataset):
+    """
+    Dataset where the annotation in the format of SA-V json files
+    """
+
+    def __init__(
+        self,
+        img_folder,
+        gt_folder,
+        file_list_txt=None,
+        excluded_videos_list_txt=None,
+        sample_rate=1,
+        rm_unannotated=True,
+        ann_every=1,
+        frames_fps=24,
+    ):
+        self.gt_folder = gt_folder
+        self.img_folder = img_folder
+        self.sample_rate = sample_rate
+        self.rm_unannotated = rm_unannotated
+        self.ann_every = ann_every
+        self.frames_fps = frames_fps
+
+        # Read and process excluded files if provided
+        excluded_files = []
+        if excluded_videos_list_txt is not None:
+            if isinstance(excluded_videos_list_txt, str):
+                excluded_videos_lists = [excluded_videos_list_txt]
+            elif isinstance(excluded_videos_list_txt, ListConfig):
+                excluded_videos_lists = list(excluded_videos_list_txt)
+            else:
+                raise NotImplementedError
+
+            for excluded_videos_list_txt in excluded_videos_lists:
+                with open(excluded_videos_list_txt, "r") as f:
+                    excluded_files.extend([os.path.splitext(line.strip())[0] for line in f])
+        excluded_files = set(excluded_files)
+
+        # Read the subset defined in file_list_txt
+        if file_list_txt is not None:
+            with g_pathmgr.open(file_list_txt, "r") as f:
+                subset = [os.path.splitext(line.strip())[0] for line in f]
+        else:
+            subset = os.listdir(self.img_folder)
+
+        self.video_names = sorted([video_name for video_name in subset if video_name not in excluded_files])
+
+    def get_video(self, video_idx):
+        """
+        Given a VOSVideo object, return the mask tensors.
+        """
+        video_name = self.video_names[video_idx]
+        video_json_path = os.path.join(self.gt_folder, video_name + "_manual.json")
+        segment_loader = JSONSegmentLoader(
+            video_json_path=video_json_path,
+            ann_every=self.ann_every,
+            frames_fps=self.frames_fps,
+        )
+
+        frame_ids = [
+            int(os.path.splitext(frame_name)[0])
+            for frame_name in sorted(os.listdir(os.path.join(self.img_folder, video_name)))
+        ]
+
+        frames = [
+            VOSFrame(
+                frame_id,
+                image_path=os.path.join(self.img_folder, f"{video_name}/%05d.jpg" % (frame_id)),
+            )
+            for frame_id in frame_ids[:: self.sample_rate]
+        ]
+
+        if self.rm_unannotated:
+            # Eliminate the frames that have not been annotated
+            valid_frame_ids = [
+                i * segment_loader.ann_every
+                for i, annot in enumerate(segment_loader.frame_annots)
+                if annot is not None and None not in annot
+            ]
+            frames = [f for f in frames if f.frame_idx in valid_frame_ids]
+
+        video = VOSVideo(video_name, video_idx, frames)
+        return video, segment_loader
+
+    def __len__(self):
+        return len(self.video_names)

nnInteractive/supervoxel/src/sam2/training/dataset/vos_sampler.py

@@ -0,0 +1,103 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import random
+from dataclasses import dataclass
+from typing import List
+
+from training.dataset.vos_segment_loader import LazySegments
+
+MAX_RETRIES = 1000
+
+
+@dataclass
+class SampledFramesAndObjects:
+    frames: List[int]
+    object_ids: List[int]
+
+
+class VOSSampler:
+    def __init__(self, sort_frames=True):
+        # frames are ordered by frame id when sort_frames is True
+        self.sort_frames = sort_frames
+
+    def sample(self, video):
+        raise NotImplementedError()
+
+
+class RandomUniformSampler(VOSSampler):
+    def __init__(
+        self,
+        num_frames,
+        max_num_objects,
+        reverse_time_prob=0.0,
+    ):
+        self.num_frames = num_frames
+        self.max_num_objects = max_num_objects
+        self.reverse_time_prob = reverse_time_prob
+
+    def sample(self, video, segment_loader, epoch=None):
+
+        for retry in range(MAX_RETRIES):
+            if len(video.frames) < self.num_frames:
+                raise Exception(
+                    f"Cannot sample {self.num_frames} frames from video {video.video_name} as it only has {len(video.frames)} annotated frames."
+                )
+            start = random.randrange(0, len(video.frames) - self.num_frames + 1)
+            frames = [video.frames[start + step] for step in range(self.num_frames)]
+            if random.uniform(0, 1) < self.reverse_time_prob:
+                # Reverse time
+                frames = frames[::-1]
+
+            # Get first frame object ids
+            visible_object_ids = []
+            loaded_segms = segment_loader.load(frames[0].frame_idx)
+            if isinstance(loaded_segms, LazySegments):
+                # LazySegments for SA1BRawDataset
+                visible_object_ids = list(loaded_segms.keys())
+            else:
+                for object_id, segment in segment_loader.load(frames[0].frame_idx).items():
+                    if segment.sum():
+                        visible_object_ids.append(object_id)
+
+            # First frame needs to have at least a target to track
+            if len(visible_object_ids) > 0:
+                break
+            if retry >= MAX_RETRIES - 1:
+                raise Exception("No visible objects")
+
+        object_ids = random.sample(
+            visible_object_ids,
+            min(len(visible_object_ids), self.max_num_objects),
+        )
+        return SampledFramesAndObjects(frames=frames, object_ids=object_ids)
+
+
+class EvalSampler(VOSSampler):
+    """
+    VOS Sampler for evaluation: sampling all the frames and all the objects in a video
+    """
+
+    def __init__(
+        self,
+    ):
+        super().__init__()
+
+    def sample(self, video, segment_loader, epoch=None):
+        """
+        Sampling all the frames and all the objects
+        """
+        if self.sort_frames:
+            # ordered by frame id
+            frames = sorted(video.frames, key=lambda x: x.frame_idx)
+        else:
+            # use the original order
+            frames = video.frames
+        object_ids = segment_loader.load(frames[0].frame_idx).keys()
+        if len(object_ids) == 0:
+            raise Exception("First frame of the video has no objects")
+
+        return SampledFramesAndObjects(frames=frames, object_ids=object_ids)