[Enhancement] Fix the build_segmentor in colab demo (#1551)

* [Enhancement] Delete DS_Store file * [Enhancement] Fix the build_segmentor

[Enhancement] Fix the build_segmentor in colab demo (#1551)
* [Enhancement] Delete DS_Store file * [Enhancement] Fix the build_segmentor
9b3338fa · Miao Zheng · GitHub · 46309187 · 9b3338fa
Commit 9b3338fa authored 2 years ago by Miao Zheng Committed by GitHub 2 years ago
--- a/demo/MMSegmentation_Tutorial.ipynb
+++ b/demo/MMSegmentation_Tutorial.ipynb
@@ -561,8 +561,7 @@
    "datasets = [build_dataset(cfg.data.train)]\n",
    "\n",
    "# Build the detector\n",
-    "model = build_segmentor(\n",
-    "    cfg.model, train_cfg=cfg.get('train_cfg'), test_cfg=cfg.get('test_cfg'))\n",
+    "model = build_segmentor(cfg.model)\n",
    "# Add an attribute for visualization convenience\n",
    "model.CLASSES = datasets[0].CLASSES\n",
    "\n",

 %% Cell type:markdown id: tags:

 <a href="https://colab.research.google.com/github/open-mmlab/mmsegmentation/blob/master/demo/MMSegmentation_Tutorial.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>

 %% Cell type:markdown id: tags:

 # MMSegmentation Tutorial
 Welcome to MMSegmentation!

 In this tutorial, we demo
 * How to do inference with MMSeg trained weight
 * How to train on your own dataset and visualize the results.

 %% Cell type:markdown id: tags:

 ## Install MMSegmentation
 This step may take several minutes.

 We use PyTorch 1.10 and CUDA 11.1 for this tutorial. You may install other versions by change the version number in pip install command.

 %% Cell type:code id: tags:

 ``` python
 # Check nvcc version
 !nvcc -V
 # Check GCC version
 !gcc --version
 ```

 %% Cell type:code id: tags:

 ``` python
 # Install PyTorch
 !conda install pytorch=1.10.0 torchvision cudatoolkit=11.1 -c pytorch
 # Install MMCV
 !pip install mmcv-full -f https://download.openmmlab.com/mmcv/dist/cu111/torch1.10/index.html
 ```

 %% Cell type:code id: tags:

 ``` python
 !rm -rf mmsegmentation
 !git clone https://github.com/open-mmlab/mmsegmentation.git
 %cd mmsegmentation
 !pip install -e .
 ```

 %% Cell type:code id: tags:

 ``` python
 # Check Pytorch installation
 import torch, torchvision
 print(torch.__version__, torch.cuda.is_available())

 # Check MMSegmentation installation
 import mmseg
 print(mmseg.__version__)
 ```

 %% Cell type:markdown id: tags:

 ## Run Inference with MMSeg trained weight

 %% Cell type:code id: tags:

 ``` python
 !mkdir checkpoints
 !wget https://download.openmmlab.com/mmsegmentation/v0.5/pspnet/pspnet_r50-d8_512x1024_40k_cityscapes/pspnet_r50-d8_512x1024_40k_cityscapes_20200605_003338-2966598c.pth -P checkpoints
 ```

 %% Cell type:code id: tags:

 ``` python
 from mmseg.apis import inference_segmentor, init_segmentor, show_result_pyplot
 from mmseg.core.evaluation import get_palette
 ```

 %% Cell type:code id: tags:

 ``` python
 config_file = 'configs/pspnet/pspnet_r50-d8_512x1024_40k_cityscapes.py'
 checkpoint_file = 'checkpoints/pspnet_r50-d8_512x1024_40k_cityscapes_20200605_003338-2966598c.pth'
 ```

 %% Cell type:code id: tags:

 ``` python
 # build the model from a config file and a checkpoint file
 model = init_segmentor(config_file, checkpoint_file, device='cuda:0')
 ```

 %% Cell type:code id: tags:

 ``` python
 # test a single image
 img = 'demo/demo.png'
 result = inference_segmentor(model, img)
 ```

 %% Cell type:code id: tags:

 ``` python
 # show the results
 show_result_pyplot(model, img, result, get_palette('cityscapes'))
 ```

 %% Cell type:markdown id: tags:

 ## Train a semantic segmentation model on a new dataset

 To train on a customized dataset, the following steps are necessary.
 1. Add a new dataset class.
 2. Create a config file accordingly.
 3. Perform training and evaluation.

 %% Cell type:markdown id: tags:

 ### Add a new dataset

 Datasets in MMSegmentation require image and semantic segmentation maps to be placed in folders with the same prefix. To support a new dataset, we may need to modify the original file structure.

 In this tutorial, we give an example of converting the dataset. You may refer to [docs](https://github.com/open-mmlab/mmsegmentation/blob/master/docs/en/tutorials/customize_datasets.md#customize-datasets-by-reorganizing-data) for details about dataset reorganization.

 We use [Stanford Background Dataset](http://dags.stanford.edu/projects/scenedataset.html) as an example. The dataset contains 715 images chosen from existing public datasets [LabelMe](http://labelme.csail.mit.edu), [MSRC](http://research.microsoft.com/en-us/projects/objectclassrecognition), [PASCAL VOC](http://pascallin.ecs.soton.ac.uk/challenges/VOC) and [Geometric Context](http://www.cs.illinois.edu/homes/dhoiem/). Images from these datasets are mainly outdoor scenes, each containing approximately 320-by-240 pixels.
 In this tutorial, we use the region annotations as labels. There are 8 classes in total, i.e. sky, tree, road, grass, water, building, mountain, and foreground object.

 %% Cell type:code id: tags:

 ``` python
 # download and unzip
 !wget http://dags.stanford.edu/data/iccv09Data.tar.gz -O stanford_background.tar.gz
 !tar xf stanford_background.tar.gz
 ```

 %% Cell type:code id: tags:

 ``` python
 # Let's take a look at the dataset
 import mmcv
 import matplotlib.pyplot as plt

 img = mmcv.imread('iccv09Data/images/6000124.jpg')
 plt.figure(figsize=(8, 6))
 plt.imshow(mmcv.bgr2rgb(img))
 plt.show()
 ```

 %% Cell type:markdown id: tags:

 We need to convert the annotation into semantic map format as an image.

 %% Cell type:code id: tags:

 ``` python
 import os.path as osp
 import numpy as np
 from PIL import Image
 # convert dataset annotation to semantic segmentation map
 data_root = 'iccv09Data'
 img_dir = 'images'
 ann_dir = 'labels'
 # define class and plaette for better visualization
 classes = ('sky', 'tree', 'road', 'grass', 'water', 'bldg', 'mntn', 'fg obj')
 palette = [[128, 128, 128], [129, 127, 38], [120, 69, 125], [53, 125, 34],
           [0, 11, 123], [118, 20, 12], [122, 81, 25], [241, 134, 51]]
 for file in mmcv.scandir(osp.join(data_root, ann_dir), suffix='.regions.txt'):
  seg_map = np.loadtxt(osp.join(data_root, ann_dir, file)).astype(np.uint8)
  seg_img = Image.fromarray(seg_map).convert('P')
  seg_img.putpalette(np.array(palette, dtype=np.uint8))
  seg_img.save(osp.join(data_root, ann_dir, file.replace('.regions.txt',
                                                         '.png')))
 ```

 %% Cell type:code id: tags:

 ``` python
 # Let's take a look at the segmentation map we got
 import matplotlib.patches as mpatches
 img = Image.open('iccv09Data/labels/6000124.png')
 plt.figure(figsize=(8, 6))
 im = plt.imshow(np.array(img.convert('RGB')))

 # create a patch (proxy artist) for every color
 patches = [mpatches.Patch(color=np.array(palette[i])/255.,
                          label=classes[i]) for i in range(8)]
 # put those patched as legend-handles into the legend
 plt.legend(handles=patches, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.,
           fontsize='large')

 plt.show()
 ```

 %% Cell type:code id: tags:

 ``` python
 # split train/val set randomly
 split_dir = 'splits'
 mmcv.mkdir_or_exist(osp.join(data_root, split_dir))
 filename_list = [osp.splitext(filename)[0] for filename in mmcv.scandir(
    osp.join(data_root, ann_dir), suffix='.png')]
 with open(osp.join(data_root, split_dir, 'train.txt'), 'w') as f:
  # select first 4/5 as train set
  train_length = int(len(filename_list)*4/5)
  f.writelines(line + '\n' for line in filename_list[:train_length])
 with open(osp.join(data_root, split_dir, 'val.txt'), 'w') as f:
  # select last 1/5 as train set
  f.writelines(line + '\n' for line in filename_list[train_length:])
 ```

 %% Cell type:markdown id: tags:

 After downloading the data, we need to implement `load_annotations` function in the new dataset class `StanfordBackgroundDataset`.

 %% Cell type:code id: tags:

 ``` python
 from mmseg.datasets.builder import DATASETS
 from mmseg.datasets.custom import CustomDataset

 @DATASETS.register_module()
 class StanfordBackgroundDataset(CustomDataset):
  CLASSES = classes
  PALETTE = palette
  def __init__(self, split, **kwargs):
    super().__init__(img_suffix='.jpg', seg_map_suffix='.png',
                     split=split, **kwargs)
    assert osp.exists(self.img_dir) and self.split is not None


 ```

 %% Cell type:markdown id: tags:

 ### Create a config file
 In the next step, we need to modify the config for the training. To accelerate the process, we finetune the model from trained weights.

 %% Cell type:code id: tags:

 ``` python
 from mmcv import Config
 cfg = Config.fromfile('configs/pspnet/pspnet_r50-d8_512x1024_40k_cityscapes.py')
 ```

 %% Cell type:markdown id: tags:

 Since the given config is used to train PSPNet on the cityscapes dataset, we need to modify it accordingly for our new dataset.

 %% Cell type:code id: tags:

 ``` python
 from mmseg.apis import set_random_seed

 # Since we use only one GPU, BN is used instead of SyncBN
 cfg.norm_cfg = dict(type='BN', requires_grad=True)
 cfg.model.backbone.norm_cfg = cfg.norm_cfg
 cfg.model.decode_head.norm_cfg = cfg.norm_cfg
 cfg.model.auxiliary_head.norm_cfg = cfg.norm_cfg
 # modify num classes of the model in decode/auxiliary head
 cfg.model.decode_head.num_classes = 8
 cfg.model.auxiliary_head.num_classes = 8

 # Modify dataset type and path
 cfg.dataset_type = 'StanfordBackgroundDataset'
 cfg.data_root = data_root

 cfg.data.samples_per_gpu = 8
 cfg.data.workers_per_gpu=8

 cfg.img_norm_cfg = dict(
    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
 cfg.crop_size = (256, 256)
 cfg.train_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='LoadAnnotations'),
    dict(type='Resize', img_scale=(320, 240), ratio_range=(0.5, 2.0)),
    dict(type='RandomCrop', crop_size=cfg.crop_size, cat_max_ratio=0.75),
    dict(type='RandomFlip', flip_ratio=0.5),
    dict(type='PhotoMetricDistortion'),
    dict(type='Normalize', **cfg.img_norm_cfg),
    dict(type='Pad', size=cfg.crop_size, pad_val=0, seg_pad_val=255),
    dict(type='DefaultFormatBundle'),
    dict(type='Collect', keys=['img', 'gt_semantic_seg']),
 ]

 cfg.test_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(
        type='MultiScaleFlipAug',
        img_scale=(320, 240),
        # img_ratios=[0.5, 0.75, 1.0, 1.25, 1.5, 1.75],
        flip=False,
        transforms=[
            dict(type='Resize', keep_ratio=True),
            dict(type='RandomFlip'),
            dict(type='Normalize', **cfg.img_norm_cfg),
            dict(type='ImageToTensor', keys=['img']),
            dict(type='Collect', keys=['img']),
        ])
 ]


 cfg.data.train.type = cfg.dataset_type
 cfg.data.train.data_root = cfg.data_root
 cfg.data.train.img_dir = img_dir
 cfg.data.train.ann_dir = ann_dir
 cfg.data.train.pipeline = cfg.train_pipeline
 cfg.data.train.split = 'splits/train.txt'

 cfg.data.val.type = cfg.dataset_type
 cfg.data.val.data_root = cfg.data_root
 cfg.data.val.img_dir = img_dir
 cfg.data.val.ann_dir = ann_dir
 cfg.data.val.pipeline = cfg.test_pipeline
 cfg.data.val.split = 'splits/val.txt'

 cfg.data.test.type = cfg.dataset_type
 cfg.data.test.data_root = cfg.data_root
 cfg.data.test.img_dir = img_dir
 cfg.data.test.ann_dir = ann_dir
 cfg.data.test.pipeline = cfg.test_pipeline
 cfg.data.test.split = 'splits/val.txt'

 # We can still use the pre-trained Mask RCNN model though we do not need to
 # use the mask branch
 cfg.load_from = 'checkpoints/pspnet_r50-d8_512x1024_40k_cityscapes_20200605_003338-2966598c.pth'

 # Set up working dir to save files and logs.
 cfg.work_dir = './work_dirs/tutorial'

 cfg.runner.max_iters = 200
 cfg.log_config.interval = 10
 cfg.evaluation.interval = 200
 cfg.checkpoint_config.interval = 200

 # Set seed to facitate reproducing the result
 cfg.seed = 0
 set_random_seed(0, deterministic=False)
 cfg.gpu_ids = range(1)

 # Let's have a look at the final config used for training
 print(f'Config:\n{cfg.pretty_text}')
 ```

 %% Cell type:markdown id: tags:

 ### Train and Evaluation

 %% Cell type:code id: tags:

 ``` python
 from mmseg.datasets import build_dataset
 from mmseg.models import build_segmentor
 from mmseg.apis import train_segmentor


 # Build the dataset
 datasets = [build_dataset(cfg.data.train)]

 # Build the detector
-model = build_segmentor(
-    cfg.model, train_cfg=cfg.get('train_cfg'), test_cfg=cfg.get('test_cfg'))
+model = build_segmentor(cfg.model)
 # Add an attribute for visualization convenience
 model.CLASSES = datasets[0].CLASSES

 # Create work_dir
 mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))
 train_segmentor(model, datasets, cfg, distributed=False, validate=True,
                meta=dict())
 ```

 %% Cell type:markdown id: tags:

 Inference with trained model

 %% Cell type:code id: tags:

 ``` python
 img = mmcv.imread('iccv09Data/images/6000124.jpg')

 model.cfg = cfg
 result = inference_segmentor(model, img)
 plt.figure(figsize=(8, 6))
 show_result_pyplot(model, img, result, palette)
 ```

 %% Cell type:code id: tags:

 ``` python
 ```