Deeplab实战:使用deeplabv3实现对人物的抠图

语言: CN / TW / HK

持续创作,加速成长!这是我参与「掘金日新计划 · 6 月更文挑战」的第10天,点击查看活动详情

摘要

在上一篇文章中我们使用UNet实现了二分类分割,训练了150个epoch,最后dice得分在0.87左右。今天我们使用更优秀的网络deeplabv3实现图像的二分类分割,dice得分大约在0.97左右。

关于二分类一般有两种做法:

第一种输出是单通道,即网络的输出 output 为 [batch_size, 1, height, width] 形状。其中 batch_szie 为批量大小,1 表示输出一个通道,heightwidth 与输入图像的高和宽保持一致。

在训练时,输出通道数是 1,网络得到的 output 包含的数值是任意的数。给定的 target ,是一个单通道标签图,数值只有 0 和 1 这两种。为了让网络输出 output 不断逼近这个标签,首先会让 output 经过一个sigmoid 函数,使其数值归一化到[0, 1],得到 output1 ,然后让这个 output1target 进行交叉熵计算,得到损失值,反向传播更新网络权重。最终,网络经过学习,会使得 output1 逼近target

训练结束后,网络已经具备让输出的 output 经过转换从而逼近 target 的能力。首先将输出的 output 通过sigmoid 函数,然后取一个阈值(一般设置为0.5),大于阈值则取1反之则取0,从而得到预测图 predict。后续则是一些评估相关的计算。

如果网络的最后一层使用sigmoid,则选用BCELoss,如果没有则选择用BCEWithLogitsLoss,例:

最后一层没有sigmod

output = net(input) # net的最后一层没有使用sigmoid loss_func1 = torch.nn.BCEWithLogitsLoss() loss = loss_func1(output, target)

加上sigmod

output = net(input) # net的最后一层没有使用sigmoid output = F.sigmoid(output) loss_func1 = torch.nn.BCEWithLoss() loss = loss_func1(output, target)

预测的时:

output = net(input) # net的最后一层没有使用sigmoid output = F.sigmoid(output) predict=torch.where(output>0.5,torch.ones_like(output),torch.zeros_like(output))

第二种输出是多通道,即网络的输出 output 为 [batch_size, num_class, height, width] 形状。其中 batch_szie 为批量大小,num_class 表示输出的通道数与分类数量一致,heightwidth 与输入图像的高和宽保持一致。

在训练时,输出通道数是 num_class(这里取2)。给定的 target ,是一个单通道标签图,数值只有 0 和 1 这两种。为了让网络输出 output 不断逼近这个标签,首先会让 output 经过一个 softmax 函数,使其数值归一化到[0, 1],得到 output1 ,在各通道中,这个数值加起来会等于1。对于target 他是一个单通道图,首先使用onehot编码,转换成 num_class个通道的图像,每个通道中的取值是根据单通道中的取值计算出来的,例如单通道中的第一个像素取值为1(0<= 1 <=num_class-1,这里num_class=2),那么onehot编码后,在第一个像素的位置上,两个通道的取值分别为0,1。也就是说像素的取值决定了对应序号的通道取1,其他的通道取0,这个非常关键。上面的操作执行完后得到target1,让这个 output1target1 进行交叉熵计算,得到损失值,反向传播更新网路权重。最终,网络经过学习,会使得 output1 逼近target1(在各通道层面上)。

训练结束后,网络已经具备让输出的 output 经过转换从而逼近 target 的能力。计算 output 中各通道每一个像素位置上,取值最大的那个对应的通道序号,从而得到预测图 predict

训练选择用的loss是加插上损失函数,例:

output = net(input) # net的最后一层没有使用sigmoid loss_func = torch.nn.CrossEntropyLoss() loss = loss_func(output, target)

预测时

output = net(input) # net的最后一层没有使用sigmoid predict = output.argmax(dim=1)

本次实战选用的第二种做法。

选用的代码地址:milesial/Pytorch-UNet: PyTorch implementation of the U-Net for image semantic segmentation with high quality images (github.com)

