text_plain_output_1.png
from PIL import Image

from mmcv import Config

from mmseg.apis import inference_segmentor, init_segmentor, show_result_pyplot

from mmseg.apis import set_random_seed

from mmseg.apis import train_segmentor

from mmseg.datasets import build_dataset

from mmseg.datasets.builder import DATASETS

from mmseg.datasets.custom import CustomDataset

from mmseg.models import build_segmentor

import matplotlib.patches as mpatches

import matplotlib.pyplot as plt

import mmcv

import numpy as np

import os.path as osp

import mmcv
import matplotlib.pyplot as plt
img = mmcv.imread('nails/images/2C29D473-CCB4-458C-926B-99D0042161E6.jpg')

import os.path as osp
import numpy as np
from PIL import Image
data_root = 'nails'
img_dir = 'images'
ann_dir = 'labels'
classes = ('background', 'nails')
palette = [[0, 0, 0], [225, 25, 38]]
for file in mmcv.scandir(osp.join(data_root, ann_dir), suffix='.jpg'):
    seg_img = Image.open(osp.join(data_root, ann_dir, file)).convert('P')
    seg_img.putpalette(np.array(palette, dtype=np.uint8))
    seg_img.save(osp.join(data_root, ann_dir, file.replace('.jpg', '.png')))

# Let's take a look at the segmentation map we got
import matplotlib.patches as mpatches
img = Image.open('/content/mmsegmentation/nails/labels/09aefeec-e05f-11e8-87a6-0242ac1c0002.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(2)]
# 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()

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:
    train_length = int(len(filename_list) * 0.85)
    f.writelines((line + '\n' for line in filename_list[:train_length]))
with open(osp.join(data_root, split_dir, 'val.txt'), 'w') as f:
    f.writelines((line + '\n' for line in filename_list[train_length:]))

from mmseg.datasets.builder import DATASETS
from mmseg.datasets.custom import CustomDataset
@DATASETS.register_module()
class NailsDataset(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

from mmcv import Config
cfg = Config.fromfile('configs/pspnet/pspnet_r101-d8_512x1024_80k_cityscapes.py ')

from mmseg.apis import set_random_seed
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
cfg.model.decode_head.num_classes = 2
cfg.model.auxiliary_head.num_classes = 2
cfg.model.decode_head.loss_decode = [dict(type='CrossEntropyLoss', loss_name='loss_ce', loss_weight=0.5), dict(type='DiceLoss', loss_name='loss_dice', loss_weight=0.5)]
cfg.model.auxiliary_head.loss_decode = [dict(type='CrossEntropyLoss', loss_name='loss_ce', loss_weight=0.5), dict(type='DiceLoss', loss_name='loss_dice', loss_weight=0.5)]
cfg.dataset_type = 'NailsDataset'
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), 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'
cfg.load_from = 'checkpoints/deeplabv3plus_r101-d8_512x1024_80k_cityscapes_20200605_094614-3769eecf.pth'
cfg.work_dir = './work_dirs/nails'
cfg.optimizer.lr = 0.01 / 8
cfg.runner.max_iters = 600
cfg.log_config.interval = 10
cfg.evaluation.interval = 200
cfg.checkpoint_config.interval = 200
cfg.seed = 1111
set_random_seed(1111, deterministic=False)
cfg.gpu_ids = range(1)

from mmseg.datasets import build_dataset
from mmseg.models import build_segmentor
from mmseg.apis import train_segmentor
datasets = [build_dataset(cfg.data.train)]
model = build_segmentor(cfg.model)
model.CLASSES = datasets[0].CLASSES
mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))
train_segmentor(model, datasets, cfg, distributed=False, validate=True, meta=dict())

img = mmcv.imread('/content/mmsegmentation/nails/images/09aefeec-e05f-11e8-87a6-0242ac1c0002.jpg')
model.cfg = cfg
result = inference_segmentor(model, img)
plt.figure(figsize=(8, 6))
show_result_pyplot(model, img, result, palette)