README.md 3.06 KB
Newer Older
Alexander Nicholas Leye's avatar
Alexander Nicholas Leye committed
1 2
# Scene-Copy-Paste
Unofficial implementation of the copy-paste augmentation from [Simple Copy-Paste is a Strong Data Augmentation Method for Instance Segmentation](https://arxiv.org/abs/2012.07177v1), adapted to paste object instances from an image onto an empty scene. Inspiration from [Cut, Paste and Learn: Surprisingly Easy Synthesis for Instance Detection](https://arxiv.org/abs/1708.01642).
conradry's avatar
conradry committed
3

conradry's avatar
v1  
conradry committed
4
The augmentation function is built to integrate easily with albumentations. An example for creating a compatible torchvision dataset is given for COCO. Core functionality for image, masks, and bounding boxes is finished; keypoints are not yet supported. In general, you can use the CopyPaste augmentation just as you would any other albumentations augmentation function. There are a few usage limitations of note.
conradry's avatar
conradry committed
5

conradry's avatar
v1  
conradry committed
6 7 8
<figure>
  <img src="./example.png"></img>
</figure>
conradry's avatar
conradry committed
9

conradry's avatar
v1  
conradry committed
10
## Usage Notes
conradry's avatar
conradry committed
11

conradry's avatar
v1  
conradry committed
12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31
1. BboxParams cannot have label_fields. To attach class labels to a bounding box, directly append it to the bounding box coordinates. (I.e. (x1, y1, x2, y2, class_id)).
2. Bounding boxes passed to the CopyPaste augmentation must also include the index of the corresponding mask in the 'masks' list. (I.e. the bounding box looks like (x1, y1, x2, y2, class_id, mask_index)). An example is given for COCO.
3. The CopyPaste augmentation expects 6 keyword arguments instead of three:

```python
output = transforms(image=image, masks=masks, bboxes=bboxes)
--->instead
output = transforms(
    image=image, masks=masks, bboxes=bboxes,
    paste_image=paste_image, paste_masks=paste_masks, paste_bboxes=paste_bboxes
  )
```

4. After pasting objects, the original bounding boxes may be occluded. To make things easier,
the original bounding boxes are just extracted from the updated masks.

## Integration with Torchvision datasets

The idea is to have a standard torchvision dataset that's decorated to add seamlessly integrate the
copy-paste functionality.
conradry's avatar
conradry committed
32 33

The dataset class looks like:
conradry's avatar
conradry committed
34 35

```python
conradry's avatar
conradry committed
36 37 38 39 40
from copy_paste import copy_paste_class
from torch.utils.data import Dataset

@copy_paste_class
class SomeVisionDataset(Dataset):
41 42
    def __init__(self, *args):
        super(SomeVisionDataset, self).__init__(*args)
conradry's avatar
conradry committed
43

conradry's avatar
conradry committed
44 45
    def __len__(self):
        return length
conradry's avatar
conradry committed
46

conradry's avatar
conradry committed
47 48 49 50 51 52 53 54 55
    def load_example(self, idx):
        image_data_dict = load_some_data(idx)
        transformed_data_dict = self.transforms(**image_data_dict)
        return transformed_data_dict

```
The only difference from a regular torchvision dataset is the decorator and the ```load_example``` method
instead of ```__getitem__```.

conradry's avatar
v1  
conradry committed
56
To compose transforms with copy-paste augmentation (bbox params are NOT optional):
conradry's avatar
conradry committed
57

58

conradry's avatar
conradry committed
59
```python
conradry's avatar
conradry committed
60 61 62 63 64
import albumentations as A
from albumentations.pytorch.transforms import ToTensorV2
from copy_paste import CopyPaste

transform = A.Compose([
65 66 67 68 69
      A.RandomScale(scale_limit=(-0.9, 1), p=1), #LargeScaleJitter from scale of 0.1 to 2
      A.PadIfNeeded(256, 256, border_mode=0), #constant 0 border
      A.RandomCrop(256, 256),
      A.HorizontalFlip(p=0.5),
      CopyPaste(blend=True, sigma=1, pct_objects_paste=0.5, p=1)
conradry's avatar
conradry committed
70
    ], bbox_params=A.BboxParams(format="coco")
conradry's avatar
conradry committed
71 72
)
```