下载代码后,解压到本地,如下图:

image-20220406094337124

数据集

数据集地址:http://www.cse.cuhk.edu.hk/~leojia/projects/automatting/,发布于2016年。

数据集包含2000张图,训练集1700张,测试集300张,数据都是来源于Flickr的肖像图,图像原始分辨率大小为600×800,其中Matting用closed-form matting和KNN matting方法生成。

由于肖像分割数据集商业价值较高,因此公开的大规模数据集很少,这个数据集是其中发布较早,使用范围也较广的一个数据集,它有几个比较重要的特点:

(1) 图像分辨率统一,拍摄清晰,质量很高。

(2) 所有图像均为上半身的肖像图,人像区域在长度和宽度均至少占据图像的2/3。

(3) 人物的姿态变化很小,都为小角度的正面图,背景较为简单。

img

img

img

[1] Shen X, Tao X, Gao H, et al. Deep Automatic Portrait Matting[M]// ComputerVision – ECCV 2016. Springer International Publishing, 2016:92-107.

将数据集下载后放到将训练集放到data文件夹中,其中图片放到imgs文件夹中,mask放到masks文件夹中,测试集放到test文件夹下面:

image-20220406094225993

由于原程序是用于Carvana Image Masking Challenge,所以我们需要修改加载数据集的逻辑,打开utils/data_loading.py文件:

python class CarvanaDataset(BasicDataset): def __init__(self, images_dir, masks_dir, scale=1): super().__init__(images_dir, masks_dir, scale, mask_suffix='_matte')

将mask_suffix改为“_matte”

训练

打开train.py,先查看全局参数:

Python def get_args(): parser = argparse.ArgumentParser(description='Train the UNet on images and target masks') parser.add_argument('--epochs', '-e', metavar='E', type=int, default=300, help='Number of epochs') parser.add_argument('--batch-size', '-b', dest='batch_size', metavar='B', type=int, default=16, help='Batch size') parser.add_argument('--learning-rate', '-l', metavar='LR', type=float, default=0.001, help='Learning rate', dest='lr') parser.add_argument('--load', '-f', type=str, default=False, help='Load model from a .pth file') parser.add_argument('--scale', '-s', type=float, default=0.5, help='Downscaling factor of the images') parser.add_argument('--validation', '-v', dest='val', type=float, default=10.0, help='Percent of the data that is used as validation (0-100)') parser.add_argument('--amp', action='store_true', default=False, help='Use mixed precision') return parser.parse_args()

epochs:epoch的个数,一般设置为300。

batch-size:批处理的大小,根据显存的大小设置。

learning-rate:学习率,一般设置为0.001,如果优化器不同,初始的学习率也要做相应的调整。

load:加载模型的路径,如果接着上次的训练,就需要设置上次训练的权重文件路径,如果有预训练权重,则设置预训练权重的路径。

scale:放大的倍数,这里设置为0.5,把图片大小变为原来的一半。

validation:验证验证集的百分比。

amp:是否使用混合精度?

比较重要的参数是epochs、batch-size和learning-rate,可以反复调整做实验,达到最好的精度。

接下来是设置模型:

```Python net = deeplabv3_resnet50(pretrained=False,num_classes=2) print(net)

if args.load:
    net.load_state_dict(torch.load(args.load, map_location=device))
    logging.info(f'Model loaded from {args.load}')

```

设置deeplabv3_resnet50参数:

导入方式:

python from torchvision.models.segmentation import deeplabv3_resnet50

pretrained:是否使用预训练权重,我们选用false。如果选用false则默认加载resnet50的预训练权重,是true则会加载deeplabv3_resnet50_coco的预训练权重。如果你想用预训练权重可以这样做:

python net = deeplabv3_resnet50(pretrained=True) print(net) net.classifier[4] = torch.nn.Conv2d(256, 2, kernel_size=(1, 1), stride=(1, 1)) net.aux_classifier[4] = torch.nn.Conv2d(256, 2, kernel_size=(1, 1), stride=(1, 1)) print(net)

