Deeplab實戰：使用deeplabv3實現對人物的摳圖
持續創作，加速成長！這是我參與「掘金日新計劃 · 6 月更文挑戰」的第10天，點擊查看活動詳情
摘要
在上一篇文章中我們使用UNet實現了二分類分割，訓練了150個epoch，最後dice得分在0.87左右。今天我們使用更優秀的網絡deeplabv3實現圖像的二分類分割，dice得分大約在0.97左右。
關於二分類一般有兩種做法：
第一種輸出是單通道，即網絡的輸出 output 為 [batch_size, 1, height, width] 形狀。其中 batch_szie 為批量大小，1 表示輸出一個通道，height 和 width 與輸入圖像的高和寬保持一致。
在訓練時，輸出通道數是 1，網絡得到的 output 包含的數值是任意的數。給定的 target ，是一個單通道標籤圖，數值只有 0 和 1 這兩種。為了讓網絡輸出 output 不斷逼近這個標籤，首先會讓 output 經過一個sigmoid 函數，使其數值歸一化到[0, 1]，得到 output1 ，然後讓這個 output1 與 target 進行交叉熵計算，得到損失值，反向傳播更新網絡權重。最終，網絡經過學習，會使得 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 表示輸出的通道數與分類數量一致，height 和 width 與輸入圖像的高和寬保持一致。
在訓練時，輸出通道數是 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，讓這個 output1 與 target1 進行交叉熵計算，得到損失值，反向傳播更新網路權重。最終，網絡經過學習，會使得 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)
本次實戰選用的第二種做法。
下載代碼後，解壓到本地，如下圖：
數據集
數據集地址：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) 人物的姿態變化很小，都為小角度的正面圖，背景較為簡單。
[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文件夾下面：
由於原程序是用於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是一款非常好用的可視化工具。安裝和使用方法見：http://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，由於官方的模型沒有類別數量的定義，所以只能自己定義類別數。
修改完上面的邏輯就可以開始訓練了。
測試
完成訓練後就可以測試了。打開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
輸出結果：
到這裏我們已經實現將人物從背景圖片中完整的摳出來了！
總結
本文實現了用deeplabv3對圖像做分割，通過本文，你可以學習到：
1、如何使用pytorch自帶deeplabv3對圖像對二分類的語義分割。pytorch自帶的deeplabv3，除了deeplabv3_resnet50,還有deeplabv3_resnet101,deeplabv3_mobilenet_v3_large，大家可以嘗試更換模型做測試。
2、如何使用wandb可視化。
3、如何使用交叉熵和dice_loss組合。
4、如何實現二分類語義分割的預測。
完整的代碼： http://download.csdn.net/download/hhhhhhhhhhwwwwwwwwww/85093662
- YoloV5實戰：手把手教物體檢測——YoloV5
- 基於阿里Semantatic Human Matting算法，實現精細化人物摳圖
- PPv3-OCR自定義數據從訓練到部署
- 如何下載pytorch的歷史版本？
- WinForm——Button總結
- WinForm——MDI窗體
- 升級 pip
- 將8位的tif圖片改為png圖片
- RepLKNet實戰：使用RepLKNet實現對植物幼苗的分類（非官方）（二）
- 關於OpenCV imread和imdecode讀取圖片是BGR的證明
- opencv讀取圖片通道以及顯示
- 萬字整理聯邦學習系統架構設計參考
- 編譯器堆空間不足
- 【圖像分類】實戰——使用EfficientNetV2實現圖像分類（Pytorch）
- MMDetection實戰：MMDetection訓練與測試
- UNet語義分割實戰：使用UNet實現對人物的摳圖
- MobileVIT實戰：使用MobileVIT實現圖像分類
- SwinIR實戰：如何使用SwinIR和預訓練模型實現圖片的超分
- 【圖像分類】手撕ResNet——復現ResNet（Pytorch）
- Deeplab實戰：使用deeplabv3實現對人物的摳圖