Update dataset_mapper.py

odise/data/dataset_mapper.py

@@ -26,6 +26,252 @@
 from panopticapi.utils import rgb2id
 
 
+class ADEPanopticDatasetMapper:
+    """
+    A callable which takes a dataset dict in Detectron2 Dataset format,
+    and map it into a format used by MaskFormer.
+
+    This dataset mapper applies the same transformation as DETR for COCO panoptic segmentation.
+
+    The callable currently does the following:
+
+    1. Read the image from "file_name"
+    2. Applies geometric transforms to the image and annotation
+    3. Find and applies suitable cropping to the image and annotation
+    4. Prepare image and annotation to Tensors
+    """
+
+    def __init__(
+        self,
+        is_train: bool = True,
+        *,
+        augmentations: List[Union[T.Augmentation, T.Transform]],
+        image_format: str,
+        segmentation_format: str = "L",
+        caption_key: str = "coco_captions",
+    ):
+        """
+        NOTE: this interface is experimental.
+        Args:
+            is_train: for training or inference
+            augmentations: a list of augmentations or deterministic transforms to apply
+            crop_gen: crop augmentation
+            tfm_gens: data augmentation
+            image_format: an image format supported by :func:`detection_utils.read_image`.
+        """
+        self.augmentations = T.AugmentationList(augmentations)
+        logging.getLogger(__name__).info(
+            f"[{self.__class__.__name__}] Full TransformGens used in training: {self.augmentations}"
+        )
+
+        self.img_format = image_format
+        self.seg_format = segmentation_format
+        self.cap_key = caption_key
+        self.is_train = is_train
+
+
+    def __call__(self, dataset_dict):
+
+        dataset_dict = copy.deepcopy(dataset_dict)  # it will be modified by code below
+        image = utils.read_image(dataset_dict["file_name"], format=self.img_format)
+        utils.check_image_size(dataset_dict, image)
+
+
+        # USER: Remove if you don't do semantic/panoptic segmentation.
+        if "sem_seg_file_name" in dataset_dict:
+            sem_seg_gt = utils.read_image(
+                dataset_dict.pop("sem_seg_file_name"), format=self.seg_format
+            )
+            if self.seg_format == "L":
+                sem_seg_gt = sem_seg_gt.squeeze(2)
+        else:
+            sem_seg_gt = None
+
+
+        # image, transforms = T.apply_augmentations(self.augmentations, image)
+        aug_input = T.AugInput(image, sem_seg=sem_seg_gt)
+        transforms = self.augmentations(aug_input)
+        image, sem_seg_gt = aug_input.image, aug_input.sem_seg
+
+        image_shape = image.shape[:2]  # h, w
+
+        dataset_dict["image"] = torch.as_tensor(np.ascontiguousarray(image.transpose(2, 0, 1)))
+
+        if sem_seg_gt is not None:
+            dataset_dict["sem_seg"] = torch.as_tensor(sem_seg_gt.astype("long"))
+
+        # if not self.is_train:
+        #     # USER: Modify this if you want to keep them for some reason.
+        #     dataset_dict.pop("annotations", None)
+        #     return dataset_dict
+
+        pan_seg_gt = utils.read_image(dataset_dict.pop("pan_seg_file_name"), "RGB")
+        segments_info = dataset_dict["segments_info"]
+
+        # apply the same transformation to panoptic segmentation
+        pan_seg_gt = transforms.apply_segmentation(pan_seg_gt)
+
+        pan_seg_gt = rgb2id(pan_seg_gt)
+        dataset_dict["pan_seg_gt"] = torch.from_numpy(np.ascontiguousarray(pan_seg_gt))
+
+        instances = Instances(image_shape)
+        classes = []
+        masks = []
+        for segment_info in segments_info:
+            class_id = segment_info["category_id"]
+            if not segment_info["iscrowd"]:
+                classes.append(class_id)
+                masks.append(pan_seg_gt == segment_info["id"])
+
+        classes = np.array(classes)
+        instances.gt_classes = torch.tensor(classes, dtype=torch.int64)
+        if len(masks) == 0:
+            # Some image does not have annotation (all ignored)
+            instances.gt_masks = torch.zeros((0, pan_seg_gt.shape[-2], pan_seg_gt.shape[-1]))
+            instances.gt_boxes = Boxes(torch.zeros((0, 4)))
+        else:
+            masks = BitMasks(
+                torch.stack([torch.from_numpy(np.ascontiguousarray(x.copy())) for x in masks])
+            )
+            instances.gt_masks = masks.tensor
+            instances.gt_boxes = masks.get_bounding_boxes()
+
+        if self.cap_key in dataset_dict:
+            dataset_dict["captions"] = dataset_dict.pop(self.cap_key)
+
+        dataset_dict["instances"] = instances
+
+
+
+
+        return dataset_dict 
+
+
+class LVISPanopticDatasetMapper:
+    """
+    A callable which takes a dataset dict in Detectron2 Dataset format,
+    and map it into a format used by MaskFormer.
+
+    This dataset mapper applies the same transformation as DETR for COCO panoptic segmentation.
+
+    The callable currently does the following:
+
+    1. Read the image from "file_name"
+    2. Applies geometric transforms to the image and annotation
+    3. Find and applies suitable cropping to the image and annotation
+    4. Prepare image and annotation to Tensors
+    """
+
+    def __init__(
+        self,
+        is_train: bool = True,
+        *,
+        augmentations: List[Union[T.Augmentation, T.Transform]],
+        image_format: str,
+        segmentation_format: str = "L",
+        caption_key: str = "coco_captions",
+    ):
+        """
+        NOTE: this interface is experimental.
+        Args:
+            is_train: for training or inference
+            augmentations: a list of augmentations or deterministic transforms to apply
+            crop_gen: crop augmentation
+            tfm_gens: data augmentation
+            image_format: an image format supported by :func:`detection_utils.read_image`.
+        """
+        self.augmentations = T.AugmentationList(augmentations)
+        logging.getLogger(__name__).info(
+            f"[{self.__class__.__name__}] Full TransformGens used in training: {self.augmentations}"
+        )
+
+        self.img_format = image_format
+        self.seg_format = segmentation_format
+        self.cap_key = caption_key
+        self.is_train = is_train
+
+
+    def __call__(self, dataset_dict):
+        """
+        Args:
+            dataset_dict (dict): Metadata of one image, in Detectron2 Dataset format.
+
+        Returns:
+            dict: a format that builtin models in detectron2 accept
+        """
+        dataset_dict = copy.deepcopy(dataset_dict)  # it will be modified by code below
+
+        image = utils.read_image(dataset_dict["file_name"], format=self.img_format)
+        utils.check_image_size(dataset_dict, image)
+
+        # USER: Remove if you don't do semantic/panoptic segmentation.
+        if "sem_seg_file_name" in dataset_dict:
+            sem_seg_gt = utils.read_image(
+                dataset_dict.pop("sem_seg_file_name"), format=self.seg_format
+            )
+            if self.seg_format == "L":
+                sem_seg_gt = sem_seg_gt.squeeze(2)
+        else:
+            sem_seg_gt = None
+
+        # image, transforms = T.apply_augmentations(self.augmentations, image)
+        aug_input = T.AugInput(image, sem_seg=sem_seg_gt)
+        transforms = self.augmentations(aug_input)
+        image, sem_seg_gt = aug_input.image, aug_input.sem_seg
+
+        image_shape = image.shape[:2]  # h, w
+
+        # Pytorch's dataloader is efficient on torch.Tensor due to shared-memory,
+        # but not efficient on large generic data structures due to the use of pickle & mp.Queue.
+        # Therefore it's important to use torch.Tensor.
+        dataset_dict["image"] = torch.as_tensor(np.ascontiguousarray(image.transpose(2, 0, 1)))
+        if sem_seg_gt is not None:
+            dataset_dict["sem_seg"] = torch.as_tensor(sem_seg_gt.astype("long"))
+
+        # if not self.is_train:
+        #     # USER: Modify this if you want to keep them for some reason.
+        #     dataset_dict.pop("annotations", None)
+        #     return dataset_dict
+
+        pan_seg_gt = utils.read_image(dataset_dict.pop("pan_seg_file_name"), "RGB")
+        segments_info = dataset_dict["segments_info"]
+
+        # apply the same transformation to panoptic segmentation
+        pan_seg_gt = transforms.apply_segmentation(pan_seg_gt)
+
+        pan_seg_gt = rgb2id(pan_seg_gt)
+        dataset_dict["pan_seg_gt"] = torch.from_numpy(np.ascontiguousarray(pan_seg_gt))
+
+        instances = Instances(image_shape)
+        classes = []
+        masks = []
+        for segment_info in segments_info:
+            class_id = segment_info["category_id"]
+            if not segment_info["iscrowd"]:
+                classes.append(class_id)
+                masks.append(pan_seg_gt == segment_info["id"])
+
+        classes = np.array(classes)
+        instances.gt_classes = torch.tensor(classes, dtype=torch.int64)
+        if len(masks) == 0:
+            # Some image does not have annotation (all ignored)
+            instances.gt_masks = torch.zeros((0, pan_seg_gt.shape[-2], pan_seg_gt.shape[-1]))
+            instances.gt_boxes = Boxes(torch.zeros((0, 4)))
+        else:
+            masks = BitMasks(
+                torch.stack([torch.from_numpy(np.ascontiguousarray(x.copy())) for x in masks])
+            )
+            instances.gt_masks = masks.tensor
+            instances.gt_boxes = masks.get_bounding_boxes()
+
+        if self.cap_key in dataset_dict:
+            dataset_dict["captions"] = dataset_dict.pop(self.cap_key)
+
+        dataset_dict["instances"] = instances
+
+        return dataset_dict
+
+
 class COCOPanopticDatasetMapper:
     """
     A callable which takes a dataset dict in Detectron2 Dataset format,
@@ -69,6 +315,7 @@ def __init__(
         self.cap_key = caption_key
         self.is_train = is_train
 
+
     def __call__(self, dataset_dict):
         """
         Args:
@@ -77,20 +324,27 @@ def __call__(self, dataset_dict):
         Returns:
             dict: a format that builtin models in detectron2 accept
         """
+
+
         dataset_dict = copy.deepcopy(dataset_dict)  # it will be modified by code below
         image = utils.read_image(dataset_dict["file_name"], format=self.img_format)
         utils.check_image_size(dataset_dict, image)
 
+
         # USER: Remove if you don't do semantic/panoptic segmentation.
         if "sem_seg_file_name" in dataset_dict:
+
             sem_seg_gt = utils.read_image(
                 dataset_dict.pop("sem_seg_file_name"), format=self.seg_format
             )
+
             if self.seg_format == "L":
                 sem_seg_gt = sem_seg_gt.squeeze(2)
+
         else:
             sem_seg_gt = None
 
+
         # image, transforms = T.apply_augmentations(self.augmentations, image)
         aug_input = T.AugInput(image, sem_seg=sem_seg_gt)
         transforms = self.augmentations(aug_input)
@@ -146,4 +400,5 @@ def __call__(self, dataset_dict):
 
         dataset_dict["instances"] = instances
 
+
         return dataset_dict