但是要注意,选择预训练选中后,aux_classifier也会包含在内,也要修改aux_classifier的类别个数,2代表num_classes。

num_classes:类别,我们这里把背景当作一类,所以设置为2。

``` try: dataset = CarvanaDataset(dir_img, dir_mask, img_scale) except (AssertionError, RuntimeError): dataset = BasicDataset(dir_img, dir_mask, img_scale)

2. Split into train / validation partitions

n_val = int(len(dataset) * val_percent) n_train = len(dataset) - n_val train_set, val_set = random_split(dataset, [n_train, n_val], generator=torch.Generator().manual_seed(0))

3. Create data loaders

loader_args = dict(batch_size=batch_size, num_workers=4, pin_memory=True) train_loader = DataLoader(train_set, shuffle=True, loader_args) val_loader = DataLoader(val_set, shuffle=False, drop_last=True, loader_args) ```

1、加载数据集。

2、按照比例切分训练集和验证集。

3、将训练集和验证集放入DataLoader中。

```python # (Initialize logging) experiment = wandb.init(project='deeplabv3', resume='allow', anonymous='must') experiment.config.update(dict(epochs=epochs, batch_size=batch_size, learning_rate=learning_rate, val_percent=val_percent, save_checkpoint=save_checkpoint, img_scale=img_scale, amp=amp))

```

设置wandb,wandb是一款非常好用的可视化工具。安装和使用方法见:https://blog.csdn.net/hhhhhhhhhhwwwwwwwwww/article/details/116124285。

python # 4. Set up the optimizer, the loss, the learning rate scheduler and the loss scaling for AMP optimizer = optim.RMSprop(net.parameters(), lr=learning_rate, weight_decay=1e-8, momentum=0.9) scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'max', patience=2) # goal: maximize Dice score grad_scaler = torch.cuda.amp.GradScaler(enabled=amp) criterion = nn.CrossEntropyLoss() global_step = 0

1、设置优化器optimizer为RMSprop,我也尝试了改为SGD,通常情况下SGD的表现好一些。但是在训练时发现,二者最终的结果都差不多。

2、ReduceLROnPlateau学习率调整策略,和keras的类似。本次选择用的是Dice score,所以将mode设置为max,当得分不再上升时,则降低学习率。

3、设置loss为 nn.CrossEntropyLoss()。交叉熵,多分类常用的loss。

接下来是train部分的逻辑,这里需要修改的如下:

```python masks_pred = net(images)['out'] print(masks_pred.shape)

for name,x in masks_pred.items():

true_masks = F.one_hot(true_masks.squeeze_(1), 2).permute(0, 3, 1, 2).float() ```

masks_pred = net(images)计算出来的结果是:collections.OrderedDict类型的。如果不理解可以看这里:

```python class _SimpleSegmentationModel(nn.Module): constants = ["aux_classifier"]

def __init__(self, backbone: nn.Module, classifier: nn.Module, aux_classifier: Optional[nn.Module] = None) -> None:
    super().__init__()
    _log_api_usage_once(self)
    self.backbone = backbone
    self.classifier = classifier
    self.aux_classifier = aux_classifier

def forward(self, x: Tensor) -> Dict[str, Tensor]:
    input_shape = x.shape[-2:]
    # contract: features is a dict of tensors
    features = self.backbone(x)
    result = OrderedDict()
    x = features["out"]
    x = self.classifier(x)
    x = F.interpolate(x, size=input_shape, mode="bilinear", align_corners=False)
    result["out"] = x

    if self.aux_classifier is not None:
        x = features["aux"]
        x = self.aux_classifier(x)
        x = F.interpolate(x, size=input_shape, mode="bilinear", align_corners=False)
        result["aux"] = x
    return result

```

这个类是DeepLabV3的父类,返回值是result["out"],如果有aux_classifier,再返回 result["aux"]。

