Is it possible to view/save images in the train/test pipelines?

[FranzEricSchneider]

I have been trying to experiment with augmentations in the dataset config for a new dataset, and I've realized it would be extremely helpful to be able to save out a few of the images in the training process. The training pipeline applies a series of augmentations (very handy), but it would be great to be able to save some (maybe every 100th, or 1 per iteration, or 1 per epoch, something low frequency like that). The same thing is true for the test/validation pipeline, although the augmentations are simpler.

For example, here is the training pipeline as I'm experimenting with it at the moment. For some of the more complicated ones like Resize, RandomCutOut, RandomRotate, and CLAHE I think I know what the augmentation is doing, but I would love to confirm that by checking out some of the augmented training images and the corresponding augmented labels.

train_pipeline = [
    dict(type="LoadImageFromFile"),
    dict(type="LoadAnnotations"),
    dict(type="Resize", img_scale=(848, 480), ratio_range=(0.5, 2.0)),
    dict(type="ResizeToMultiple", size_divisor=16),
    dict(type="RandomCrop", crop_size=crop_size, cat_max_ratio=0.75),
    dict(type="RandomFlip", prob=0.5),
    dict(type="RandomRotate", prob=0.75, degree=180),
    dict(type="RandomCutOut",
         prob=0.25,
         n_holes=(1, 8),
         cutout_ratio=[(0.05, 0.05), (0.05, 0.1), (0.05, 0.15),
                       (0.1, 0.05), (0.1, 0.1), (0.1, 0.15),
                       (0.15, 0.05), (0.15, 0.1), (0.15, 0.15)]),
    dict(type="CLAHE"),
    dict(type="PhotoMetricDistortion"),
    dict(type="Normalize", **img_norm_cfg),
    dict(type="Pad", size=crop_size, pad_val=0, seg_pad_val=255),
    dict(type="DefaultFormatBundle"),
    dict(type="Collect", keys=["img", "gt_semantic_seg"]),
]

Is there any option to automatically save some of the training images and labels for later viewing? If not does that seem like a simple feature I could take a look at?

[xiexinch]

Hi @FranzEricSchneider
Sorry for the late reply.
Your requirement is similar to the visualization feature in MMGeneration.
https://github.com/open-mmlab/mmgeneration/blob/master/mmgen/core/hooks/visualize_training_samples.py#L12

I'll try to test this feature in MMSegmentation, if it works, I'll update to you.

[xiexinch]

Hi @FranzEricSchneider
We just need to make a little modification to the visualization hook implemented in MMGeneration.

    # Modify after_train_iter to before_train_iter
    @master_only
    def before_train_iter(self, runner):
        """The behavior before each train iteration.

        Args:
            runner (object): The runner.
        """
        if not self.every_n_iters(runner, self.interval):
            return
        # get images from data container
        imgs = runner.data_batch['img'].data
        img_metas = runner.data_batch['img_metas'].data[0]

        filename = self.filename_tmpl.format(runner.iter + 1)

        # img_list = [x for k, x in results.items() if k in self.res_name_list]
        img_list = [img for img in imgs]
        for img in img_list:
            print(img.shape)
        img_cat = torch.cat(img_list, dim=3).detach()
        if self.rerange:
            img_cat = ((img_cat + 1) / 2)
        if self.bgr2rgb:
            img_cat = img_cat[:, [2, 1, 0], ...]
        img_cat = img_cat.clamp_(0, 1)

        if not hasattr(self, '_out_dir'):
            self._out_dir = osp.join(runner.work_dir, self.output_dir)
        mmcv.mkdir_or_exist(self._out_dir)
        save_image(
            img_cat,
            osp.join(self._out_dir, filename),
            nrow=self.nrow,
            padding=self.padding)
        filename = filename.replace('.png', '_meta_infos.json')
        mmcv.dump(img_metas, osp.join(self._out_dir, filename))

Add custom hook config to your config file.

custom_hooks = [
    dict(
        type='MMSegVisualizationHook',
        output_dir='training_samples',
        bgr2rgb=True,
        interval=50)
]

[FranzEricSchneider]

Sorry I don't totally follow, I don't have mmgeneration installed at the moment. Do I need to start by copying some code from the mmgeneration source into a checkout of mmsegmentation? Such as the before_train_iter function perhaps? If so, where would it go in mmsegmentation?

Thanks!

[cheul0518]

Hi @xiexinch , can you elaborate how to add the custom hook that you wrote? I read customize-hooks but haven't figured out how to implement your code into my config file. This is what happened when adding your code into my cfg.

>>> 'MMSegVisualizationHook is not in the hook registry'

Please enlighten me.

[xiexinch]

Hi @cheul0518,
Just like MMSegWandbHook , write your own hook class then register it by HOOK, thus the runner might find your custom hook.

mmsegmentation/mmseg/core/hook/wandblogger_hook.py

Line 14 in 7a0f45e

@HOOKS.register_module()

[xiexinch]

Might try to use MMSegWandbHook, you can find more info in thePR #1603 that added this feature.