最近在研究语义分割，于是想用它做点有趣的事情，很喜欢狗狗的我决定用它来识别图片中猫的位置。我会介绍从原始数据到结果的整个过程，整个过程包括 数据标注-模型训练-验证结果。
数据标注
在工业应用中，我们拿到的原始数据一般都是不带标注的，不像参加竞赛给你的都是打好标注的干净的数据。所以你可能没有很多时间去优化算法，大部分时间都是在处理数据上，包括数据的清洗，预处理，数据标注等等，如果数据量比较大的话可能需要雇佣第三方平台来做数据标注工作。
这次我使用的数据集是 kaggle 猫狗分类 比赛的数据集，比赛的内容是区分图片中是猫还是狗，所以数据集只标注了每张图片的类别，而语义分割需要对每个像素做标注。由于我的时间有限(主要是手酸)，我挑选了其中近 200 张猫的图像进行标注，就这 200 张也标了我好几个小时，T. T。
介绍一个专门做标注的工具 - labelme，github 地址: https://github.com/wkentaro/labelme
直接使用 pip 安装:
pip install labelme
安装完成之后 cmd 输入 labelme 就启动了(我是 windows 操作系统)。用法如下:
- 第一步，选择打开文件或者文件夹，打开文件夹的话会自动遍历文件夹下所有图片。
- 第二步，选择标注的形状，图中框出来的是多边形(Polygon)，你也可以点击菜单栏上的 edit 选择其他图形例如矩形(Rectangle)，圆形(Circle)等。多边形是我最常用的。
[站外图片上传中...(image-dd5577-1568188101476)]
第三步，标注完成之后点击保存，会将标注信息保存为 json 格式。
第四步，将保存的 json 文件转换为 image 以及 label 图像，labelme 自带了转换工具，运行下面的 python 代码可以将指定文件夹下所有 json 文件统一转换。
import os
from glob import glob
json_files = glob(os.path.join(r"your_dir", "*.json"))
for json_file in json_files:
os.system(f"labelme_json_to_dataset {json_file} -o {''.join(json_file.split('.')[:-1])}")
转换之后每个 json 文件会生成一个对应的文件夹，文件夹下包括了 image.png-原图，label.png-标注图像，label_viz.png-标注可视化图像:
至此，数据的准备工作就完成了，下面我们用代码来读取这些数据并进行训练。
模型训练
环境:
- Win10
- python 3.7.3
- pytorch 1.1.0
- cuda 10.1
首先构建数据类 CatDataset，负责读取图像数据以及数据增强，这里我只用了简单的水平和垂直翻转:
def get_transforms(phase, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
list_transforms = []
if phase == "train":
list_transforms.extend(
[
HorizontalFlip(),
VerticalFlip()
]
)
list_transforms.extend(
[
Resize(128, 128, interpolation=Image.BILINEAR),
Normalize(mean=mean, std=std, p=1),
ToTensor(),
]
)
list_trfms = Compose(list_transforms)
return list_trfms
class CatDataset(Dataset):
def __init__(self, idx, image_list, phase="train"):
assert phase in ("train", "val", "test")
self.idx = idx
self.image_list = image_list
self.phase = phase
self.transform = get_transforms(phase)
def __getitem__(self, index):
real_idx = self.idx[index]
image_path = os.path.join(self.image_list[real_idx], "img.png")
mask_path = os.path.join(self.image_list[real_idx], "label.png")
image = cv2.imread(image_path)
mask = cv2.imread(mask_path)
augmented = self.transform(image=image, mask=mask/128)
return augmented["image"], augmented["mask"][:, :, :, 2]
def __len__(self):
return len(self.idx)
模型选择之前刚研究过的 Unet，由于数据量较少，在选用 U-Net 的 backbone 时应该避免参数量太大的模型，所以我先选择了 Resnet-18 做为 Encoder 部分。
import segmentation_models_pytorch as smp
model = smp.Unet('resnet18', classes=1, activation=None)
给大家介绍一个 pytorch 写的第三方库(没错我就是调库工程师): segmentation_models_pytorch，该库把基于 Encoder-Decoder 的几个语义分割模型做了一个整合，包括 FPN，UNet，PSPNet 等，你可以根据自己的需要配置这些模型的 Encoder 部分，支持包括 Resnet，DenseNet，VGG 等 backbone。注意安装时要从 github 上下载源码安装，用 pip 安装会有小 bug，github 地址: https://github.com/qubvel/segmentation_models.pytorch
数据以及模型准备完毕，下面可以开始训练了，我的训练参数:
- batch size=32
- lr=1e-3
- optimizer=Adam
- loss=BCEWithLogitsLoss
- epoch=30
因为数据量比较小，训练只花了一首歌的时间，在经过了 30 轮训练之后，模型收敛，可以看看模型的表现了。
训练结果
Resnet18 做为 backbone 的模型训练结果如下，最好的表现在第 15 个 epoch，训练集上的 IoU 为 0.8777，验证集为 0.8583。
Starting epoch: 15 | phase: train | ⏰: 08:48:16
Loss: 0.127440 | IoU: 0.8777 | dice: 0.9312 | dice_neg: 0.0000 | dice_pos: 0.9312
Starting epoch: 15 | phase: val | ⏰: 08:48:20
Loss: 0.165582 | IoU: 0.8583 | dice: 0.9203 | dice_neg: 0.0000 | dice_pos: 0.9203
******** New optimal found, saving state ********
发现训练集拟合的不够好，于是尝试用一个稍复杂点的 backbone - Resnet50，训练大概花了一首 《情歌王》的时间。结果如下:
Starting epoch: 40 | phase: train | ⏰: 09:25:38
Loss: 0.073017 | IoU: 0.9285 | dice: 0.9624 | dice_neg: 0.0000 | dice_pos: 0.9624
Starting epoch: 40 | phase: val | ⏰: 09:25:43
Loss: 0.150036 | IoU: 0.8785 | dice: 0.9330 | dice_neg: 0.0000 | dice_pos: 0.9330
******** New optimal found, saving state ********
表现有明显的提升，且计算时间也在我们可以接受的范围内，这个模型是更加适合的。我们在选择模型的时候，需要综合考虑其计算量与精度，在两者之间达到平衡是最好的。
这个模型收敛用了 40 个 epoch，这很符合常理，因为更复杂的模型意味这更多的参数，在训练时需要训练更多 epoch 才能收敛。
我在验证集上用 visdom 做了可视化，见证奇迹的时刻:
图中黑白的就是模型预测的猫的位置，在训练集只有 200 张图片不到的情况下，模型的表现已经很不错了。当然还有很大的提升空间，可以尝试的包括增加数据集，更复杂的模型，多模型融合等方法。
整个项目的代码已经放在我的 github，欢迎 star，地址: https://github.com/FroyoZzz/CV-Papers-Codes
最后，求赞求关注，欢迎关注我的微信公众号[MachineLearning学习之路] ，深度学习 & CV 方向的童鞋不要错过！！