所以我们需要的结果放在result["out"]中。masks_pred的shape是[batch, 2, 400, 300],2对应的是类别。true_masks.shape是[batch, 1, 400, 300],所以要对true_masks做onehot处理。如果直接对true_masks做onehot处理,你会发现处理后的shape是[batch, 1, 400, 300,2],这样就和masks_pred 对不上了,所以在做onehot之前,先将第二维(也就是1这一维度)去掉,这样onehot后的shape是[batch, 400, 300,2],然后调整顺序,和masks_pred 的维度对上。

接下来就要计算loss,loss分为两部分,一部分时交叉熵,另一部分是dice_loss,这两个loss各有优势,组合使用效果更优。dice_loss在utils/dice_sorce.py文件中,代码如下:

```python import torch from torch import Tensor

def dice_coeff(input: Tensor, target: Tensor, reduce_batch_first: bool = False, epsilon=1e-6): # Average of Dice coefficient for all batches, or for a single mask assert input.size() == target.size() if input.dim() == 2 and reduce_batch_first: raise ValueError(f'Dice: asked to reduce batch but got tensor without batch dimension (shape {input.shape})') if input.dim() == 2 or reduce_batch_first: inter = torch.dot(input.reshape(-1), target.reshape(-1)) sets_sum = torch.sum(input) + torch.sum(target) if sets_sum.item() == 0: sets_sum = 2 * inter return (2 * inter + epsilon) / (sets_sum + epsilon) else: # compute and average metric for each batch element dice = 0 for i in range(input.shape[0]): dice += dice_coeff(input[i, ...], target[i, ...]) return dice / input.shape[0]

def dice_coeff_1(pred, target): smooth = 1. num = pred.size(0) m1 = pred.view(num, -1) # Flatten m2 = target.view(num, -1) # Flatten intersection = (m1 * m2).sum() return 1 - (2. * intersection + smooth) / (m1.sum() + m2.sum() + smooth)

def multiclass_dice_coeff(input: Tensor, target: Tensor, reduce_batch_first: bool = False, epsilon=1e-6): # Average of Dice coefficient for all classes assert input.size() == target.size() dice = 0 for channel in range(input.shape[1]): dice += dice_coeff(input[:, channel, ...], target[:, channel, ...], reduce_batch_first, epsilon)

return dice / input.shape[1]

def dice_loss(input: Tensor, target: Tensor, multiclass: bool = False): # Dice loss (objective to minimize) between 0 and 1 assert input.size() == target.size() fn = multiclass_dice_coeff if multiclass else dice_coeff return 1 - fn(input, target, reduce_batch_first=True)

```

导入到train.py中,然后和交叉熵组合作为本项目的loss。

python loss = criterion(masks_pred, true_masks) \ + dice_loss(F.softmax(masks_pred, dim=1).float(), true_masks, multiclass=True)

接下来是对evaluate函数的逻辑做修改。

python mask_true = mask_true.to(device=device, dtype=torch.long) mask_true = F.one_hot(mask_true.squeeze_(1), net.n_classes).permute(0, 3, 1, 2).float() with torch.no_grad(): # predict the mask mask_pred = net(image)['out'] num_classes=2 # convert to one-hot format if num_classes == 1: mask_pred = (F.sigmoid(mask_pred) > 0.5).float() # compute the Dice score dice_score += dice_coeff(mask_pred, mask_true, reduce_batch_first=False) else: mask_pred = F.one_hot(mask_pred.argmax(dim=1), 2).permute(0, 3, 1, 2).float() # compute the Dice score, ignoring background dice_score += multiclass_dice_coeff(mask_pred[:, 1:, ...], mask_true[:, 1:, ...], reduce_batch_first=False)

增加对mask_trued的onehot逻辑。

定义num_classes为2,由于官方的模型没有类别数量的定义,所以只能自己定义类别数。

修改完上面的逻辑就可以开始训练了。

image-20220408130729405

测试

完成训练后就可以测试了。打开predict.py,修改全局参数:

