How to use detectron2's augmentation with datasets loaded using register_coco_instances

Question

I've trained a detectron2 model on custom data I labeled and exported in the coco format, but I now want to apply augmentation and train using the augmented data. How can I do that if I'm not using a custom DataLoader, but the register_coco_instances function.

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
predictor = DefaultPredictor(cfg)
outputs = predictor(im)

train_annotations_path = "./data/cvat-corn-train-coco-1.0/annotations/instances_default.json"
train_images_path = "./data/cvat-corn-train-coco-1.0/images"
validation_annotations_path = "./data/cvat-corn-validation-coco-1.0/annotations/instances_default.json"
validation_images_path = "./data/cvat-corn-validation-coco-1.0/images"

register_coco_instances(
    "train-corn",
    {},
    train_annotations_path,
    train_images_path
)
register_coco_instances(
    "validation-corn",
    {},
    validation_annotations_path,
    validation_images_path
)
metadata_train = MetadataCatalog.get("train-corn")
dataset_dicts = DatasetCatalog.get("train-corn")

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = ("train-corn",)
cfg.DATASETS.TEST = ("validation-corn",)
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")  # Let training initialize from model zoo
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025
cfg.SOLVER.MAX_ITER = 10000
cfg.SOLVER.STEPS = []
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128 
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 4
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg) 
trainer.resume_or_load(resume=False)
trainer.train()

I saw in the documentation you can load a dataset and apply augmentation like this:

dataloader = build_detection_train_loader(cfg,
   mapper=DatasetMapper(cfg, is_train=True, augmentations=[
      T.Resize((800, 800))
   ]))

But I'm not using a custom dataloader, what is the best approach to do this?

score 5 · Accepted Answer · answered Apr 08 '22 at 15:18

From my experience, how you register your datasets (i.e., tell Detectron2 how to obtain a dataset named "my_dataset") has no bearing on what dataloader to use during training (i.e., how to load information from a registered dataset and process it into a format needed by the model).

So, you can register your dataset however you want - either by using the register_coco_instances function or by using the dataset APIs (DatasetCatalog, MetadataCatalog) directly; it doesn't matter. What matters is that you want to apply some transformation(s) during the data loading part.

Basically, you want to customise the data loading part which can only be achieved by using a custom dataloader (unless you perform offline augmentation which is likely not what you want).

Now, you don't need to define and use a custom dataloader directly in your top-level code. You can just create your own trainer deriving from DefaultTrainer, and override its build_train_loader method. This is as simple as the following.

class MyTrainer(DefaultTrainer):

    @classmethod
    def build_train_loader(cls, cfg):
        mapper = DatasetMapper(cfg, is_train=True, augmentations=[T.Resize((800, 800))])
        return build_detection_train_loader(cfg, mapper=mapper)

In your top-level code then, the only change required would be to use MyTrainer instead of DefaultTrainer.

trainer = MyTrainer(cfg) 
trainer.resume_or_load(resume=False)
trainer.train()

Thx finally a working approach! I first tried https://github.com/facebookresearch/detectron2/issues/1578 but when I started training, the "Augmentations used in training: ... " log line did not reflected my augmentations list. With your approach, I finally see my augmentations list as expected. — Greg7000, Apr 14 '22 at 13:47
@zepman If I had a custom augmentation function that takes an image as in input (it's not a method derived from `T.`), how would that work? Can I just place my function as: `augmentations=[some_func]`? — Omar AlSuwaidi, Jul 24 '22 at 22:13
@OmarAlSuwaidi No, unfortunately. The `augmentations` argument of the `DatasetMapper` constructor expects a list of `T.Augmentation` / `T.Transform`. So, you will need to wrap your custom function in classes derived from those. For example, [this code](https://gist.github.com/mallyagirish/1885443642b3b7bf438a20f216df0dc3#file-custom_aug_for_detectron2-py) shows a random Gaussian Blur augmentation being defined via a `GaussianBlurTransform` that derives from `T.Transform`. You can see that your custom function would sit in the overloaded `apply_image` method that takes an image as input. — zepman, Jul 25 '22 at 12:26
@zepman Using this approach gives me a CUDA out-of-memory error. I think detectron2 is storing the augmented data as tensor's in GPU's memory (makes sense). I ended up decreasing the image size to fit it in my 16GB Tesla T4. — Muneeb Ahmad Khurram, Dec 07 '22 at 12:57

How to use detectron2's augmentation with datasets loaded using register_coco_instances

1 Answers1

Linked