remove transforms from class

updated the labels
added gaussian blur

Author: lezwon

references/detection/transforms.py

@@ -442,25 +442,9 @@ def forward(
 
 
 class SimpleCopyPaste(torch.nn.Module):
-    def __init__(self, jittering_type: str = "LSJ"):
+    def __init__(self):
         super().__init__()
 
-        if jittering_type == "LSJ":
-            scale_range = (0.1, 2.0)
-        elif jittering_type == "SSJ":
-            scale_range = (0.8, 1.25)
-        else:
-            # TODO: add invalid option error
-            raise ValueError("Invalid jittering type")
-
-        self.transforms = Compose(
-            [
-                ScaleJitter(target_size=(1024, 1024), scale_range=scale_range),
-                FixedSizeCrop(size=(1024, 1024), fill=105),
-                RandomHorizontalFlip(0.5),
-            ]
-        )
-
     def combine_masks(self, masks):
         return masks.sum(dim=0).greater(0)
 
@@ -472,22 +456,18 @@ def forward(self, batch: torch.Tensor, target: torch.Tensor) -> Tuple[torch.Tens
         if not batch.is_floating_point():
             raise TypeError(f"Batch dtype should be a float tensor. Got {batch.dtype}.")
 
-        for i, (image, mask) in enumerate(zip(batch, target)):
-            batch[i], target[i] = self.transforms(image, mask)
-
         # create copy of batch and target as the original will be modified
         batch_rolled = batch.roll(1, 0).detach().clone()
         target_rolled = copy.deepcopy(target[-1:] + target[:-1])
 
-        # TODO: select a random subset of objects from one of the images and paste them onto the other image
-
-        # TODO: Smooth out the edges of the pasted objects using a Gaussian filter on the mask
-
         # collect binary paste masks for all images
         paste_masks = []
 
         for source_image, paste_image, source_data, paste_data in zip(batch, batch_rolled, target, target_rolled):
-            paste_alpha_mask = self.combine_masks(paste_data["masks"])
+            number_of_masks = len(paste_data["masks"])
+            random_selection = torch.randint(0, number_of_masks, (number_of_masks,)).unique()
+
+            paste_alpha_mask = self.combine_masks(paste_data["masks"][random_selection])
             paste_masks.append(paste_alpha_mask)
 
             # update original masks
@@ -496,21 +476,24 @@ def forward(self, batch: torch.Tensor, target: torch.Tensor) -> Tuple[torch.Tens
 
             # remove masks where no annotations are present (all values are 0)
             mask_filter = source_data["masks"].sum((2, 1)).not_equal(0)
-            filtered_masks = source_data["masks"][mask_filter]
+            source_data["masks"] = source_data["masks"][mask_filter]
+            source_data["boxes"] = ops.masks_to_boxes(source_data["masks"])
+            source_data["labels"] = source_data["labels"][mask_filter]
+            source_data["area"] = source_data["area"][mask_filter]
+            source_data["iscrowd"] = source_data["iscrowd"][mask_filter]
 
-            # update bboxes based on new masks
-            source_data["boxes"] = ops.masks_to_boxes(filtered_masks)
             # TODO: update area
 
             # concatenate paste data with original data
-            source_data["masks"] = torch.cat((source_data["masks"], paste_data["masks"]))
-            source_data["boxes"] = torch.cat((source_data["boxes"], paste_data["boxes"]))
-            source_data["labels"] = torch.cat((source_data["labels"], paste_data["labels"]))
-            source_data["area"] = torch.cat((source_data["area"], paste_data["area"]))
-            source_data["iscrowd"] = torch.cat((source_data["iscrowd"], paste_data["iscrowd"]))
+            source_data["masks"] = torch.cat((source_data["masks"], paste_data["masks"][random_selection]))
+            source_data["boxes"] = torch.cat((source_data["boxes"], paste_data["boxes"][random_selection]))
+            source_data["labels"] = torch.cat((source_data["labels"], paste_data["labels"][random_selection]))
+            source_data["area"] = torch.cat((source_data["area"], paste_data["area"][random_selection]))
+            source_data["iscrowd"] = torch.cat((source_data["iscrowd"], paste_data["iscrowd"][random_selection]))
 
         # update the original images with paste images
-        paste_masks = torch.stack(paste_masks)
+        paste_masks = torch.stack(paste_masks).to(torch.uint8)
+        paste_masks = T.GaussianBlur((5, 5), sigma=2)(paste_masks)  # Adds Gaussian Filter
         batch.mul_(torch.unsqueeze(torch.logical_not(paste_masks), 1))
 
         paste_images = batch_rolled * torch.unsqueeze(paste_masks, 1)