Python def get_args(): parser = argparse.ArgumentParser(description='Predict masks from input images') parser.add_argument('--model', '-m', default='checkpoints/checkpoint_epoch150.pth', metavar='FILE', help='Specify the file in which the model is stored') parser.add_argument('--input', '-i', metavar='INPUT',default='test/00002.png', nargs='+', help='Filenames of input images') parser.add_argument('--output', '-o', metavar='INPUT',default='00001.png', nargs='+', help='Filenames of output images') parser.add_argument('--viz', '-v', action='store_true', help='Visualize the images as they are processed') parser.add_argument('--no-save', '-n', action='store_true',default=False, help='Do not save the output masks') parser.add_argument('--mask-threshold', '-t', type=float, default=0.5, help='Minimum probability value to consider a mask pixel white') parser.add_argument('--scale', '-s', type=float, default=0.5, help='Scale factor for the input images')

model:设置权重文件路径。这里要改为自己训练的路径。

scale:0.5,和训练的参数对应上。

其他的参数,通过命令输入。

python net = deeplabv3_resnet50(pretrained=False,num_classes=2) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') logging.info(f'Loading model {args.model}') logging.info(f'Using device {device}') net.to(device=device) net.load_state_dict(torch.load(args.model, map_location=device))

定义网络deeplabv3_resnet50。

加载训练好的权重文件。

接下来看预测部分:

```Python with torch.no_grad(): output = net(img)['out'] n_classes=2 if n_classes > 1: probs = F.softmax(output, dim=1)[0] else: probs = torch.sigmoid(output)[0] tf = transforms.Compose([ transforms.ToPILImage(), transforms.Resize((full_img.size[1], full_img.size[0])), transforms.ToTensor() ])

    full_mask = tf(probs.cpu()).squeeze()


if n_classes == 1:
    return (full_mask > out_threshold).numpy()
else:
    return F.one_hot(full_mask.argmax(dim=0), n_classes).permute(2, 0, 1).numpy()

```

输出预测结果。

定义类别为2.

然后输入到softmax中。

将softmax输出的结果,做one_hot,然后返回。

Python def mask_to_image(mask: np.ndarray): if mask.ndim == 2: return Image.fromarray((mask * 255).astype(np.uint8)) elif mask.ndim == 3: img_np=(np.argmax(mask, axis=0) * 255 / (mask.shape[0]-1)).astype(np.uint8) print(img_np.shape) print(np.max(img_np)) return Image.fromarray(img_np)

img_np=(np.argmax(mask, axis=0) * 255 / (mask.shape[0]-1)).astype(np.uint8)这里的逻辑需要修改。

源代码:

return Image.fromarray((np.argmax(mask, axis=0) * 255 / mask.shape[0]).astype(np.uint8))

我们增加了一类背景,所以mask.shape[0]为2,需要减去背景。

展示结果的方法也需要修改;

python def plot_img_and_mask(img, mask): print(mask.shape) classes = mask.shape[0] if len(mask.shape) > 2 else 1 fig, ax = plt.subplots(1, classes + 1) ax[0].set_title('Input image') ax[0].imshow(img) if classes > 1: for i in range(classes): ax[i + 1].set_title(f'Output mask (class {i + 1})') ax[i + 1].imshow(mask[i, :, :]) else: ax[1].set_title(f'Output mask') ax[1].imshow(mask) plt.xticks([]), plt.yticks([]) plt.show()

将原来的ax[i + 1].imshow(mask[:, :, i])改为:ax[i + 1].imshow(mask[i, :, :])。

执行命令:

bash python predict.py -i test/00002.png -o output.png -v

输出结果:

image-20220408131901870

到这里我们已经实现将人物从背景图片中完整的抠出来了!

总结

本文实现了用deeplabv3对图像做分割,通过本文,你可以学习到:

1、如何使用pytorch自带deeplabv3对图像对二分类的语义分割。pytorch自带的deeplabv3,除了deeplabv3_resnet50,还有deeplabv3_resnet101,deeplabv3_mobilenet_v3_large,大家可以尝试更换模型做测试。

2、如何使用wandb可视化。

3、如何使用交叉熵和dice_loss组合。

4、如何实现二分类语义分割的预测。

完整的代码: https://download.csdn.net/download/hhhhhhhhhhwwwwwwwwww/85093662