nnInteractive 2.0.0__py3-none-any.whl

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

@@ -0,0 +1,289 @@
+# 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 json
+import os
+
+import numpy as np
+import pandas as pd
+import torch
+
+from PIL import Image as PILImage
+
+try:
+    from pycocotools import mask as mask_utils
+except:
+    pass
+
+
+class JSONSegmentLoader:
+    def __init__(self, video_json_path, ann_every=1, frames_fps=24, valid_obj_ids=None):
+        # Annotations in the json are provided every ann_every th frame
+        self.ann_every = ann_every
+        # Ids of the objects to consider when sampling this video
+        self.valid_obj_ids = valid_obj_ids
+        with open(video_json_path, "r") as f:
+            data = json.load(f)
+            if isinstance(data, list):
+                self.frame_annots = data
+            elif isinstance(data, dict):
+                masklet_field_name = "masklet" if "masklet" in data else "masks"
+                self.frame_annots = data[masklet_field_name]
+                if "fps" in data:
+                    if isinstance(data["fps"], list):
+                        annotations_fps = int(data["fps"][0])
+                    else:
+                        annotations_fps = int(data["fps"])
+                    assert frames_fps % annotations_fps == 0
+                    self.ann_every = frames_fps // annotations_fps
+            else:
+                raise NotImplementedError
+
+    def load(self, frame_id, obj_ids=None):
+        assert frame_id % self.ann_every == 0
+        rle_mask = self.frame_annots[frame_id // self.ann_every]
+
+        valid_objs_ids = set(range(len(rle_mask)))
+        if self.valid_obj_ids is not None:
+            # Remove the masklets that have been filtered out for this video
+            valid_objs_ids &= set(self.valid_obj_ids)
+        if obj_ids is not None:
+            # Only keep the objects that have been sampled
+            valid_objs_ids &= set(obj_ids)
+        valid_objs_ids = sorted(list(valid_objs_ids))
+
+        # Construct rle_masks_filtered that only contains the rle masks we are interested in
+        id_2_idx = {}
+        rle_mask_filtered = []
+        for obj_id in valid_objs_ids:
+            if rle_mask[obj_id] is not None:
+                id_2_idx[obj_id] = len(rle_mask_filtered)
+                rle_mask_filtered.append(rle_mask[obj_id])
+            else:
+                id_2_idx[obj_id] = None
+
+        # Decode the masks
+        raw_segments = torch.from_numpy(mask_utils.decode(rle_mask_filtered)).permute(2, 0, 1)  # （num_obj, h, w）
+        segments = {}
+        for obj_id in valid_objs_ids:
+            if id_2_idx[obj_id] is None:
+                segments[obj_id] = None
+            else:
+                idx = id_2_idx[obj_id]
+                segments[obj_id] = raw_segments[idx]
+        return segments
+
+    def get_valid_obj_frames_ids(self, num_frames_min=None):
+        # For each object, find all the frames with a valid (not None) mask
+        num_objects = len(self.frame_annots[0])
+
+        # The result dict associates each obj_id with the id of its valid frames
+        res = {obj_id: [] for obj_id in range(num_objects)}
+
+        for annot_idx, annot in enumerate(self.frame_annots):
+            for obj_id in range(num_objects):
+                if annot[obj_id] is not None:
+                    res[obj_id].append(int(annot_idx * self.ann_every))
+
+        if num_frames_min is not None:
+            # Remove masklets that have less than num_frames_min valid masks
+            for obj_id, valid_frames in list(res.items()):
+                if len(valid_frames) < num_frames_min:
+                    res.pop(obj_id)
+
+        return res
+
+
+class PalettisedPNGSegmentLoader:
+    def __init__(self, video_png_root):
+        """
+        SegmentLoader for datasets with masks stored as palettised PNGs.
+        video_png_root: the folder contains all the masks stored in png
+        """
+        self.video_png_root = video_png_root
+        # build a mapping from frame id to their PNG mask path
+        # note that in some datasets, the PNG paths could have more
+        # than 5 digits, e.g. "00000000.png" instead of "00000.png"
+        png_filenames = os.listdir(self.video_png_root)
+        self.frame_id_to_png_filename = {}
+        for filename in png_filenames:
+            frame_id, _ = os.path.splitext(filename)
+            self.frame_id_to_png_filename[int(frame_id)] = filename
+
+    def load(self, frame_id):
+        """
+        load the single palettised mask from the disk (path: f'{self.video_png_root}/{frame_id:05d}.png')
+        Args:
+            frame_id: int, define the mask path
+        Return:
+            binary_segments: dict
+        """
+        # check the path
+        mask_path = os.path.join(self.video_png_root, self.frame_id_to_png_filename[frame_id])
+
+        # load the mask
+        masks = PILImage.open(mask_path).convert("P")
+        masks = np.array(masks)
+
+        object_id = pd.unique(masks.flatten())
+        object_id = object_id[object_id != 0]  # remove background (0)
+
+        # convert into N binary segmentation masks
+        binary_segments = {}
+        for i in object_id:
+            bs = masks == i
+            binary_segments[i] = torch.from_numpy(bs)
+
+        return binary_segments
+
+    def __len__(self):
+        return
+
+
+class MultiplePNGSegmentLoader:
+    def __init__(self, video_png_root, single_object_mode=False):
+        """
+        video_png_root: the folder contains all the masks stored in png
+        single_object_mode: whether to load only a single object at a time
+        """
+        self.video_png_root = video_png_root
+        self.single_object_mode = single_object_mode
+        # read a mask to know the resolution of the video
+        if self.single_object_mode:
+            tmp_mask_path = glob.glob(os.path.join(video_png_root, "*.png"))[0]
+        else:
+            tmp_mask_path = glob.glob(os.path.join(video_png_root, "*", "*.png"))[0]
+        tmp_mask = np.array(PILImage.open(tmp_mask_path))
+        self.H = tmp_mask.shape[0]
+        self.W = tmp_mask.shape[1]
+        if self.single_object_mode:
+            self.obj_id = int(video_png_root.split("/")[-1]) + 1  # offset by 1 as bg is 0
+        else:
+            self.obj_id = None
+
+    def load(self, frame_id):
+        if self.single_object_mode:
+            return self._load_single_png(frame_id)
+        else:
+            return self._load_multiple_pngs(frame_id)
+
+    def _load_single_png(self, frame_id):
+        """
+        load single png from the disk (path: f'{self.obj_id}/{frame_id:05d}.png')
+        Args:
+            frame_id: int, define the mask path
+        Return:
+            binary_segments: dict
+        """
+        mask_path = os.path.join(self.video_png_root, f"{frame_id:05d}.png")
+        binary_segments = {}
+
+        if os.path.exists(mask_path):
+            mask = np.array(PILImage.open(mask_path))
+        else:
+            # if png doesn't exist, empty mask
+            mask = np.zeros((self.H, self.W), dtype=bool)
+        binary_segments[self.obj_id] = torch.from_numpy(mask > 0)
+        return binary_segments
+
+    def _load_multiple_pngs(self, frame_id):
+        """
+        load multiple png masks from the disk (path: f'{obj_id}/{frame_id:05d}.png')
+        Args:
+            frame_id: int, define the mask path
+        Return:
+            binary_segments: dict
+        """
+        # get the path
+        all_objects = sorted(glob.glob(os.path.join(self.video_png_root, "*")))
+        num_objects = len(all_objects)
+        assert num_objects > 0
+
+        # load the masks
+        binary_segments = {}
+        for obj_folder in all_objects:
+            # obj_folder is {video_name}/{obj_id}, obj_id is specified by the name of the folder
+            obj_id = int(obj_folder.split("/")[-1])
+            obj_id = obj_id + 1  # offset 1 as bg is 0
+            mask_path = os.path.join(obj_folder, f"{frame_id:05d}.png")
+            if os.path.exists(mask_path):
+                mask = np.array(PILImage.open(mask_path))
+            else:
+                mask = np.zeros((self.H, self.W), dtype=bool)
+            binary_segments[obj_id] = torch.from_numpy(mask > 0)
+
+        return binary_segments
+
+    def __len__(self):
+        return
+
+
+class LazySegments:
+    """
+    Only decodes segments that are actually used.
+    """
+
+    def __init__(self):
+        self.segments = {}
+        self.cache = {}
+
+    def __setitem__(self, key, item):
+        self.segments[key] = item
+
+    def __getitem__(self, key):
+        if key in self.cache:
+            return self.cache[key]
+        rle = self.segments[key]
+        mask = torch.from_numpy(mask_utils.decode([rle])).permute(2, 0, 1)[0]
+        self.cache[key] = mask
+        return mask
+
+    def __contains__(self, key):
+        return key in self.segments
+
+    def __len__(self):
+        return len(self.segments)
+
+    def keys(self):
+        return self.segments.keys()
+
+
+class SA1BSegmentLoader:
+    def __init__(
+        self,
+        video_mask_path,
+        mask_area_frac_thresh=1.1,
+        video_frame_path=None,
+        uncertain_iou=-1,
+    ):
+        with open(video_mask_path, "r") as f:
+            self.frame_annots = json.load(f)
+
+        if mask_area_frac_thresh <= 1.0:
+            # Lazily read frame
+            orig_w, orig_h = PILImage.open(video_frame_path).size
+            area = orig_w * orig_h
+
+        self.frame_annots = self.frame_annots["annotations"]
+
+        rle_masks = []
+        for frame_annot in self.frame_annots:
+            if not frame_annot["area"] > 0:
+                continue
+            if ("uncertain_iou" in frame_annot) and (frame_annot["uncertain_iou"] < uncertain_iou):
+                # uncertain_iou is stability score
+                continue
+            if mask_area_frac_thresh <= 1.0 and (frame_annot["area"] / area) >= mask_area_frac_thresh:
+                continue
+            rle_masks.append(frame_annot["segmentation"])
+
+        self.segments = LazySegments()
+        for i, rle in enumerate(rle_masks):
+            self.segments[i] = rle
+
+    def load(self, frame_idx):
+        return self.segments

nnInteractive/supervoxel/src/sam2/training/loss_fns.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.
+
+from collections import defaultdict
+from typing import Dict, List
+
+import torch
+import torch.distributed
+import torch.nn as nn
+import torch.nn.functional as F
+
+from training.trainer import CORE_LOSS_KEY
+
+from training.utils.distributed import get_world_size, is_dist_avail_and_initialized
+
+
+def dice_loss(inputs, targets, num_objects, loss_on_multimask=False):
+    """
+    Compute the DICE loss, similar to generalized IOU for masks
+    Args:
+        inputs: A float tensor of arbitrary shape.
+                The predictions for each example.
+        targets: A float tensor with the same shape as inputs. Stores the binary
+                 classification label for each element in inputs
+                (0 for the negative class and 1 for the positive class).
+        num_objects: Number of objects in the batch
+        loss_on_multimask: True if multimask prediction is enabled
+    Returns:
+        Dice loss tensor
+    """
+    inputs = inputs.sigmoid()
+    if loss_on_multimask:
+        # inputs and targets are [N, M, H, W] where M corresponds to multiple predicted masks
+        assert inputs.dim() == 4 and targets.dim() == 4
+        # flatten spatial dimension while keeping multimask channel dimension
+        inputs = inputs.flatten(2)
+        targets = targets.flatten(2)
+        numerator = 2 * (inputs * targets).sum(-1)
+    else:
+        inputs = inputs.flatten(1)
+        numerator = 2 * (inputs * targets).sum(1)
+    denominator = inputs.sum(-1) + targets.sum(-1)
+    loss = 1 - (numerator + 1) / (denominator + 1)
+    if loss_on_multimask:
+        return loss / num_objects
+    return loss.sum() / num_objects
+
+
+def sigmoid_focal_loss(
+    inputs,
+    targets,
+    num_objects,
+    alpha: float = 0.25,
+    gamma: float = 2,
+    loss_on_multimask=False,
+):
+    """
+    Loss used in RetinaNet for dense detection: https://arxiv.org/abs/1708.02002.
+    Args:
+        inputs: A float tensor of arbitrary shape.
+                The predictions for each example.
+        targets: A float tensor with the same shape as inputs. Stores the binary
+                 classification label for each element in inputs
+                (0 for the negative class and 1 for the positive class).
+        num_objects: Number of objects in the batch
+        alpha: (optional) Weighting factor in range (0,1) to balance
+                positive vs negative examples. Default = -1 (no weighting).
+        gamma: Exponent of the modulating factor (1 - p_t) to
+               balance easy vs hard examples.
+        loss_on_multimask: True if multimask prediction is enabled
+    Returns:
+        focal loss tensor
+    """
+    prob = inputs.sigmoid()
+    ce_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction="none")
+    p_t = prob * targets + (1 - prob) * (1 - targets)
+    loss = ce_loss * ((1 - p_t) ** gamma)
+
+    if alpha >= 0:
+        alpha_t = alpha * targets + (1 - alpha) * (1 - targets)
+        loss = alpha_t * loss
+
+    if loss_on_multimask:
+        # loss is [N, M, H, W] where M corresponds to multiple predicted masks
+        assert loss.dim() == 4
+        return loss.flatten(2).mean(-1) / num_objects  # average over spatial dims
+    return loss.mean(1).sum() / num_objects
+
+
+def iou_loss(inputs, targets, pred_ious, num_objects, loss_on_multimask=False, use_l1_loss=False):
+    """
+    Args:
+        inputs: A float tensor of arbitrary shape.
+                The predictions for each example.
+        targets: A float tensor with the same shape as inputs. Stores the binary
+                 classification label for each element in inputs
+                (0 for the negative class and 1 for the positive class).
+        pred_ious: A float tensor containing the predicted IoUs scores per mask
+        num_objects: Number of objects in the batch
+        loss_on_multimask: True if multimask prediction is enabled
+        use_l1_loss: Whether to use L1 loss is used instead of MSE loss
+    Returns:
+        IoU loss tensor
+    """
+    assert inputs.dim() == 4 and targets.dim() == 4
+    pred_mask = inputs.flatten(2) > 0
+    gt_mask = targets.flatten(2) > 0
+    area_i = torch.sum(pred_mask & gt_mask, dim=-1).float()
+    area_u = torch.sum(pred_mask | gt_mask, dim=-1).float()
+    actual_ious = area_i / torch.clamp(area_u, min=1.0)
+
+    if use_l1_loss:
+        loss = F.l1_loss(pred_ious, actual_ious, reduction="none")
+    else:
+        loss = F.mse_loss(pred_ious, actual_ious, reduction="none")
+    if loss_on_multimask:
+        return loss / num_objects
+    return loss.sum() / num_objects
+
+
+class MultiStepMultiMasksAndIous(nn.Module):
+    def __init__(
+        self,
+        weight_dict,
+        focal_alpha=0.25,
+        focal_gamma=2,
+        supervise_all_iou=False,
+        iou_use_l1_loss=False,
+        pred_obj_scores=False,
+        focal_gamma_obj_score=0.0,
+        focal_alpha_obj_score=-1,
+    ):
+        """
+        This class computes the multi-step multi-mask and IoU losses.
+        Args:
+            weight_dict: dict containing weights for focal, dice, iou losses
+            focal_alpha: alpha for sigmoid focal loss
+            focal_gamma: gamma for sigmoid focal loss
+            supervise_all_iou: if True, back-prop iou losses for all predicted masks
+            iou_use_l1_loss: use L1 loss instead of MSE loss for iou
+            pred_obj_scores: if True, compute loss for object scores
+            focal_gamma_obj_score: gamma for sigmoid focal loss on object scores
+            focal_alpha_obj_score: alpha for sigmoid focal loss on object scores
+        """
+
+        super().__init__()
+        self.weight_dict = weight_dict
+        self.focal_alpha = focal_alpha
+        self.focal_gamma = focal_gamma
+        assert "loss_mask" in self.weight_dict
+        assert "loss_dice" in self.weight_dict
+        assert "loss_iou" in self.weight_dict
+        if "loss_class" not in self.weight_dict:
+            self.weight_dict["loss_class"] = 0.0
+
+        self.focal_alpha_obj_score = focal_alpha_obj_score
+        self.focal_gamma_obj_score = focal_gamma_obj_score
+        self.supervise_all_iou = supervise_all_iou
+        self.iou_use_l1_loss = iou_use_l1_loss
+        self.pred_obj_scores = pred_obj_scores
+
+    def forward(self, outs_batch: List[Dict], targets_batch: torch.Tensor):
+        assert len(outs_batch) == len(targets_batch)
+        num_objects = torch.tensor(
+            (targets_batch.shape[1]), device=targets_batch.device, dtype=torch.float
+        )  # Number of objects is fixed within a batch
+        if is_dist_avail_and_initialized():
+            torch.distributed.all_reduce(num_objects)
+        num_objects = torch.clamp(num_objects / get_world_size(), min=1).item()
+
+        losses = defaultdict(int)
+        for outs, targets in zip(outs_batch, targets_batch):
+            cur_losses = self._forward(outs, targets, num_objects)
+            for k, v in cur_losses.items():
+                losses[k] += v
+
+        return losses
+
+    def _forward(self, outputs: Dict, targets: torch.Tensor, num_objects):
+        """
+        Compute the losses related to the masks: the focal loss and the dice loss.
+        and also the MAE or MSE loss between predicted IoUs and actual IoUs.
+
+        Here "multistep_pred_multimasks_high_res" is a list of multimasks (tensors
+        of shape [N, M, H, W], where M could be 1 or larger, corresponding to
+        one or multiple predicted masks from a click.
+
+        We back-propagate focal, dice losses only on the prediction channel
+        with the lowest focal+dice loss between predicted mask and ground-truth.
+        If `supervise_all_iou` is True, we backpropagate ious losses for all predicted masks.
+        """
+
+        target_masks = targets.unsqueeze(1).float()
+        assert target_masks.dim() == 4  # [N, 1, H, W]
+        src_masks_list = outputs["multistep_pred_multimasks_high_res"]
+        ious_list = outputs["multistep_pred_ious"]
+        object_score_logits_list = outputs["multistep_object_score_logits"]
+
+        assert len(src_masks_list) == len(ious_list)
+        assert len(object_score_logits_list) == len(ious_list)
+
+        # accumulate the loss over prediction steps
+        losses = {"loss_mask": 0, "loss_dice": 0, "loss_iou": 0, "loss_class": 0}
+        for src_masks, ious, object_score_logits in zip(src_masks_list, ious_list, object_score_logits_list):
+            self._update_losses(losses, src_masks, target_masks, ious, num_objects, object_score_logits)
+        losses[CORE_LOSS_KEY] = self.reduce_loss(losses)
+        return losses
+
+    def _update_losses(self, losses, src_masks, target_masks, ious, num_objects, object_score_logits):
+        target_masks = target_masks.expand_as(src_masks)
+        # get focal, dice and iou loss on all output masks in a prediction step
+        loss_multimask = sigmoid_focal_loss(
+            src_masks,
+            target_masks,
+            num_objects,
+            alpha=self.focal_alpha,
+            gamma=self.focal_gamma,
+            loss_on_multimask=True,
+        )
+        loss_multidice = dice_loss(src_masks, target_masks, num_objects, loss_on_multimask=True)
+        if not self.pred_obj_scores:
+            loss_class = torch.tensor(0.0, dtype=loss_multimask.dtype, device=loss_multimask.device)
+            target_obj = torch.ones(
+                loss_multimask.shape[0],
+                1,
+                dtype=loss_multimask.dtype,
+                device=loss_multimask.device,
+            )
+        else:
+            target_obj = torch.any((target_masks[:, 0] > 0).flatten(1), dim=-1)[..., None].float()
+            loss_class = sigmoid_focal_loss(
+                object_score_logits,
+                target_obj,
+                num_objects,
+                alpha=self.focal_alpha_obj_score,
+                gamma=self.focal_gamma_obj_score,
+            )
+
+        loss_multiiou = iou_loss(
+            src_masks,
+            target_masks,
+            ious,
+            num_objects,
+            loss_on_multimask=True,
+            use_l1_loss=self.iou_use_l1_loss,
+        )
+        assert loss_multimask.dim() == 2
+        assert loss_multidice.dim() == 2
+        assert loss_multiiou.dim() == 2
+        if loss_multimask.size(1) > 1:
+            # take the mask indices with the smallest focal + dice loss for back propagation
+            loss_combo = loss_multimask * self.weight_dict["loss_mask"] + loss_multidice * self.weight_dict["loss_dice"]
+            best_loss_inds = torch.argmin(loss_combo, dim=-1)
+            batch_inds = torch.arange(loss_combo.size(0), device=loss_combo.device)
+            loss_mask = loss_multimask[batch_inds, best_loss_inds].unsqueeze(1)
+            loss_dice = loss_multidice[batch_inds, best_loss_inds].unsqueeze(1)
+            # calculate the iou prediction and slot losses only in the index
+            # with the minimum loss for each mask (to be consistent w/ SAM)
+            if self.supervise_all_iou:
+                loss_iou = loss_multiiou.mean(dim=-1).unsqueeze(1)
+            else:
+                loss_iou = loss_multiiou[batch_inds, best_loss_inds].unsqueeze(1)
+        else:
+            loss_mask = loss_multimask
+            loss_dice = loss_multidice
+            loss_iou = loss_multiiou
+
+        # backprop focal, dice and iou loss only if obj present
+        loss_mask = loss_mask * target_obj
+        loss_dice = loss_dice * target_obj
+        loss_iou = loss_iou * target_obj
+
+        # sum over batch dimension (note that the losses are already divided by num_objects)
+        losses["loss_mask"] += loss_mask.sum()
+        losses["loss_dice"] += loss_dice.sum()
+        losses["loss_iou"] += loss_iou.sum()
+        losses["loss_class"] += loss_class
+
+    def reduce_loss(self, losses):
+        reduced_loss = 0.0
+        for loss_key, weight in self.weight_dict.items():
+            if loss_key not in losses:
+                raise ValueError(f"{type(self)} doesn't compute {loss_key}")
+            if weight != 0:
+                reduced_loss += losses[loss_key] * weight
+
+        return reduced_loss

nnInteractive/supervoxel/src/sam2/training/model/__init__.py

@@ -0,0 +1,5 @@
+# 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.