{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Train a Segmentation Model\n",
"\n",
"This segmentation task example will be divided into the following steps:\n",
"\n",
"- [Download Camvid Dataset](#download-camvid-dataset)\n",
"- [Implement Camvid Dataset](#implement-the-camvid-dataset)\n",
"- [Implement a Segmentation Model](#implement-the-segmentation-model)\n",
"- [Train with Runner](#training-with-runner)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Download Camvid Dataset\n",
"\n",
"First, you should download the Camvid dataset from OpenDataLab:\n",
"\n",
"```bash\n",
"# https://opendatalab.com/CamVid\n",
"# Configure install\n",
"pip install opendatalab\n",
"# Upgraded version\n",
"pip install -U opendatalab\n",
"# Login\n",
"odl login\n",
"# Download this dataset\n",
"mkdir data\n",
"odl get CamVid -d data\n",
"# Preprocess data in Linux. You should extract the files to data manually in\n",
"# Windows\n",
"tar -xzvf data/CamVid/raw/CamVid.tar.gz.00 -C ./data\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Implement the Camvid Dataset\n",
"\n",
"We have implemented the CamVid class here, which inherits from VisionDataset. Within this class, we have overridden the `__getitem__` and `__len__` methods to ensure that each index returns a dict of images and labels. Additionally, we have implemented the color_to_class dictionary to map the mask's color to the class index.\""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"import numpy as np\n",
"from torchvision.datasets import VisionDataset\n",
"from PIL import Image\n",
"import csv\n",
"\n",
"\n",
"def create_palette(csv_filepath):\n",
" color_to_class = {}\n",
" with open(csv_filepath, newline='') as csvfile:\n",
" reader = csv.DictReader(csvfile)\n",
" for idx, row in enumerate(reader):\n",
" r, g, b = int(row['r']), int(row['g']), int(row['b'])\n",
" color_to_class[(r, g, b)] = idx\n",
" return color_to_class\n",
"\n",
"class CamVid(VisionDataset):\n",
"\n",
" def __init__(self,\n",
" root,\n",
" img_folder,\n",
" mask_folder,\n",
" transform=None,\n",
" target_transform=None):\n",
" super().__init__(\n",
" root, transform=transform, target_transform=target_transform)\n",
" self.img_folder = img_folder\n",
" self.mask_folder = mask_folder\n",
" self.images = list(\n",
" sorted(os.listdir(os.path.join(self.root, img_folder))))\n",
" self.masks = list(\n",
" sorted(os.listdir(os.path.join(self.root, mask_folder))))\n",
" self.color_to_class = create_palette(\n",
" os.path.join(self.root, 'class_dict.csv'))\n",
"\n",
" def __getitem__(self, index):\n",
" img_path = os.path.join(self.root, self.img_folder, self.images[index])\n",
" mask_path = os.path.join(self.root, self.mask_folder,\n",
" self.masks[index])\n",
"\n",
" img = Image.open(img_path).convert('RGB')\n",
" mask = Image.open(mask_path).convert('RGB') # Convert to RGB\n",
"\n",
" if self.transform is not None:\n",
" img = self.transform(img)\n",
"\n",
" # Convert the RGB values to class indices\n",
" mask = np.array(mask)\n",
" mask = mask[:, :, 0] * 65536 + mask[:, :, 1] * 256 + mask[:, :, 2]\n",
" labels = np.zeros_like(mask, dtype=np.int64)\n",
" for color, class_index in self.color_to_class.items():\n",
" rgb = color[0] * 65536 + color[1] * 256 + color[2]\n",
" labels[mask == rgb] = class_index\n",
"\n",
" if self.target_transform is not None:\n",
" labels = self.target_transform(labels)\n",
" data_samples = dict(\n",
" labels=labels, img_path=img_path, mask_path=mask_path)\n",
" return img, data_samples\n",
"\n",
" def __len__(self):\n",
" return len(self.images)\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We utilize the Camvid dataset to create the `train_dataloader` and `val_dataloader`, which serve as the data loaders for training and validation in the subsequent Runner."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"import torchvision.transforms as transforms\n",
"\n",
"norm_cfg = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
"transform = transforms.Compose(\n",
" [transforms.ToTensor(),\n",
" transforms.Normalize(**norm_cfg)])\n",
"\n",
"target_transform = transforms.Lambda(\n",
" lambda x: torch.tensor(np.array(x), dtype=torch.long))\n",
"\n",
"train_set = CamVid(\n",
" 'data/CamVid',\n",
" img_folder='train',\n",
" mask_folder='train_labels',\n",
" transform=transform,\n",
" target_transform=target_transform)\n",
"\n",
"valid_set = CamVid(\n",
" 'data/CamVid',\n",
" img_folder='val',\n",
" mask_folder='val_labels',\n",
" transform=transform,\n",
" target_transform=target_transform)\n",
"\n",
"train_dataloader = dict(\n",
" batch_size=3,\n",
" dataset=train_set,\n",
" sampler=dict(type='DefaultSampler', shuffle=True),\n",
" collate_fn=dict(type='default_collate'))\n",
"\n",
"val_dataloader = dict(\n",
" batch_size=3,\n",
" dataset=valid_set,\n",
" sampler=dict(type='DefaultSampler', shuffle=False),\n",
" collate_fn=dict(type='default_collate'))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Implement the Segmentation Model\n",
"\n",
"The provided code defines a model class named `MMDeeplabV3`. This class is derived from `BaseModel` and incorporates the segmentation model of the DeepLabV3 architecture. It overrides the `forward` method to handle both input images and labels and supports computing losses and returning predictions in both training and prediction modes.\n",
"\n",
"For additional information about `BaseModel`, you can refer to the [Model tutorial](../tutorials/model.md)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from mmengine.model import BaseModel\n",
"from torchvision.models.segmentation import deeplabv3_resnet50\n",
"import torch.nn.functional as F\n",
"\n",
"\n",
"class MMDeeplabV3(BaseModel):\n",
"\n",
" def __init__(self, num_classes):\n",
" super().__init__()\n",
" self.deeplab = deeplabv3_resnet50(num_classes=num_classes)\n",
"\n",
" def forward(self, imgs, data_samples=None, mode='tensor'):\n",
" x = self.deeplab(imgs)['out']\n",
" if mode == 'loss':\n",
" return {'loss': F.cross_entropy(x, data_samples['labels'])}\n",
" elif mode == 'predict':\n",
" return x, data_samples"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Training with Runner\n",
"\n",
"Before training with the Runner, we need to implement the IoU (Intersection over Union) metric to evaluate the model's performance."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from mmengine.evaluator import BaseMetric\n",
"\n",
"class IoU(BaseMetric):\n",
"\n",
" def process(self, data_batch, data_samples):\n",
" preds, labels = data_samples[0], data_samples[1]['labels']\n",
" preds = torch.argmax(preds, dim=1)\n",
" intersect = (labels == preds).sum()\n",
" union = (torch.logical_or(preds, labels)).sum()\n",
" iou = (intersect / union).cpu()\n",
" self.results.append(\n",
" dict(batch_size=len(labels), iou=iou * len(labels)))\n",
"\n",
" def compute_metrics(self, results):\n",
" total_iou = sum(result['iou'] for result in self.results)\n",
" num_samples = sum(result['batch_size'] for result in self.results)\n",
" return dict(iou=total_iou / num_samples)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Implementing a visualization hook is also important to facilitate easier comparison between predictions and labels."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from mmengine.hooks import Hook\n",
"import shutil\n",
"import cv2\n",
"import os.path as osp\n",
"\n",
"\n",
"class SegVisHook(Hook):\n",
"\n",
" def __init__(self, data_root, vis_num=1) -> None:\n",
" super().__init__()\n",
" self.vis_num = vis_num\n",
" self.palette = create_palette(osp.join(data_root, 'class_dict.csv'))\n",
"\n",
" def after_val_iter(self,\n",
" runner,\n",
" batch_idx: int,\n",
" data_batch=None,\n",
" outputs=None) -> None:\n",
" if batch_idx > self.vis_num:\n",
" return\n",
" preds, data_samples = outputs\n",
" img_paths = data_samples['img_path']\n",
" mask_paths = data_samples['mask_path']\n",
" _, C, H, W = preds.shape\n",
" preds = torch.argmax(preds, dim=1)\n",
" for idx, (pred, img_path,\n",
" mask_path) in enumerate(zip(preds, img_paths, mask_paths)):\n",
" pred_mask = np.zeros((H, W, 3), dtype=np.uint8)\n",
" runner.visualizer.set_image(pred_mask)\n",
" for color, class_id in self.palette.items():\n",
" runner.visualizer.draw_binary_masks(\n",
" pred == class_id,\n",
" colors=[color],\n",
" alphas=1.0,\n",
" )\n",
" # Convert RGB to BGR\n",
" pred_mask = runner.visualizer.get_image()[..., ::-1]\n",
" saved_dir = osp.join(runner.log_dir, 'vis_data', str(idx))\n",
" os.makedirs(saved_dir, exist_ok=True)\n",
"\n",
" shutil.copyfile(img_path,\n",
" osp.join(saved_dir, osp.basename(img_path)))\n",
" shutil.copyfile(mask_path,\n",
" osp.join(saved_dir, osp.basename(mask_path)))\n",
" cv2.imwrite(\n",
" osp.join(saved_dir, f'pred_{osp.basename(img_path)}'),\n",
" pred_mask)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Finnaly, just train the model with Runner!"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from torch.optim import AdamW\n",
"from mmengine.optim import AmpOptimWrapper\n",
"from mmengine.runner import Runner\n",
"\n",
"\n",
"num_classes = 32 # Modify to actual number of categories.\n",
"\n",
"runner = Runner(\n",
" model=MMDeeplabV3(num_classes),\n",
" work_dir='./work_dir',\n",
" train_dataloader=train_dataloader,\n",
" optim_wrapper=dict(\n",
" type=AmpOptimWrapper, optimizer=dict(type=AdamW, lr=2e-4)),\n",
" train_cfg=dict(by_epoch=True, max_epochs=10, val_interval=10),\n",
" val_dataloader=val_dataloader,\n",
" val_cfg=dict(),\n",
" val_evaluator=dict(type=IoU),\n",
" custom_hooks=[SegVisHook('data/CamVid')],\n",
" default_hooks=dict(checkpoint=dict(type='CheckpointHook', interval=1)),\n",
")\n",
"runner.train()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Finnaly, you can check the training results in the folder `./work_dir/{timestamp}/vis_data`.\n",
"\n",
"<table class=\"docutils\">\n",
"<thead>\n",
"<tr>\n",
" <th>image</th>\n",
" <th>prediction</th>\n",
" <th>label</th>\n",
"</tr>\n",
"<tr>\n",
" <th><img src=\"https://github.com/open-mmlab/mmengine/assets/57566630/de70c138-fb8e-402c-9497-574b01725b6c\" width=\"200\"></th>\n",
" <th><img src=\"https://github.com/open-mmlab/mmengine/assets/57566630/ea9221e7-48ca-4515-8815-56b5ff091f53\" width=\"200\"></th>\n",
" <th><img src=\"https://github.com/open-mmlab/mmengine/assets/57566630/dcb2324f-a2df-4e5c-a038-df896dde2471\" width=\"200\"></th>\n",
"</tr>\n",
"</thead>\n",
"</table>"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "py310torch20",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.9"
},
"orig_nbformat": 4
},
"nbformat": 4,
"nbformat_minor": 2
}