正如论文里所给出的 BackgroundMattingV2 网络结构图如下：网路结构包含两个网络组成：[1] - base 网络在将采样后的分辨率上输出粗略抠图结果;[2] - refine 网络仅对容易出错(error-prone)的图像部分进行处理，输出全分辨率的结果. 其能够节省冗余计算量，有助于调整耗时.1. 总体模型定义model.pyimport torch from torch import nn from torch.nn import functional as F from torchvision.models.segmentation.deeplabv3 import ASPP from .decoder import Decoder from .mobilenet import MobileNetV2Encoder from .refiner import Refiner from .resnet import ResNetEncoder from .utils import load_matched_state_dict class Base(nn.Module): """ A generic implementation of the base encoder-decoder network inspired by DeepLab. Accepts arbitrary channels for input and output. """ def __init__(self, backbone: str, in_channels: int, out_channels: int): super().__init__() #支持的 Encoder-Decoder Base 网络 assert backbone in ["resnet50", "resnet101", "mobilenetv2"] if backbone in ['resnet50', 'resnet101']: self.backbone = ResNetEncoder(in_channels, variant=backbone) self.aspp = ASPP(2048, [3, 6, 9]) self.decoder = Decoder([256, 128, 64, 48, out_channels], [512, 256, 64, in_channels]) else: self.backbone = MobileNetV2Encoder(in_channels) self.aspp = ASPP(320, [3, 6, 9]) self.decoder = Decoder([256, 128, 64, 48, out_channels], [32, 24, 16, in_channels]) def forward(self, x): x, *shortcuts = self.backbone(x) x = self.aspp(x) x = self.decoder(x, *shortcuts) return x def load_pretrained_deeplabv3_state_dict(self, state_dict, print_stats=True): # Pretrained DeepLabV3 models are provided by <https://github.com/VainF/DeepLabV3Plus-Pytorch>. # This method converts and loads their pretrained state_dict to match with our model structure. # This method is not needed if you are not planning to train from deeplab weights. # Use load_state_dict() for normal weight loading. # Convert state_dict naming for aspp module state_dict = {k.replace('classifier.classifier.0', 'aspp'): v for k, v in state_dict.items()} if isinstance(self.backbone, ResNetEncoder): # ResNet backbone does not need change. load_matched_state_dict(self, state_dict, print_stats) else: # Change MobileNetV2 backbone to state_dict format, then change back after loading. backbone_features = self.backbone.features self.backbone.low_level_features = backbone_features[:4] self.backbone.high_level_features = backbone_features[4:] del self.backbone.features # load_matched_state_dict(self, state_dict, print_stats) self.backbone.features = backbone_features del self.backbone.low_level_features del self.backbone.high_level_features #网络结构中的 Base 网络 class MattingBase(Base): """ MattingBase is used to produce coarse global results at a lower resolution. MattingBase extends Base. Args: backbone: ["resnet50", "resnet101", "mobilenetv2"] Input: src: (B, 3, H, W) the source image. Channels are RGB values normalized to 0 ~ 1. bgr: (B, 3, H, W) the background image . Channels are RGB values normalized to 0 ~ 1. Output: pha: (B, 1, H, W) the alpha prediction. Normalized to 0 ~ 1. fgr: (B, 3, H, W) the foreground prediction. Channels are RGB values normalized to 0 ~ 1. err: (B, 1, H, W) the error prediction. Normalized to 0 ~ 1. hid: (B, 32, H, W) the hidden encoding. Used for connecting refiner module. Example: model = MattingBase(backbone='resnet50') pha, fgr, err, hid = model(src, bgr) # for training pha, fgr = model(src, bgr)[:2] # for inference """ def __init__(self, backbone: str): super().__init__(backbone, in_channels=6, out_channels=(1 + 3 + 1 + 32)) def forward(self, src, bgr): x = torch.cat([src, bgr], dim=1) x, *shortcuts = self.backbone(x) x = self.aspp(x) x = self.decoder(x, *shortcuts) pha = x[:, 0:1].clamp_(0., 1.) fgr = x[:, 1:4].add(src).clamp_(0., 1.) err = x[:, 4:5].clamp_(0., 1.) hid = x[:, 5: ].relu_() return pha, fgr, err, hid #网络结构中 Refine 网络 class MattingRefine(MattingBase): """ MattingRefine includes the refiner module to upsample coarse result to full resolution. MattingRefine extends MattingBase. Args: backbone: ["resnet50", "resnet101", "mobilenetv2"] backbone_scale: The image downsample scale for passing through backbone, default 1/4 or 0.25. Must not be greater than 1/2. 如 1920x1080 分辨率降低到 480x270. refine_mode: refine area selection mode. Options: "full" - No area selection, refine everywhere using regular Conv2d. "sampling" - Refine fixed amount of pixels ranked by the top most errors. "thresholding" - Refine varying amount of pixels that has more error than the threshold. refine_sample_pixels: number of pixels to refine. Only used when mode == "sampling". refine_threshold: error threshold ranged from 0 ~ 1. Refine where err > threshold. Only used when mode == "thresholding". refine_kernel_size: the refiner's convolutional kernel size. Options: [1, 3] refine_prevent_oversampling: prevent sampling more pixels than needed for sampling mode. Set False only for speedtest. Input: src: (B, 3, H, W) the source image. Channels are RGB values normalized to 0 ~ 1. bgr: (B, 3, H, W) the background image. Channels are RGB values normalized to 0 ~ 1. Output: pha: (B, 1, H, W) the alpha prediction. Normalized to 0 ~ 1. fgr: (B, 3, H, W) the foreground prediction. Channels are RGB values normalized to 0 ~ 1. pha_sm: (B, 1, Hc, Wc) the coarse alpha prediction from matting base. Normalized to 0 ~ 1. fgr_sm: (B, 3, Hc, Hc) the coarse foreground prediction from matting base. Normalized to 0 ~ 1. err_sm: (B, 1, Hc, Wc) the coarse error prediction from matting base. Normalized to 0 ~ 1. ref_sm: (B, 1, H/4, H/4) the quarter resolution refinement map. 1 indicates refined 4x4 patch locations. Example: model = MattingRefine(backbone='resnet50', backbone_scale=1/4, refine_mode='sampling', refine_sample_pixels=80_000) model = MattingRefine(backbone='resnet50', backbone_scale=1/4, refine_mode='thresholding', refine_threshold=0.1) model = MattingRefine(backbone='resnet50', backbone_scale=1/4, refine_mode='full') pha, fgr, pha_sm, fgr_sm, err_sm, ref_sm = model(src, bgr) # for training pha, fgr = model(src, bgr)[:2] # for inference """ def __init__(self, backbone: str, backbone_scale: float = 1/4, refine_mode: str = 'sampling', refine_sample_pixels: int = 80_000, refine_threshold: float = 0.1, refine_kernel_size: int = 3, refine_prevent_oversampling: bool = True, refine_patch_crop_method: str = 'unfold', refine_patch_replace_method: str = 'scatter_nd'): assert backbone_scale <= 1/2, 'backbone_scale should not be greater than 1/2' super().__init__(backbone) self.backbone_scale = backbone_scale self.refiner = Refiner(refine_mode, refine_sample_pixels, refine_threshold, refine_kernel_size, refine_prevent_oversampling, refine_patch_crop_method, refine_patch_replace_method) def forward(self, src, bgr): assert src.size() == bgr.size(), 'src and bgr must have the same shape' assert src.size(2) // 4 * 4 == src.size(2) and src.size(3) // 4 * 4 == src.size(3), \ 'src and bgr must have width and height that are divisible by 4' # Downsample src and bgr for backbone src_sm = F.interpolate(src, scale_factor=self.backbone_scale, mode='bilinear', align_corners=False, recompute_scale_factor=True) bgr_sm = F.interpolate(bgr, scale_factor=self.backbone_scale, mode='bilinear', align_corners=False, recompute_scale_factor=True) # Base x = torch.cat([src_sm, bgr_sm], dim=1) x, *shortcuts = self.backbone(x) x = self.aspp(x) x = self.decoder(x, *shortcuts) pha_sm = x[:, 0:1].clamp_(0., 1.) fgr_sm = x[:, 1:4] err_sm = x[:, 4:5].clamp_(0., 1.) hid_sm = x[:, 5: ].relu_() # Refiner pha, fgr, ref_sm = self.refiner(src, bgr, pha_sm, fgr_sm, err_sm, hid_sm) # Clamp outputs pha = pha.clamp_(0., 1.) fgr = fgr.add_(src).clamp_(0., 1.) fgr_sm = src_sm.add_(fgr_sm).clamp_(0., 1.) return pha, fgr, pha_sm, fgr_sm, err_sm, ref_sm2. Base 网络之 EncoderEncoder 支持三种：ResNet50、ResNet101、MobileNetV2.2.1. resnetresnet.pyfrom torch import nn from torchvision.models.resnet import ResNet, Bottleneck class ResNetEncoder(ResNet): """ ResNetEncoder inherits from torchvision's official ResNet. It is modified to use dilation on the last block to maintain output stride 16, and deleted the global average pooling layer and the fully connected layer that was originally used for classification. The forward method additionally returns the feature maps at all resolutions for decoder's use. """ layers = { 'resnet50': [3, 4, 6, 3], 'resnet101': [3, 4, 23, 3], } def __init__(self, in_channels, variant='resnet101', norm_layer=None): super().__init__( block=Bottleneck, layers=self.layers[variant], replace_stride_with_dilation=[False, False, True], norm_layer=norm_layer) # Replace first conv layer if in_channels doesn't match. if in_channels != 3: self.conv1 = nn.Conv2d(in_channels, 64, 7, 2, 3, bias=False) # Delete fully-connected layer del self.avgpool del self.fc def forward(self, x): x0 = x # 1/1 x = self.conv1(x) x = self.bn1(x) x = self.relu(x) x1 = x # 1/2 x = self.maxpool(x) x = self.layer1(x) x2 = x # 1/4 x = self.layer2(x) x3 = x # 1/8 x = self.layer3(x) x = self.layer4(x) x4 = x # 1/16 return x4, x3, x2, x1, x02.2. mobilenetmobilenet.pyfrom torch import nn from torchvision.models import MobileNetV2 class MobileNetV2Encoder(MobileNetV2): """ MobileNetV2Encoder inherits from torchvision's official MobileNetV2. It is modified to use dilation on the last block to maintain output stride 16, and deleted the classifier block that was originally used for classification. The forward method additionally returns the feature maps at all resolutions for decoder's use. """ def __init__(self, in_channels, norm_layer=None): super().__init__() # Replace first conv layer if in_channels doesn't match. if in_channels != 3: self.features[0][0] = nn.Conv2d(in_channels, 32, 3, 2, 1, bias=False) # Remove last block self.features = self.features[:-1] # Change to use dilation to maintain output stride = 16 self.features[14].conv[1][0].stride = (1, 1) for feature in self.features[15:]: feature.conv[1][0].dilation = (2, 2) feature.conv[1][0].padding = (2, 2) # Delete classifier del self.classifier def forward(self, x): x0 = x # 1/1 x = self.features[0](x) x = self.features[1](x) x1 = x # 1/2 x = self.features[2](x) x = self.features[3](x) x2 = x # 1/4 x = self.features[4](x) x = self.features[5](x) x = self.features[6](x) x3 = x # 1/8 x = self.features[7](x) x = self.features[8](x) x = self.features[9](x) x = self.features[10](x) x = self.features[11](x) x = self.features[12](x) x = self.features[13](x) x = self.features[14](x) x = self.features[15](x) x = self.features[16](x) x = self.features[17](x) x4 = x # 1/16 return x4, x3, x2, x1, x03. Base 网络之 Decoderdecoder.pyimport torch import torch.nn as nn import torch.nn.functional as F class Decoder(nn.Module): """ Decoder upsamples the image by combining the feature maps at all resolutions from the encoder. Input: x4: (B, C, H/16, W/16) feature map at 1/16 resolution. x3: (B, C, H/8, W/8) feature map at 1/8 resolution. x2: (B, C, H/4, W/4) feature map at 1/4 resolution. x1: (B, C, H/2, W/2) feature map at 1/2 resolution. x0: (B, C, H, W) feature map at full resolution. Output: x: (B, C, H, W) upsampled output at full resolution. """ def __init__(self, channels, feature_channels): super().__init__() self.conv1 = nn.Conv2d(feature_channels[0] + channels[0], channels[1], 3, padding=1, bias=False) self.bn1 = nn.BatchNorm2d(channels[1]) self.conv2 = nn.Conv2d(feature_channels[1] + channels[1], channels[2], 3, padding=1, bias=False) self.bn2 = nn.BatchNorm2d(channels[2]) self.conv3 = nn.Conv2d(feature_channels[2] + channels[2], channels[3], 3, padding=1, bias=False) self.bn3 = nn.BatchNorm2d(channels[3]) self.conv4 = nn.Conv2d(feature_channels[3] + channels[3], channels[4], 3, padding=1) self.relu = nn.ReLU(True) def forward(self, x4, x3, x2, x1, x0): x = F.interpolate(x4, size=x3.shape[2:], mode='bilinear', align_corners=False) x = torch.cat([x, x3], dim=1) x = self.conv1(x) x = self.bn1(x) x = self.relu(x) x = F.interpolate(x, size=x2.shape[2:], mode='bilinear', align_corners=False) x = torch.cat([x, x2], dim=1) x = self.conv2(x) x = self.bn2(x) x = self.relu(x) x = F.interpolate(x, size=x1.shape[2:], mode='bilinear', align_corners=False) x = torch.cat([x, x1], dim=1) x = self.conv3(x) x = self.bn3(x) x = self.relu(x) x = F.interpolate(x, size=x0.shape[2:], mode='bilinear', align_corners=False) x = torch.cat([x, x0], dim=1) x = self.conv4(x) return x4. Refine 网络refiner.pyimport torch import torchvision from torch import nn from torch.nn import functional as F from typing import Tuple class Refiner(nn.Module): """ Refiner refines the coarse output to full resolution. Args: mode: area selection mode. Options: "full" - No area selection, refine everywhere using regular Conv2d. "sampling" - Refine fixed amount of pixels ranked by the top most errors. "thresholding" - Refine varying amount of pixels that have greater error than the threshold. sample_pixels: number of pixels to refine. Only used when mode == "sampling". threshold: error threshold ranged from 0 ~ 1. Refine where err > threshold. Only used when mode == "thresholding". kernel_size: The convolution kernel_size. Options: [1, 3] prevent_oversampling: True for regular cases, False for speedtest. Compatibility Args: patch_crop_method: the operation for cropping patches. Options: "unfold" - Best performance for PyTorch and TorchScript. "roi_align" - Best compatibility for ONNX export. patch_replace_method: the operation for replacing patches. Options: "scatter_nd" - Best performance for PyTorch and TorchScript. "scatter_element" - Best compatibility for ONNX export. Input: src: (B, 3, H, W) full resolution source image. bgr: (B, 3, H, W) full resolution background image. pha: (B, 1, Hc, Wc) coarse alpha prediction. fgr: (B, 3, Hc, Wc) coarse foreground residual prediction. err: (B, 1, Hc, Hc) coarse error prediction. hid: (B, 32, Hc, Hc) coarse hidden encoding. Output: pha: (B, 1, H, W) full resolution alpha prediction. fgr: (B, 3, H, W) full resolution foreground residual prediction. ref: (B, 1, H/4, W/4) quarter resolution refinement selection map. 1 indicates refined 4x4 patch locations. """ # For TorchScript export optimization. __constants__ = ['kernel_size', 'patch_crop_method', 'patch_replace_method'] def __init__(self, mode: str, sample_pixels: int, threshold: float, kernel_size: int = 3, prevent_oversampling: bool = True, patch_crop_method: str = 'unfold', patch_replace_method: str = 'scatter_nd'): super().__init__() assert mode in ['full', 'sampling', 'thresholding'] assert kernel_size in [1, 3] assert patch_crop_method in ['unfold', 'roi_align'] assert patch_replace_method in ['scatter_nd', 'scatter_element'] self.mode = mode self.sample_pixels = sample_pixels self.threshold = threshold self.kernel_size = kernel_size self.prevent_oversampling = prevent_oversampling self.patch_crop_method = patch_crop_method self.patch_replace_method = patch_replace_method channels = [32, 24, 16, 12, 4] self.conv1 = nn.Conv2d(channels[0] + 6 + 4, channels[1], kernel_size, bias=False) self.bn1 = nn.BatchNorm2d(channels[1]) self.conv2 = nn.Conv2d(channels[1], channels[2], kernel_size, bias=False) self.bn2 = nn.BatchNorm2d(channels[2]) self.conv3 = nn.Conv2d(channels[2] + 6, channels[3], kernel_size, bias=False) self.bn3 = nn.BatchNorm2d(channels[3]) self.conv4 = nn.Conv2d(channels[3], channels[4], kernel_size, bias=True) self.relu = nn.ReLU(True) def forward(self, src: torch.Tensor, bgr: torch.Tensor, pha: torch.Tensor, fgr: torch.Tensor, err: torch.Tensor, hid: torch.Tensor): H_full, W_full = src.shape[2:] H_half, W_half = H_full // 2, W_full // 2 H_quat, W_quat = H_full // 4, W_full // 4 src_bgr = torch.cat([src, bgr], dim=1) if self.mode != 'full': err = F.interpolate(err, (H_quat, W_quat), mode='bilinear', align_corners=False) ref = self.select_refinement_regions(err) idx = torch.nonzero(ref.squeeze(1)) idx = idx[:, 0], idx[:, 1], idx[:, 2] if idx[0].size(0) > 0: x = torch.cat([hid, pha, fgr], dim=1) x = F.interpolate(x, (H_half, W_half), mode='bilinear', align_corners=False) x = self.crop_patch(x, idx, 2, 3 if self.kernel_size == 3 else 0) y = F.interpolate(src_bgr, (H_half, W_half), mode='bilinear', align_corners=False) y = self.crop_patch(y, idx, 2, 3 if self.kernel_size == 3 else 0) x = self.conv1(torch.cat([x, y], dim=1)) x = self.bn1(x) x = self.relu(x) x = self.conv2(x) x = self.bn2(x) x = self.relu(x) x = F.interpolate(x, 8 if self.kernel_size == 3 else 4, mode='nearest') y = self.crop_patch(src_bgr, idx, 4, 2 if self.kernel_size == 3 else 0) x = self.conv3(torch.cat([x, y], dim=1)) x = self.bn3(x) x = self.relu(x) x = self.conv4(x) out = torch.cat([pha, fgr], dim=1) out = F.interpolate(out, (H_full, W_full), mode='bilinear', align_corners=False) out = self.replace_patch(out, x, idx) pha = out[:, :1] fgr = out[:, 1:] else: pha = F.interpolate(pha, (H_full, W_full), mode='bilinear', align_corners=False) fgr = F.interpolate(fgr, (H_full, W_full), mode='bilinear', align_corners=False) else: x = torch.cat([hid, pha, fgr], dim=1) x = F.interpolate(x, (H_half, W_half), mode='bilinear', align_corners=False) y = F.interpolate(src_bgr, (H_half, W_half), mode='bilinear', align_corners=False) if self.kernel_size == 3: x = F.pad(x, (3, 3, 3, 3)) y = F.pad(y, (3, 3, 3, 3)) x = self.conv1(torch.cat([x, y], dim=1)) x = self.bn1(x) x = self.relu(x) x = self.conv2(x) x = self.bn2(x) x = self.relu(x) if self.kernel_size == 3: x = F.interpolate(x, (H_full + 4, W_full + 4)) y = F.pad(src_bgr, (2, 2, 2, 2)) else: x = F.interpolate(x, (H_full, W_full), mode='nearest') y = src_bgr x = self.conv3(torch.cat([x, y], dim=1)) x = self.bn3(x) x = self.relu(x) x = self.conv4(x) pha = x[:, :1] fgr = x[:, 1:] ref = torch.ones((src.size(0), 1, H_quat, W_quat), device=src.device, dtype=src.dtype) return pha, fgr, ref def crop_patch(self, x: torch.Tensor, idx: Tuple[torch.Tensor, torch.Tensor, torch.Tensor], size: int, padding: int): """ Crops selected patches from image given indices. Inputs: x: image (B, C, H, W). idx: selection indices Tuple[(P,), (P,), (P),], where the 3 values are (B, H, W) index. size: center size of the patch, also stride of the crop. padding: expansion size of the patch. Output: patch: (P, C, h, w), where h = w = size + 2 * padding. """ if padding != 0: x = F.pad(x, (padding,) * 4) if self.patch_crop_method == 'unfold': # Use unfold. Best performance for PyTorch and TorchScript. return x.permute(0, 2, 3, 1) \ .unfold(1, size + 2 * padding, size) \ .unfold(2, size + 2 * padding, size)[idx[0], idx[1], idx[2]] else: # Use roi_align. Best compatibility for ONNX. idx = idx[0].type_as(x), idx[1].type_as(x), idx[2].type_as(x) b = idx[0] x1 = idx[2] * size - 0.5 y1 = idx[1] * size - 0.5 x2 = idx[2] * size + size + 2 * padding - 0.5 y2 = idx[1] * size + size + 2 * padding - 0.5 boxes = torch.stack([b, x1, y1, x2, y2], dim=1) return torchvision.ops.roi_align(x, boxes, size + 2 * padding, sampling_ratio=1) def replace_patch(self, x: torch.Tensor, y: torch.Tensor, idx: Tuple[torch.Tensor, torch.Tensor, torch.Tensor]): """ Replaces patches back into image given index. Inputs: x: image (B, C, H, W) y: patches (P, C, h, w) idx: selection indices Tuple[(P,), (P,), (P,)] where the 3 values are (B, H, W) index. Output: image: (B, C, H, W), where patches at idx locations are replaced with y. """ xB, xC, xH, xW = x.shape yB, yC, yH, yW = y.shape if self.patch_replace_method == 'scatter_nd': # Use scatter_nd. Best performance for PyTorch and TorchScript. Replacing patch by patch. x = x.view(xB, xC, xH // yH, yH, xW // yW, yW).permute(0, 2, 4, 1, 3, 5) x[idx[0], idx[1], idx[2]] = y x = x.permute(0, 3, 1, 4, 2, 5).view(xB, xC, xH, xW) return x else: # Use scatter_element. Best compatibility for ONNX. Replacing pixel by pixel. iH, iW = xH // yH, xW // yW i = self.crop_patch(torch.arange(0, xB * xC * xH * xW).view(xB, xC, xH, xW).type_as(x), idx, 4, 0) i, x, y = i.view(-1), x.view(-1), y.view(-1) x.scatter_(0, i.long(), y) x = x.view(xB, xC, xH, xW) return x def select_refinement_regions(self, err: torch.Tensor): """ Select refinement regions. Input: err: error map (B, 1, H, W) Output: ref: refinement regions (B, 1, H, W). FloatTensor. 1 is selected, 0 is not. """ if self.mode == 'sampling': # Sampling mode. b, _, h, w = err.shape err = err.view(b, -1) idx = err.topk(self.sample_pixels // 16, dim=1, sorted=False).indices ref = torch.zeros_like(err) ref.scatter_(1, idx, 1.) if self.prevent_oversampling: ref.mul_(err.gt(0).float()) ref = ref.view(b, 1, h, w) else: # Thresholding mode. ref = err.gt(self.threshold).float() return ref5. load_matched_state_dictdef load_matched_state_dict(model, state_dict, print_stats=True): """ Only loads weights that matched in key and shape. Ignore other weights. 仅加载 key 和 shape 一致的权重参数,其余的忽略. """ num_matched, num_total = 0, 0 curr_state_dict = model.state_dict() for key in curr_state_dict.keys(): num_total += 1 if key in state_dict and curr_state_dict[key].shape == state_dict[key].shape: curr_state_dict[key] = state_dict[key] num_matched += 1 model.load_state_dict(curr_state_dict) if print_stats: print(f'Loaded state_dict: {num_matched}/{num_total} matched')
BackgroundMattingV2 提供了对图像、视频、网络摄像头的抠图实现，以及支持导出 ONNX 和 Torchscript.1. HomographicAlignment 图像对齐图像对齐函数，用于将背景与源图的匹配.inference_utils.pyimport numpy as np import cv2 from PIL import Image class HomographicAlignment: """ Apply homographic alignment on background to match with the source image. """ def __init__(self): self.detector = cv2.ORB_create() # ''' ORB 采用FAST(features from accelerated segment test)算法来检测特征点. FAST核心思想是, 找出那些卓尔不群的点，即拿一个点跟它周围的点比较， 如果它和其中大部分的点都不一样就可以认为它是一个特征点. ''' self.matcher = cv2.DescriptorMatcher_create(cv2.DESCRIPTOR_MATCHER_BRUTEFORCE) ''' 描述符匹配，特征匹配 ''' def __call__(self, src, bgr): src = np.asarray(src) bgr = np.asarray(bgr) keypoints_src, descriptors_src = self.detector.detectAndCompute(src, None) keypoints_bgr, descriptors_bgr = self.detector.detectAndCompute(bgr, None) matches = self.matcher.match(descriptors_bgr, descriptors_src, None) matches.sort(key=lambda x: x.distance, reverse=False) num_good_matches = int(len(matches) * 0.15) matches = matches[:num_good_matches] points_src = np.zeros((len(matches), 2), dtype=np.float32) points_bgr = np.zeros((len(matches), 2), dtype=np.float32) for i, match in enumerate(matches): points_src[i, :] = keypoints_src[match.trainIdx].pt points_bgr[i, :] = keypoints_bgr[match.queryIdx].pt # H, _ = cv2.findHomography(points_bgr, points_src, cv2.RANSAC) h, w = src.shape[:2] bgr = cv2.warpPerspective(bgr, H, (w, h)) #透视变换 msk = cv2.warpPerspective(np.ones((h, w)), H, (w, h)) # For areas that is outside of the background, # We just copy pixels from the source. bgr[msk != 1] = src[msk != 1] src = Image.fromarray(src) bgr = Image.fromarray(bgr) return src, bgr关于图像对齐可参考：[1] - Image Alignment (Feature Based) using OpenCV (C++/Python)[2] - 使用OpenCV实现基于特征的图像对齐2. 图像抠图inference_images.py对图像进行抠图(多线程).import argparse import torch import os import shutil from torch.nn import functional as F from torch.utils.data import DataLoader from torchvision import transforms as T from torchvision.transforms.functional import to_pil_image from threading import Thread from tqdm import tqdm from dataset import ImagesDataset, ZipDataset from dataset import augmentation as A from model import MattingBase, MattingRefine from inference_utils import HomographicAlignment # --------------- Arguments --------------- parser = argparse.ArgumentParser(description='Inference images') parser.add_argument('--model-type', type=str, required=True, choices=['mattingbase', 'mattingrefine']) parser.add_argument('--model-backbone', type=str, required=True, choices=['resnet101', 'resnet50', 'mobilenetv2']) parser.add_argument('--model-backbone-scale', type=float, default=0.25) parser.add_argument('--model-checkpoint', type=str, required=True) parser.add_argument('--model-refine-mode', type=str, default='sampling', choices=['full', 'sampling', 'thresholding']) parser.add_argument('--model-refine-sample-pixels', type=int, default=80_000) parser.add_argument('--model-refine-threshold', type=float, default=0.7) parser.add_argument('--model-refine-kernel-size', type=int, default=3) parser.add_argument('--images-src', type=str, required=True) parser.add_argument('--images-bgr', type=str, required=True) parser.add_argument('--device', type=str, choices=['cpu', 'cuda'], default='cuda') parser.add_argument('--preprocess-alignment', action='store_true') parser.add_argument('--output-dir', type=str, required=True) parser.add_argument('--output-types', type=str, required=True, nargs='+', choices=['com', 'pha', 'fgr', 'err', 'ref']) parser.add_argument('-y', action='store_true') args = parser.parse_args() assert 'err' not in args.output_types or args.model_type in ['mattingbase', 'mattingrefine'], \ 'Only mattingbase and mattingrefine support err output' assert 'ref' not in args.output_types or args.model_type in ['mattingrefine'], \ 'Only mattingrefine support ref output' # --------------- Main --------------- device = torch.device(args.device) #模型加载 if args.model_type == 'mattingbase': model = MattingBase(args.model_backbone) if args.model_type == 'mattingrefine': model = MattingRefine( args.model_backbone, args.model_backbone_scale, args.model_refine_mode, args.model_refine_sample_pixels, args.model_refine_threshold, args.model_refine_kernel_size) model = model.to(device).eval() model.load_state_dict(torch.load(args.model_checkpoint), strict=False) # 加载图像数据集 dataset = ZipDataset([ ImagesDataset(args.images_src), ImagesDataset(args.images_bgr),], assert_equal_length=True, transforms=A.PairCompose([HomographicAlignment() if args.preprocess_alignment else A.PairApply(nn.Identity()), A.PairApply(T.ToTensor()) ])) dataloader = DataLoader(dataset, batch_size=1, num_workers=8, pin_memory=True) # Create output directory if os.path.exists(args.output_dir): if args.y or input(f'Directory {args.output_dir} already exists. Override? [Y/N]: ').lower() == 'y': shutil.rmtree(args.output_dir) else: exit() for output_type in args.output_types: os.makedirs(os.path.join(args.output_dir, output_type)) # Worker function def writer(img, path): img = to_pil_image(img[0].cpu()) img.save(path) # Conversion loop with torch.no_grad(): for i, (src, bgr) in enumerate(tqdm(dataloader)): filename = dataset.datasets[0].filenames[i] src = src.to(device, non_blocking=True) bgr = bgr.to(device, non_blocking=True) if args.model_type == 'mattingbase': pha, fgr, err, _ = model(src, bgr) elif args.model_type == 'mattingrefine': pha, fgr, _, _, err, ref = model(src, bgr) elif args.model_type == 'mattingbm': pha, fgr = model(src, bgr) if 'com' in args.output_types: com = torch.cat([fgr * pha.ne(0), pha], dim=1) Thread(target=writer, args=(com, filename.replace(args.images_src, os.path.join(args.output_dir, 'com')).replace('.jpg', '.png'))).start() if 'pha' in args.output_types: Thread(target=writer, args=(pha, filename.replace(args.images_src, os.path.join(args.output_dir, 'pha')).replace('.png', '.jpg'))).start() if 'fgr' in args.output_types: Thread(target=writer, args=(fgr, filename.replace(args.images_src, os.path.join(args.output_dir, 'fgr')).replace('.png', '.jpg'))).start() if 'err' in args.output_types: err = F.interpolate(err, src.shape[2:], mode='bilinear', align_corners=False) Thread(target=writer, args=(err, filename.replace(args.images_src, os.path.join(args.output_dir, 'err')).replace('.png', '.jpg'))).start() if 'ref' in args.output_types: ref = F.interpolate(ref, src.shape[2:], mode='nearest') Thread(target=writer, args=(ref, filename.replace(args.images_src, os.path.join(args.output_dir, 'ref')).replace('.png', '.jpg'))).start()使用示例如：python inference_images.py \ --model-type mattingrefine \ --model-backbone resnet50 \ --model-backbone-scale 0.25 \ --model-refine-mode sampling \ --model-refine-sample-pixels 80000 \ --model-checkpoint "PATH_TO_CHECKPOINT" \ --images-src "PATH_TO_IMAGES_SRC_DIR" \ --images-bgr "PATH_TO_IMAGES_BGR_DIR" \ --output-dir "PATH_TO_OUTPUT_DIR" \ --output-type com fgr pha3. 视频抠图inference_video.py对视频进行抠图.import argparse import cv2 import torch import os import shutil from torch import nn from torch.nn import functional as F from torch.utils.data import DataLoader from torchvision import transforms as T from torchvision.transforms.functional import to_pil_image from threading import Thread from tqdm import tqdm from PIL import Image from dataset import VideoDataset, ZipDataset from dataset import augmentation as A from model import MattingBase, MattingRefine from inference_utils import HomographicAlignment # --------------- Arguments --------------- parser = argparse.ArgumentParser(description='Inference video') parser.add_argument('--model-type', type=str, required=True, choices=['mattingbase', 'mattingrefine']) parser.add_argument('--model-backbone', type=str, required=True, choices=['resnet101', 'resnet50', 'mobilenetv2']) parser.add_argument('--model-backbone-scale', type=float, default=0.25) parser.add_argument('--model-checkpoint', type=str, required=True) parser.add_argument('--model-refine-mode', type=str, default='sampling', choices=['full', 'sampling', 'thresholding']) parser.add_argument('--model-refine-sample-pixels', type=int, default=80_000) parser.add_argument('--model-refine-threshold', type=float, default=0.7) parser.add_argument('--model-refine-kernel-size', type=int, default=3) parser.add_argument('--video-src', type=str, required=True) parser.add_argument('--video-bgr', type=str, required=True) parser.add_argument('--video-resize', type=int, default=None, nargs=2) parser.add_argument('--device', type=str, choices=['cpu', 'cuda'], default='cuda') parser.add_argument('--preprocess-alignment', action='store_true') parser.add_argument('--output-dir', type=str, required=True) parser.add_argument('--output-types', type=str, required=True, nargs='+', choices=['com', 'pha', 'fgr', 'err', 'ref']) parser.add_argument('--output-format', type=str, default='video', choices=['video', 'image_sequences']) args = parser.parse_args() assert 'err' not in args.output_types or args.model_type in ['mattingbase', 'mattingrefine'], \ 'Only mattingbase and mattingrefine support err output' assert 'ref' not in args.output_types or args.model_type in ['mattingrefine'], \ 'Only mattingrefine support ref output' # --------------- Utils --------------- #视频写入类 class VideoWriter: def __init__(self, path, frame_rate, width, height): self.out = cv2.VideoWriter(path, cv2.VideoWriter_fourcc(*'mp4v'), frame_rate, (width, height)) def add_batch(self, frames): frames = frames.mul(255).byte() frames = frames.cpu().permute(0, 2, 3, 1).numpy() for i in range(frames.shape[0]): frame = frames[i] frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) self.out.write(frame) #图片序列写入类 class ImageSequenceWriter: def __init__(self, path, extension): self.path = path self.extension = extension self.index = 0 os.makedirs(path) def add_batch(self, frames): #多线程 Thread(target=self._add_batch, args=(frames, self.index)).start() self.index += frames.shape[0] def _add_batch(self, frames, index): frames = frames.cpu() for i in range(frames.shape[0]): frame = frames[i] frame = to_pil_image(frame) frame.save(os.path.join(self.path, str(index + i).zfill(5) + '.' + self.extension)) # --------------- Main --------------- device = torch.device(args.device) # Load model if args.model_type == 'mattingbase': model = MattingBase(args.model_backbone) if args.model_type == 'mattingrefine': model = MattingRefine( args.model_backbone, args.model_backbone_scale, args.model_refine_mode, args.model_refine_sample_pixels, args.model_refine_threshold, args.model_refine_kernel_size) model = model.to(device).eval() model.load_state_dict(torch.load(args.model_checkpoint), strict=False) #视频数据集和背景图像加载 vid = VideoDataset(args.video_src) bgr = [Image.open(args.video_bgr).convert('RGB')] dataset = ZipDataset( [vid, bgr], transforms=A.PairCompose([ A.PairApply(T.Resize(args.video_resize[::-1]) if args.video_resize else nn.Identity()), HomographicAlignment() if args.preprocess_alignment else A.PairApply(nn.Identity()), A.PairApply(T.ToTensor()) ])) # Create output directory if os.path.exists(args.output_dir): if input(f'Directory {args.output_dir} already exists. Override? [Y/N]: ').lower() == 'y': shutil.rmtree(args.output_dir) else: exit() os.makedirs(args.output_dir) # Prepare writers if args.output_format == 'video': h = args.video_resize[1] if args.video_resize is not None else vid.height w = args.video_resize[0] if args.video_resize is not None else vid.width if 'com' in args.output_types: com_writer = VideoWriter(os.path.join(args.output_dir, 'com.mp4'), vid.frame_rate, w, h) if 'pha' in args.output_types: pha_writer = VideoWriter(os.path.join(args.output_dir, 'pha.mp4'), vid.frame_rate, w, h) if 'fgr' in args.output_types: fgr_writer = VideoWriter(os.path.join(args.output_dir, 'fgr.mp4'), vid.frame_rate, w, h) if 'err' in args.output_types: err_writer = VideoWriter(os.path.join(args.output_dir, 'err.mp4'), vid.frame_rate, w, h) if 'ref' in args.output_types: ref_writer = VideoWriter(os.path.join(args.output_dir, 'ref.mp4'), vid.frame_rate, w, h) else: if 'com' in args.output_types: com_writer = ImageSequenceWriter(os.path.join(args.output_dir, 'com'), 'png') if 'pha' in args.output_types: pha_writer = ImageSequenceWriter(os.path.join(args.output_dir, 'pha'), 'jpg') if 'fgr' in args.output_types: fgr_writer = ImageSequenceWriter(os.path.join(args.output_dir, 'fgr'), 'jpg') if 'err' in args.output_types: err_writer = ImageSequenceWriter(os.path.join(args.output_dir, 'err'), 'jpg') if 'ref' in args.output_types: ref_writer = ImageSequenceWriter(os.path.join(args.output_dir, 'ref'), 'jpg') # Conversion loop with torch.no_grad(): for src, bgr in tqdm(DataLoader(dataset, batch_size=1, pin_memory=True)): src = src.to(device, non_blocking=True) bgr = bgr.to(device, non_blocking=True) if args.model_type == 'mattingbase': pha, fgr, err, _ = model(src, bgr) elif args.model_type == 'mattingrefine': pha, fgr, _, _, err, ref = model(src, bgr) elif args.model_type == 'mattingbm': pha, fgr = model(src, bgr) if 'com' in args.output_types: if args.output_format == 'video': #加绿幕背景输出 Output composite with green background bgr_green = torch.tensor([120/255, 255/255, 155/255], device=device).view(1, 3, 1, 1) com = fgr * pha + bgr_green * (1 - pha) com_writer.add_batch(com) else: #合成rgba图像输出 Output composite as rgba png images com = torch.cat([fgr * pha.ne(0), pha], dim=1) com_writer.add_batch(com) if 'pha' in args.output_types: pha_writer.add_batch(pha) if 'fgr' in args.output_types: fgr_writer.add_batch(fgr) if 'err' in args.output_types: err_writer.add_batch(F.interpolate(err, src.shape[2:], mode='bilinear', align_corners=False)) if 'ref' in args.output_types: ref_writer.add_batch(F.interpolate(ref, src.shape[2:], mode='nearest'))使用示例如：python inference_video.py \ --model-type mattingrefine \ --model-backbone resnet50 \ --model-backbone-scale 0.25 \ --model-refine-mode sampling \ --model-refine-sample-pixels 80000 \ --model-checkpoint "PATH_TO_CHECKPOINT" \ --video-src "PATH_TO_VIDEO_SRC" \ --video-bgr "PATH_TO_VIDEO_BGR" \ --video-resize 1920 1080 \ --output-dir "PATH_TO_OUTPUT_DIR" \ --output-type com fgr pha err ref4. 网络摄像头抠图inference_webcam.py对于网络摄像头输入来使用模型进行抠图.作者提供的实现中，当网络摄像头启动时，脚本会处在背景采集模式; 按 B 键在背景采集模式和抠图模式间进行切换. 当 B 键按下时的视频帧会作为背景进行抠图.按 Q 键推出脚本程序.import argparse, os, shutil, time import cv2 import torch from torch import nn from torch.utils.data import DataLoader from torchvision.transforms import Compose, ToTensor, Resize from torchvision.transforms.functional import to_pil_image from threading import Thread, Lock from tqdm import tqdm from PIL import Image from dataset import VideoDataset from model import MattingBase, MattingRefine # --------------- Arguments --------------- parser = argparse.ArgumentParser(description='Inference from web-cam') parser.add_argument('--model-type', type=str, required=True, choices=['mattingbase', 'mattingrefine']) parser.add_argument('--model-backbone', type=str, required=True, choices=['resnet101', 'resnet50', 'mobilenetv2']) parser.add_argument('--model-backbone-scale', type=float, default=0.25) parser.add_argument('--model-checkpoint', type=str, required=True) parser.add_argument('--model-refine-mode', type=str, default='sampling', choices=['full', 'sampling', 'thresholding']) parser.add_argument('--model-refine-sample-pixels', type=int, default=80_000) parser.add_argument('--model-refine-threshold', type=float, default=0.7) parser.add_argument('--hide-fps', action='store_true') parser.add_argument('--resolution', type=int, nargs=2, metavar=('width', 'height'), default=(1280, 720)) args = parser.parse_args() # ----------- Utility classes ------------- # A wrapper that reads data from cv2.VideoCapture in its own thread to optimize. # Use .read() in a tight loop to get the newest frame class Camera: def __init__(self, device_id=0, width=1280, height=720): self.capture = cv2.VideoCapture(device_id) self.capture.set(cv2.CAP_PROP_FRAME_WIDTH, width) self.capture.set(cv2.CAP_PROP_FRAME_HEIGHT, height) self.width = int(self.capture.get(cv2.CAP_PROP_FRAME_WIDTH)) self.height = int(self.capture.get(cv2.CAP_PROP_FRAME_HEIGHT)) # self.capture.set(cv2.CAP_PROP_BUFFERSIZE, 2) self.success_reading, self.frame = self.capture.read() self.read_lock = Lock() self.thread = Thread(target=self.__update, args=()) self.thread.daemon = True self.thread.start() def __update(self): while self.success_reading: grabbed, frame = self.capture.read() with self.read_lock: self.success_reading = grabbed self.frame = frame def read(self): with self.read_lock: frame = self.frame.copy() return frame def __exit__(self, exec_type, exc_value, traceback): self.capture.release() # An FPS tracker that computes exponentialy moving average FPS class FPSTracker: def __init__(self, ratio=0.5): self._last_tick = None self._avg_fps = None self.ratio = ratio def tick(self): if self._last_tick is None: self._last_tick = time.time() return None t_new = time.time() fps_sample = 1.0 / (t_new - self._last_tick) self._avg_fps = self.ratio * fps_sample + (1 - self.ratio) * self._avg_fps if self._avg_fps is not None else fps_sample self._last_tick = t_new return self.get() def get(self): return self._avg_fps # Wrapper for playing a stream with cv2.imshow(). 视频流播放的封装函数 # It can accept an image and return keypress info for basic interactivity. # It also tracks FPS and optionally overlays info onto the stream. class Displayer: def __init__(self, title, width=None, height=None, show_info=True): self.title, self.width, self.height = title, width, height self.show_info = show_info self.fps_tracker = FPSTracker() cv2.namedWindow(self.title, cv2.WINDOW_NORMAL) if width is not None and height is not None: cv2.resizeWindow(self.title, width, height) # Update the currently showing frame and return key press char code def step(self, image): fps_estimate = self.fps_tracker.tick() if self.show_info and fps_estimate is not None: message = f"{int(fps_estimate)} fps | {self.width}x{self.height}" cv2.putText(image, message, (10, 40), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 0)) cv2.imshow(self.title, image) return cv2.waitKey(1) & 0xFF # --------------- Main --------------- # Load model if args.model_type == 'mattingbase': model = MattingBase(args.model_backbone) if args.model_type == 'mattingrefine': model = MattingRefine( args.model_backbone, args.model_backbone_scale, args.model_refine_mode, args.model_refine_sample_pixels, args.model_refine_threshold) model = model.cuda().eval() model.load_state_dict(torch.load(args.model_checkpoint), strict=False) width, height = args.resolution cam = Camera(width=width, height=height) dsp = Displayer('MattingV2', cam.width, cam.height, show_info=(not args.hide_fps)) def cv2_frame_to_cuda(frame): frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) return ToTensor()(Image.fromarray(frame)).unsqueeze_(0).cuda() with torch.no_grad(): while True: bgr = None while True: # grab bgr frame = cam.read() key = dsp.step(frame) if key == ord('b'): bgr = cv2_frame_to_cuda(cam.read()) break elif key == ord('q'): exit() while True: # matting frame = cam.read() src = cv2_frame_to_cuda(frame) pha, fgr = model(src, bgr)[:2] res = pha * fgr + (1 - pha) * torch.ones_like(fgr) res = res.mul(255).byte().cpu().permute(0, 2, 3, 1).numpy()[0] res = cv2.cvtColor(res, cv2.COLOR_RGB2BGR) key = dsp.step(res) if key == ord('b'): break elif key == ord('q'): exit()使用示例如：python inference_webcam.py \ --model-type mattingrefine \ --model-backbone resnet50 \ --model-checkpoint "PATH_TO_CHECKPOINT" \ --resolution 1280 7205. 抠图速度测试inference_speed_test.py测试抠图的速度，主要包括两种：[1] - 采用随机噪声输入，进行固定计算测试.[2] - 采用提供的图像输入，进行动态计算测试.import argparse import torch from torchvision.transforms.functional import to_tensor from tqdm import tqdm from PIL import Image from model import MattingBase, MattingRefine # --------------- Arguments --------------- parser = argparse.ArgumentParser() parser.add_argument('--model-type', type=str, required=True, choices=['mattingbase', 'mattingrefine']) parser.add_argument('--model-backbone', type=str, required=True, choices=['resnet101', 'resnet50', 'mobilenetv2']) parser.add_argument('--model-backbone-scale', type=float, default=0.25) parser.add_argument('--model-checkpoint', type=str, default=None) parser.add_argument('--model-refine-mode', type=str, default='sampling', choices=['full', 'sampling', 'thresholding']) parser.add_argument('--model-refine-sample-pixels', type=int, default=80_000) parser.add_argument('--model-refine-threshold', type=float, default=0.7) parser.add_argument('--model-refine-kernel-size', type=int, default=3) parser.add_argument('--batch-size', type=int, default=1) parser.add_argument('--resolution', type=int, default=None, nargs=2) parser.add_argument('--precision', type=str, default='float32', choices=['float32', 'float16']) parser.add_argument('--backend', type=str, default='pytorch', choices=['pytorch', 'torchscript']) parser.add_argument('--device', type=str, choices=['cpu', 'cuda'], default='cuda') parser.add_argument('--image-src', type=str, default=None) parser.add_argument('--image-bgr', type=str, default=None) args = parser.parse_args() assert type(args.image_src) == type(args.image_bgr), 'Image source and background must be provided together.' assert (not args.image_src) != (not args.resolution), 'Must provide either a resolution or an image and not both.' # --------------- Run Loop --------------- device = torch.device(args.device) # Load model if args.model_type == 'mattingbase': model = MattingBase(args.model_backbone) if args.model_type == 'mattingrefine': model = MattingRefine( args.model_backbone, args.model_backbone_scale, args.model_refine_mode, args.model_refine_sample_pixels, args.model_refine_threshold, args.model_refine_kernel_size, refine_prevent_oversampling=False) if args.model_checkpoint: model.load_state_dict(torch.load(args.model_checkpoint), strict=False) if args.precision == 'float32': precision = torch.float32 else: precision = torch.float16 if args.backend == 'torchscript': model = torch.jit.script(model) model = model.eval().to(device=device, dtype=precision) # Load data if not args.image_src: src = torch.rand((args.batch_size, 3, *args.resolution[::-1]), device=device, dtype=precision) bgr = torch.rand((args.batch_size, 3, *args.resolution[::-1]), device=device, dtype=precision) else: src = to_tensor(Image.open(args.image_src)).unsqueeze(0).repeat(args.batch_size, 1, 1, 1).to(device=device, dtype=precision) bgr = to_tensor(Image.open(args.image_bgr)).unsqueeze(0).repeat(args.batch_size, 1, 1, 1).to(device=device, dtype=precision) # Loop with torch.no_grad(): for _ in tqdm(range(1000)): model(src, bgr)使用示例如：#Run inference on random noise input for fixed computation setting. #(i.e. mode in ['full', 'sampling']) python inference_speed_test.py \ --model-type mattingrefine \ --model-backbone resnet50 \ --model-backbone-scale 0.25 \ --model-refine-mode sampling \ --model-refine-sample-pixels 80000 \ --batch-size 1 \ --resolution 1920 1080 \ --backend pytorch \ --precision float32 #Run inference on provided image input for dynamic computation setting. #(i.e. mode in ['thresholding']) python inference_speed_test.py \ --model-type mattingrefine \ --model-backbone resnet50 \ --model-backbone-scale 0.25 \ --model-checkpoint "PATH_TO_CHECKPOINT" \ --model-refine-mode thresholding \ --model-refine-threshold 0.7 \ --batch-size 1 \ --backend pytorch \ --precision float32 \ --image-src "PATH_TO_IMAGE_SRC" \ --image-bgr "PATH_TO_IMAGE_BGR"
论文阅读 - Real-Time High-Resolution Background Matting
论文：Real-Time High-Resolution Background Matting作者：Shanchuan Lin Andrey Ryabtsev Soumyadip Sengupta团队：cs.washington.eduHomepage - Real-Time High Resolution Background MattingGithub - PeterL1n/BackgroundMattingV2针对视频背景抠图，该论文的几个着重点：[1] - 实时且高分辨率, 在 GPU 上，4K 分辨率(3840x2160) 30fps; HD 分辨率(1920x1080) 60 fps.[2] - 基于背景抠图(background matting)，其需要单独捕捉一帧背景图，以用于重建 alpha matte 和前景层.[3] - 最关键的挑战是，高分辨率 alpha matte 的计算，在实时处理高分辨率图像的同时，保持发丝级细节.[4] - 采用两个神经网络，base 网络计算低分辨率结果; refine 网络对高分辨率图像所选择的图片块( selective patches) 上进一步处理.[5] - 两个视频和图像抠图数据集：VideoMatte240K 和 PhotoMatte 13K/85.1. 抠图数据集获取大量、高分辨率、高质量的、由人工标注的抠图数据集是非常难的.1.1. 公开数据集Adobe Deep Image Matting 数据集提供了 269 张人物训练样本和 11 张测试图像，分辨率平均为 1000x1000.Attention-Guided Hierarchical Structure Aggregation for Image Matting 中362 张人物相关的训练样本和 11 张测试样本，分辨率平均为 1700x2000.以上两个数据集共 631 张训练样本，是远远不够的.1.2. VideoMatte240KVideoMatte240K 中采集了 484 段高分辨率绿幕(green screen) 视频，并采用色度键抠图(Chroma-key)软件 Adobe After Effects 生成了 240709 张 alpha mattes 和 foregrounds 的视频帧.其中，384 段视频的分辨率为 4K, 100 段视频的分辨率为 HD. 根据 479:5 的比例将视频划分为 tran 和 validataion 数据集. 如图：1.3. PhotoMatte13K/85PhotoMatte13K/85 中采集了 13665 张图像集合，这些图像是采用摄影棚质量的光照，在绿幕前的相机所拍摄的，并采用色度键抠图算法(Chroma-key) 和手工调整以及错误修正所提取的 mattes.根据 13165:500 的比例划分为 train 和 validation 数据集. 数据集图像分辨率平均为 2000x2500，包括了一些头发.2. 算法细节给定图像 I 以及捕捉的背景 B, 算法会预测输出 alpha matte $\alpha$ 和前景 F, 基于如下公式，其可以被任何背景来组合，$$ I' = \alpha F + (1 - \alpha) B' $$其中，$B'$ 为新背景.相比于直接求解前景，这里求解的是前景残差 $F^{R} = F - I$. 因此，F 可以通过添加适当的约束将 $F^R$ 和输入图像 I 相加:$$ F = max(min( F^{R} + I, 1), 0) $$高分辨率抠图是很有挑战的，因为直接进行深度网络计算，其计算量和内存消耗是不切实际的. 如图所示：人物 mattes 往往是非常稀疏的，大面积的像素要么是背景($\alpha=0$)，要么是前景($\alpha = 1$)，仅有少量的区域包含更多的细节，如头发、眼镜、人物轮廓等周围.对此，相比于设计直接在高分辨率图像直接操作的一个网络，这里采用了两个网络：[1] - base 网络 $G_{base}$ - 对低分辨率图像进行处理;[2] - refinement 网络 $G_{refine}$ - 基于低分辨率预测结果，从原始图像分辨率选择特定图像块进行处理.网络结构如图：给定原始输入图像 I 和捕捉的背景图 B，首先采用缩放因子 $c$ 进行将采样到 $I_c$ 和 $B_c$. $G_{base}$ 采用 $I_c$ 和 $B_c$ 作为输入，并预测输出粗略的 alpha matte $\alpha_c$、foreground residual $F_c^R$、error prediction map $E_c$ 以及 hidden feature $H_c$. $G_{refine}$ 采用 $H_c$、I 和 B 仅对 $E_c$ 比较大的区域，进一步的精细化 $\alpha_c$ 和 $F_c^R$，并在原始分辨率输出 alpha $\alpha$ 和 $F^R$.网络模型是全卷积的，且能在任意尺寸和长宽比(aspect ratios) 上进行训练.2.1. Base 网络Base 网络受到 DeepLabV3 和 DeepLabV3+ 结构的启发，设计为全卷积的 encoder-decoder 网络，其中主要包括三个模块：Backbone， ASPP，Decoder.[1] - Backbone采用 ResNet50 作为 encoder backbone，其也可以采用 ResNet101 和 MobileNetV2，以平衡速度和精度.[2] - ASPP类似于 DeepLabV3，在 backbone 网络后接 ASPP(Atrous Spatial Pyramid Pooling) 模块. ASPP 模块由多个不同 dilation rates 为 3, 6, 9 的 dialted conv filters 组成.[3] - Decoder Decoder 网络的每一步采用双线性上采样(bilinear upsampling)，并通过 skip connection 从 backbone 进行连接，其后接一个 3x3 conv, Batch Normalization 和 ReLU(最后一层例外).Decoder 网络输出粗略的 alpha matte $\alpha _c$，foreground residual $F_c^R$，error prediction map $E_c$ 和 32 通道 hidden features $H_c$. $H_c$ 包含了有用于 refinement 网络的全局内容信息.2.2. Refinement 网络Refinement 网络的目标是，降低冗余计算量，并重建高分辨率抠图细节.Base 网络是对整张图像进行处理，而 Refinement 网络仅对基于 error prediction map $E_c$ 所选择的图像块进行处理.Refinement 网络是 two-stage 的，分别处理原始分辨率的二分之一和全分辨率.在推断时，精细化 k 个图像块，k 是提前设定的，或着是根据抠图质量和计算量平衡的阈值计算来设定的.给定在原始分辨率的 $\frac{1}{c}$ 时的 error prediction map $E_c$，首先将其重采样为原始分辨率的 $\frac{1}{4}$，记 $E_4$，其每个像素对应于原始图像分辨率的一个 4x4 图像块. 然后，从 $E_4$ 中选取前 k 个最高的 predicted error 的像素，以表示 k 个 4x4 图像块位置，用于 refinement 模块进一步精细化处理. 原始分辨率的 refined 像素总数是 16k.Refinement 是 two-stage 的，[1] - 首先，双线性重采样粗略的输出，如 alpha matte $\alpha_c$、foreground residual $F_c^R$、hidden features $H_c$、输入图像 I和背景图像的分辨率重采样为原始分辨率的 $\frac{1}{2}$，并将这些链接作为特征.[2] - 然后，在 $E_4$ 中选择出的 error 位置的周围裁剪出多个 8x8 图像块，并将其分别送入两个 valid padding 的 3x3 conv 层、BN 层和 ReLU 层，以将图像块的分辨率降低为 4x4. 接着，这些中间特征再一次被上采样到 8x8，并与从原始分辨率输入 I 和 B 中所提取的 8x8 图像块拼接在一起.其后，再采用一个两个 valid padding 的 3x3 conv 层、BN 层和 ReLU 层(最后一层除外)处理，得到 4x4 alpha matte 和 foreground residuals 结果.最后，上采样粗略的 alpha matte $\alpha_c$、foreground residual $F_c^R$ 到原始分辨率，并交换被 refine 后的 4x4 图像块，以得到最终的 alpha matte $\alpha$ 和 foreground residual $F^R$.2.3. Loss 函数关于 GT $\alpha^*$ 和 $\alpha$ ，采用 L1 loss 计算整个 alpha matte 和对应的 Sobel 梯度损失函数：$$ L_{\alpha} = ||\alpha - \alpha^*||_1 + ||\nabla \alpha - \nabla \alpha^*||_1 $$根据预测的 foreground residual $F^R$ 以及如下计算公式，得到前景层：$$ F = max(min(F^R + I, 1), 0) $$仅对 $\alpha^* > 0$ 的像素区域计算 L1 loss:$$ L_F = ||(\alpha^* > 0) * (F - F^*)||_1 $$其中，$\alpha^* > 0$ 为布尔表达式.对于 refinement 区域选择，首先定义 ground truth error map 为：$$ E^* = |\alpha - \alpha^*| $$然后，计算预测的 error map 和 ground truth error map 间的均方误差：$$ L_E = ||E - E^*||_2 $$Base 网络:$$ (\alpha_c, F_c^R, E_c, H_c) = G_{base}(I_c, B_c) $$对原始图像分辨率的 $\frac{1}{c}$ 进行处理，其训练的损失函数如下：$$ L_{base} = L_{\alpha_c} + L_{F_C} + L_{E_c} $$Refinement 网络：$$ (\alpha, F^R) = G_{refine}(\alpha_c, F_{c}^R, E_c, H_c, I, B) $$其训练的损失函数如下：$$ L_{refine} = L_{\alpha} + L_{F} $$2.4. 训练抠图数据集提供了 alpha matte 和 foreground 层，其可以组合到多种高分辨率背景上.数据增强：仿射变换(affine transformation)、水平翻转(horizontal flipping)、光照色调饱和度调整(brightness, hue, and saturation adjustment)、模糊(blurring)、锐化(sharpening)、随机噪声(random noise) 等，分别对前景层和背景层进行处理.此外，轻微平移背景以模拟不对齐; 生成阴影效果，以模拟真实环境的效果.如：每个 minibatch 随机裁剪图像，以使得 height 和 width 是在 [1024, 2048] 区间的均匀分布，便于支持推断时支持任意分辨率和长宽比.模型训练次序:[1] - 首先，仅训练 base 网络 $G_{base}$; 在数据集 VideoMatte240K 上，联合训练整个模型 $G_{base}$ 和 $G_{refine}$ ;[2] - 然后，在数据集 PhotoMatte 13K 联合训练模型，以提升高分辨率细节.[3] - 最后，在公开数据集(646 张训练数据)上联合训练模型.3. Results抠图效果对比如图：抠图速率对比：(NVIDIA RTX 2080Ti GPU)模型参数及大小对比：
原文：Pytorch 抠图算法 Deep Image Matting 模型实现 - 2019.07.21GitHub: deep_image_matting_pytorch作者：公输睚信 - 简书本文旨在实现抠图算法 Semantic Human Matting 的第二阶段模型 M-Net，也即 Deep Image Matting. 值得说明的是，本文实现的模型与原始论文略有出入，除了模型的输入层有细微差别之外，损失函数也作了简化(但无本质差别). 本文训练数据来源于爱分割公司开源的 aisegmentcn/matting_human_datasets，总共包含 34426 张图片和对应的 alpha 通道，数据量非常大，能公开特别值得点赞. 但同时因为标注数据的 alpha 通道精度不高，导致训练后测试效果较差. 建议使用 Deep Image Matting 的数据集训练. 图：Semantic Human Matting 抠图模型总的来说，Semantic Human Matting 论文提出的自动抠图的思路特别清晰明了(如上图)，对于一张待抠图像，首先通过语义分割模型(即 T-Net)分割出前景 Fs、背景 Bs 和未知区域 Us，Fs+Bs+Us=1. 广义的认为前景Fs和未知区域Us组成一个三分图Trimap. 此时，再利用 Deep Image Matting(即 M-Net) 即可高质量的完成抠图. 本文只关注该模型的第二阶段(M-Net). M-Net 接受待抠图像(前景与背景的 RGB 3通道合成)以及语义分割模型输出的3通道预测(Fs，Bs，Us) 拼接而成的 6 通道输入，经过编码器提取图像特征之后，由解码器得到预测Alpha($\alpha _r$). 如果语义分割模型分割的精度较高，那么可以认为 Fs 和 Bs 对应的区域已经很好的抠出了大部分的前景和背景，唯一需要提升准确率的是待抠对象的边缘区域，所以模型的第二阶段 M-Net 的目的就是细化的预测边缘区域(这正是 Deep Image Matting 要干的事情)，两部分结合即得到最终的预测：$$ F_s + U_s \alpha _r $$公式可以这样理解：$$ 预测的前景 = 确定区域上的前景 + 未知区域上的前景 $$即：$$ P(前景) = P(确定区域上前景) + P(未知区域上前景) $$根据全概率公式，符号化表达为：$$ \alpha = P(F) = P(F|known) P(known) + P(F|unknown) P(unknown) $$$$ \alpha = \frac{F_s}{F_s + B_s}(F_s + B_s) + \alpha _r U_s = F_s + U_s \alpha _r $$但上述公式存在一个缺陷，即如果待抠目标外有大块噪声，则最终的预测也消除不了这个噪声，如下图：图：语义分割之后的前景带有外部噪声(衣服左侧的小照片)为了消除第一阶段可能包含的外部噪声，本文的在实现 M-Net 的时候做了一个小的改动：第二阶段的输入改为由待抠图像 + Fs 组成的 4通道图片(此时，相当于将 Fs 看成是三分图 trimap)，并且将第二阶段的预测作为最终的预测. 另外，第二阶段损失函数简化为只用 alpha 通道的损失更好.1. DIM 模型实现Deep Image Matting 原文的模型如下：图：Deep Image Matting 论文模型模型先通过一个编码器提取特征，之后经过一个解码器预测一个初始的 alpha 通道，这个预测值效果已经很好，但作者为了进一步提升抠图的精度，又额外的接了几层细化的小网络，然后将细化后的输出作为整个模型的最终输出. 具体来说，首先将待抠图像(3 通道)以及事先准备好的三分图(trimap)合成一个 4 通道图像，然后经过 VGG16 的前 13 个卷积层以及之后的 1 个全连接层(看成是 1x1 的卷积层)，总共 14 个卷积层提取图像特征(此时已做了 5 个最大池化，因此图像分辨率下降了 32 倍，如果输入是 320x320，那么特征映射的分辨率就变成了 10x10)，这是模型的编码器阶段. 接下来对图像特征进行解码，即开始解码器阶段. 解码器使用 6 个卷积层(5x5 的卷积核)和 5 个 反池化层，每个反池化层将特征映射的分辨率提升 2 倍，因此解码器的输出与模型输入的大小一样. 这里，使用反池化层的效果要比直接使用转置卷积(deconvolution)的效果要好. 虽然他们都是为了提升图像分辨率，但使用转置卷积并不能很好的抠出细节，而使用反池化层却可以抠图头发丝等非常细的前景. 为了最求极致效果，作者又接了一个小网络，将待抠图像和编码器预测的 alpha 通道合成一个 4 通道图像，然后通过 4 个 3x3 的卷积层得到细化后的 alpha 通道预测，作为最后的输出. 图：Deep Image Matting 效果图：(a) 原图；(b) 编码器-解码器阶段结果；(c)细化阶段结果2. DIM 损失函数损失方面，总共用了 3 个分损失来合成网络的损失：[1] - 编码器阶段预测的 alpha 通道和真实的 alpha 通道的损失；[2] - 编码器阶段使用预测的 alpha 合成的图像和真实的 alpha 合成的图像的损失；[3] - 细化阶段预测的 alpha 通道和真实的 alpha 通道的损失. 这些损失都是逐点损失，即平方和误差：[1] - 预测与真实 alpha 通道之间的损失[2] - 由前景、背景和 alpha 通道合成图像之间的损失三个损失使用加权和形成整个网络最后反向传播的总损失.3. SHM M-NetDeep Image Matting 虽然论文上报告的效果很惊人，但实际实现时(在个人应用数据集上)泛化性能不够理想. Semantic Human Matting(SHM)这篇论文的 M-Net 在以上基础上做了一些简化和修改. 首先，为了防止网络容量太大造成过拟合，SHM 只使用 VGG16 的前 13 个卷积层及 4 个最大池化层来作为编码器，相应的，解码器阶段也就少一个反池化层. 另外，为了加速网络收敛，所有的卷积层(编码器以及解码器的)都带批标准化(Batch Normalization)处理. 其次，网络的输入由 4 通道变成了 6 通道，这样做一方面没有影响网络性能(论文 4.2 节)，另一方面也是为了方便与 T-Net 对接，因为 T-Net 输出前景、背景、未知 3 个预测通道，与待抠图像的 3 通道直接合成即得到 6 通道输入. 最后，SHM 直接去掉了 Deep Image Matting 网络的细化阶段，因此损失也相应的减少为 2 个分损失. 本文基本忠实的实现了 SHM 的 M-Net 结构，但如本文开始时候说的那样，将 6 通道的输入改成了 4 通道，且为了完全引入 VGG16 的预训练模型，直接在 VGG16 的最前面接了一个输入为 6 通道、输出为 4 通道的卷积层. 此外，本文将 M-Net 的预测作为最终的输出，以及训练时不再求合成图像的损失(以下模型实现时，loss 函数是支持合成图像损失的).4. Python 模型定义总的来说，Deep Image Matting (或 M-Net)网络是非常清晰明了的，实现也很简单，模型文件 model.py 如下：# -*- coding: utf-8 -*- """ Created on Sun Jul 21 07:08:58 2019 @author: shirhe-lyh Implementation of paper: Deep Image Matting, Ning Xu, eta., arxiv:1703.03872 """ import torch import torchvision as tv VGG16_BN_MODEL_URL = 'https://download.pytorch.org/models/vgg16_bn-6c64b313.pth' VGG16_BN_CONFIGS = { '13conv': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512], '10conv': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512] } def make_layers(cfg, batch_norm=False): """ Copy from: torchvision/models/vgg. Changs retrue_indices in MaxPool2d from False to True. """ layers = [] in_channels = 3 for v in cfg: if v == 'M': layers += [torch.nn.MaxPool2d(kernel_size=2, stride=2, return_indices=True)] else: conv2d = torch.nn.Conv2d(in_channels, v, kernel_size=3, padding=1) if batch_norm: layers += [conv2d, torch.nn.BatchNorm2d(v), torch.nn.ReLU(inplace=True)] else: layers += [conv2d, torch.nn.ReLU(inplace=True)] in_channels = v return torch.nn.Sequential(*layers) class VGGFeatureExtractor(torch.nn.Module): """Feature extractor by VGG network.""" def __init__(self, config=None, batch_norm=True): """ Constructor. Args: config: The convolutional architecture of VGG network. batch_norm: A boolean indicating whether the architecture include Batch Normalization layers or not. """ super(VGGFeatureExtractor, self).__init__() self._config = config if self._config is None: self._config = VGG16_BN_CONFIGS.get('10conv') self.features = make_layers(self._config, batch_norm=batch_norm) self._indices = None def forward(self, x): self._indices = [] self._pre_pool_shapes = [] for layer in self.features: if isinstance(layer, torch.nn.modules.pooling.MaxPool2d): self._pre_pool_shapes.append(x.size()) x, indices = layer(x) self._indices.append(indices) else: x = layer(x) return x def vgg16_bn_feature_extractor(config=None, pretrained=True, progress=True): model = VGGFeatureExtractor(config, batch_norm=True) if pretrained: state_dict = tv.models.utils.load_state_dict_from_url( VGG16_BN_MODEL_URL, progress=progress) model.load_state_dict(state_dict, strict=False) return model class DIM(torch.nn.Module): """Deep Image Matting.""" def __init__(self, feature_extractor): """ Constructor. Args: feature_extractor: Feature extractor, such as VGGFeatureExtractor. """ super(DIM, self).__init__() # Head convolution layer, number of channels: 4 -> 3 self._head_conv = torch.nn.Conv2d(in_channels=4, out_channels=3, kernel_size=5, padding=2) # Encoder self._feature_extractor = feature_extractor self._feature_extract_config = self._feature_extractor._config # Decoder self._decode_layers = self.decode_layers() # Prediction self._final_conv = torch.nn.Conv2d(self._feature_extract_config[0], 1, kernel_size=5, padding=2) self._sigmoid = torch.nn.Sigmoid() def forward(self, x): x = self._head_conv(x) x = self._feature_extractor(x) indices = self._feature_extractor._indices[::-1] unpool_shapes = self._feature_extractor._pre_pool_shapes[::-1] index = 0 for layer in self._decode_layers: if isinstance(layer, torch.nn.modules.pooling.MaxUnpool2d): x = layer(x, indices[index], output_size=unpool_shapes[index]) index += 1 else: x = layer(x) x = self._final_conv(x) x = self._sigmoid(x) return x def decode_layers(self): layers = [] strides = [1] channels = [] config_reversed = self._feature_extract_config[::-1] for i, v in enumerate(config_reversed): if v == 'M': strides.append(2) channels.append(config_reversed[i+1]) channels.append(channels[-1]) in_channels = self._feature_extract_config[-1] for stride, out_channels in zip(strides, channels): if stride == 2: layers += [torch.nn.MaxUnpool2d(kernel_size=2, stride=2)] layers += [torch.nn.Conv2d(in_channels, out_channels, kernel_size=5, padding=2), torch.nn.BatchNorm2d(num_features=out_channels), torch.nn.ReLU(inplace=True)] in_channels = out_channels return torch.nn.Sequential(*layers) def loss(alphas_pred, alphas_gt, images=None, epsilon=1e-12): losses = torch.sqrt( torch.mul(alphas_pred - alphas_gt, alphas_pred - alphas_gt) + epsilon) loss = torch.mean(losses) if images is not None: images_fg_gt = torch.mul(images, alphas_gt) images_fg_pred = torch.mul(images, alphas_pred) images_fg_error = images_fg_pred - images_fg_gt losses_image = torch.sqrt( torch.mul(images_fg_error, images_fg_error) + epsilon) loss += torch.mean(losses_image) return loss # def loss2(alphas_pred, alphas_gt, masks, images=None, epsilon=1e-12): diff = alphas_pred - alphas_gt diff = diff * masks num_unkowns = torch.sum(masks) + epsilon losses = torch.sqrt(torch.mul(diff, diff) + epsilon) loss = torch.sum(losses) / num_unkowns if images is not None: images_fg_gt = torch.mul(images, alphas_gt) images_fg_pred = torch.mul(images, alphas_pred) images_fg_diff = images_fg_pred - images_fg_gt images_fg_diff = images_fg_diff * masks losses_image = torch.sqrt( torch.mul(images_fg_diff, images_fg_diff) + epsilon) loss += torch.sum(losses_image) / num_unkowns return lossPytorch 的官方是带有 VGG 系列模型的，使用也很方便，比如使用带批标准化层的 VGG16 直接写为：vgg = torchvision.models.vgg16_bn(pretrained=True)其中 pretrained=True 表示导入在 ImageNet 上预训练的参数. 但因为，我们是只使用 VGG16 的前 13 个卷积层，而不需要后面的全连接层，因此，不会像上面那样直接使用，而是要从 torchvision 的官方实现中截取卷积层的部分(官方实现见文件 Python 安装路径下的 site-packages/torchvision/models/vgg.py). 我们主要复制该文件中的 make_layers 函数，但因为后面解码器阶段要用反池化，所以还要做一些修改：要把layers += [nn.MaxPool2d(kernel_size=2, stride=2)]改为layers += [torch.nn.MaxPool2d(kernel_size=2, stride=2, return_indices=True)]之所以要加上 return_indices=True，是因为后面反池化层要用到这些池化层的池化过程中的最大值的下标(从而要记下来). 也正因为池化层多返回了一个值(同时返回特征映射和最大值下标张量)，因此在重载 forward 函数时要进行如下的区别对待：def forward(self, x): self._indices = [] for layer in self.features: if isinstance(layer, torch.nn.modules.pooling.MaxPool2d): x, indices = layer(x) self._indices.append(indices) else: x = layer(x) return x除此之外，编码器阶段都是非常简单的，无需赘言. 来看解码器阶段. 只需要重点关注一下反池化层(接后续的卷积层)的实现即可. 具体也很简单：unpool = torch.nn.MaxUnpool2d(kernel_size=2, stride=2) conv = torch.nn.Conv2d(in_channels, out_channels, kernel_size=5, padding=2)即，先用反池化操作提升 2 倍分辨率，然后再接一个普通卷积层(可选操作：批标准化、整流线性单元). 实际前向传播时的计算如下：x = unpool(x, indices) x = conv(x)其中，indices 是编码器阶段对应的池化层返回的最大池化操作返回的最大值下标. 反池化层与转置卷积层都是可以训练的(都带有参数)，作用也几乎相同(提升分辨率)，但对于抠图这个任务来说，最关心的就是目标的边界区域，而这些边界因为都是前景、背景的交界区，因此表现在特征映射的响应上，就基本都是局部极大值，从而在池化操作时，返回的最大值下标就基本完整的记录下了待抠目标的边界，反池化操作因为会重点关注这些区域，所以效果较好. DIM 类的 decode_layers 函数就是整个的解码器的定义. 它看起来有点晦涩，其实不难理解. 我们看编码器阶段的配置：[64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512]其中的 M 表示最大池化层，其它的数字就是全部的 13 个卷积层对应的输出通道数. 解码器执行的操作基本上就是以上操作的逆过程：[512, 'U', 512, 'U', 256, 'U', 128, 'U', 64]其中的 U 表示 torch.nn.MaxUnpool2d 反池化操作.5. 训练过程本节训练数据来源于爱分割开源的 数据集. 该数据集内的所有 34426 张图片都类似于如下的上身模特图：可以明显看到标注的 alpha 通道是非常粗陋的，远达不到头发丝的精度. 因为这个数据集直接给出了原始图像，所以不需要前景、背景图片的合成.5.1. 数据准备当下载好爱分割开源数据集(并解压)之后，我们需要一次性将所有图片的掩码(mask)都准备好，因此需要打开 data/retrieve.py 文件，将 root_dir 改成 Matting_Human_Half 文件夹的路径，然后执行 retrieve.py 等待生成所有图片的 alpha 和 mask(在 Matting_Human_Half 文件夹内)，以及用于训练的 train.txt 和 val.txt(在 data 文件夹内，其中默认随机选择 100 张图像用于验证). 假如，训练时不再改动 Matting_Human_Half 文件夹的路径，那么就不需要再做其它处理了. 如果训练时，Matting_Human_Half 与以上制作 train.txt 和 val.txt 时指定的 root_dir 路径不一致了，那么可以使用诸如 Notepad ++ 之类的工具，将 root_dir 替换为空，形成如下的形式：train.txt 和 val.txt 分别记录了训练和验证图像的路径，每一行对应一张图像的 4 个路径，分别是原图像路径(3 通道)、透明图路径(4 通道)、alpha 通道图像路径、mask 路径，它们通过 @ 符号分隔.5.2. 训练直接在命令行执行：python3 train.py --root_dir "xxx/Matting_Human_Half" [--gpu_indices 0 1 ...]开始训练，如果从制作数据时开始， Matting_Human_Half 这个文件夹的路径始终没有改动过，那么 root_dir 这个参数也可以不指定(指定也无妨). 后面的 [--gpu_indices ...] 表示需要根据实际情况，可选的指定可用的 GPU 下标，这里默认是使用 0,1,2,3 共 4 块 GPU，如果使用一块 GPU，则指定--gpu_indices 0如果使用多块，比如使用 第 1 块和第 3 块 GPU，则指定--gpu_indices 1 3即可. 其它类似. 训练过程中的所有超参数都在 train.py 文件的开头部分，可以直接修改默认值或通过命令行指定. 训练开始几分钟后，在项目路径下执行：tensorboard --logdir ./models/logs可以打开浏览器查看训练的学习率、损失曲线，和训练过程中的分割结果图像. 这里使用的是 Pytorch 自带的类：from torch.utils.tensorboard import SummaryWriter 来调用 tensorboard，因此需要 Pytorch 1.1.0 以及之后的版本才可以. (但好像浏览器刷新不了新结果，需要不断重开 tensorboard 才可以观看训练进展)训练结束后(默认训练 30 个 epoch)，在 models 文件夹中保存了训练过程中的模型参数文件(模型使用参考 predict.py). 直接执行python3 predict.py将在 test 文件夹里生成测试图片的抠图结果.5.3. 其它数据集上训练训练 Pytorch 模型时，需要重载 torch.utils.data.Dataset 类，用来提供数据的批量生成. 重载时，只需要实现 init, getitem, len 这三个函数. 在这个项目里，我们使用的是 dataset.py 的重载类 MattingDataset. 可以按照自己的方式依据自己的标注格式来重载，也可以依照 MattingDataset 来改写. 对于只提供前、背景分离的数据，建议先一次性提前合成好合成图像，和制作好 alpha 通道图像. 此时就可以适当修改一下 get_image_mask_paths 函数即可. 这个函数需要返回一个如下格式的列表：[[image_path, alpha_path], [image_path, alpha_path], ... [image_path, alpha_path]]另外，__ getitem __ 函数数据增强的方式裁剪、缩放、水平翻转(以及 alpha 通道随机膨胀腐蚀)，如果还有其它的处理方式请自行添加或删减. 另外，这里指定的裁剪尺寸：crop_sizes = [320, 480, 600, 800]是根据爱分割提供的数据来划定的，里面所有图片都是 600x800 的分辨率. 一般来说，根据 Deep Image Matting 论文，是从 320 开始，每隔 160 像素的尺寸裁剪，最后统一缩放到 320 即可.6. DIM 数据集上的复现本节将在 Deep Image Matting 数据集上进行训练(训练数据可联系论文作者获取)，因部分参数未仔细调整，训练结果并非最优. 使用的背景图像集是 COCO/train2017.6.1. 数据准备我们将前景图像和背景图像通过 alpha 通道合成训练集. 假设已经获取了 DIM 数据集，那么进入 data_dim 文件夹，打开 composition.py，将 root_dir 替换成你保持 Combined_Dataset 文件夹的路径，bg_image_root_dir 填写背景图像的文件夹路径，output_dir 填写合成图像的保持文件夹路径. num_bg_images_per_fg 表示一张前景图像对应多少张背景图像，论文里这个值取 100(这里为了减少训练时间，本文取的是 50，读者根据具体情况修改). 当这些值都确认无误后，执行 composition.py，将花费很长一段时间来合成图片. 合成图像结束后，在当前路径下生成 train.txt 和 val.txt 两个标注文件. train.txt 文件里每一行对应 4 个路径，分别是合成图像路径、前景图像路径、alpha 通道路径、背景图像路径，val.txt 里除了缺少合成图像路径之外，其它顺序一致. 训练时，我们只需要 train.txt 中的第 1 个和第 3 个路径，因此和训练爱分割开源数据集的标注文件格式一致，从而可以共用同一个 dataset.py，只需要确保 dataset.py 里面的函数 getitem 中的crop_sizes = [320, 480, 640]即可.6.2. 训练过程执行python3 train.py --annotation_path "./data_dim/train.txt" [gpu_indices 0 1 2 ...]开始训练，期间可通过 tensorboard 查看训练进程. 训练的所有超参数可在 train.py 内修改，也可以通过命令行直接指定. 图：训练期间的 tensorboard 抠图结果展示图：训练期间的 tensorboard 学习率和损失曲线本次训练，超参数采用的都是 train.py 中的默认值，由损失曲线可以看到，如果再继续训练(已训练 200 epoch)，损失会进一步下降，抠图效果会更好.6.3. 结果展示训练结束后，执行(如果模型保存路径是默认的 ./models，否则需要修改一下 ckpt_path)：python3 predict_trimap.py会在 data_dim/test 文件夹里生成预测结果(见文件夹 preds)：图：合成图片、trimap、抠图结果、GT alpha图：合成图片、trimap、抠图结果、GT alpha从以上抠图结果可看出，在某些细节上效果还不理想. 如果想获得更好的结果，一方面可以在合成训练图片时提升 1 张前景对应的背景图片数(本文取的是 1:50，论文是 1:100)；其次，仔细调整学习率及其衰减；再次，增加训练的轮次(epoch，本文训练了 200 个 epoch)等.
贝赛尔工具在 photoshop 中叫“钢笔工具”；在 CorelDraw 中翻译成“贝赛尔工具”；而在 Fireworks 中叫“画笔”. 它是用来“画线”造型的一种专业工具. 当然还有很多工具也可以完成画线的工作，例如常用的photoshop里的直线、喷枪、画笔工具，Fireworks里的直线、铅笔和笔刷工具，CorelDraw里的自由笔，手绘工具等等.很多设计师应该都知道贝塞尔曲线，习惯用 PS 的会用钢笔工具，习惯用AI 的会用贝塞尔，因为它所绘制出来的曲线很容易受控制，也很美观. 下面深入了解下贝塞尔曲线的数学原理和公式.在数学中，贝塞尔曲线又分为很多种，如，一阶贝塞尔曲线、二阶贝塞尔曲线、三阶贝塞尔曲线，等等. 除了一阶贝塞尔曲线是直线外，剩下的多阶贝塞尔曲线都是抛物线. 贝塞尔曲线由起点、终点和控制点组成，根据控制点的个数和位置决定了这个曲线的最终样式.1. 原理贝塞尔曲线的数学原理 - 2017.01.26先在一个平面内任选 3 个不共线的点，依次用线段连接，如图：在第一条线段上任选一个点 D，计算该点到线段起点的距离 AD，与该线段总长 AB 的比例.根据上一步得到的比例，从第二条线段上找出对应的点 E，使得 AD:AB = BE:BC.连接这两点 DE.从新的线段 DE 上再次找出相同比例的点 F，使得 DF:DE = AD:AB=BE:BC.到这里，就确定了贝塞尔曲线上的一个点 F. 接下来，稍微回想一下中学所学的极限知识，让选取的点 D 在第一条线段上从起点 A 移动到终点 B，找出所有的贝塞尔曲线上的点 F. 所有的点找出来以后，也就得到了这条贝塞尔曲线.如果想象不出来，可以看下面的动画.到这里，已经大概了解贝塞尔曲线绘制的过程了.下面看公式.2. 公式以下公式中：$B(t)$ 为 $t$ 时间下点的坐标； $P_0$ 为起点；$P_n$ 为终点，$P_i$ 为控制点.2.1. 一阶贝塞尔曲线一阶贝塞尔曲线只有起点和终点，并没有控制点，所以绘制出来的图形仅仅只是一条直线，那么在时间 t 为 1 秒的情况下，其公式为：$$ B(t) = (1 - t) P_0 + tP_1, t \in [0, 1] $$意义：由 $P_0$ 至 $P_1$ 的连续点， 描述的一条线段.2.2. 二阶贝塞尔曲线二阶贝塞尔曲线只存在一个控制点，此时从起点到终点的线段发生变化，具体的变化是由控制点的位置而改变的. 图中，绿色的线段为红色曲线的切线.$$ B(t) = (1 - t)^2 P_0 + 2t(1 - t)P_1 + t^2 P_2, t \in [0, 1] $$仅仅只是简单的一元二次方程式.原理：由 $P_0$ 至 $P_1$ 的连续点 $Q_0$，描述一条线段；由 $P_1$ 至 $P_2$ 的连续点 $Q_1$，描述一条线段；由 $Q_0$ 至 $Q_1$ 的连续点 $B(t)$，描述一条二次贝塞尔曲线.Python 示例如：#!/usr/bin/python3 #!--*-- coding: utf-8 --*-- import numpy as np import matplotlib.pyplot as plt P0, P1, P2 = np.array([[0, 0], [2, 4], [5, 3]]) # 定义贝塞尔曲线 P = lambda t: (1 - t)**2 * P0 + 2 * t * (1 - t) * P1 + t**2 * P2 # 在 [0, 1] 范围内的 50 个点上验证贝塞尔曲线 points = np.array([P(t) for t in np.linspace(0, 1, 50)]) # 分别获取点的 x 坐标和 y 坐标 x, y = points[:, 0], points[:, 1] # plt.plot(x, y, 'b-') plt.plot(*P0, 'r.') plt.plot(*P1, 'r.') plt.plot(*P2, 'r.') plt.show()如：可以看到，曲线从第一个控制点处开始，到最后一个控制点处结束. 这个结果对任意数量的点都成立.从P0到P2，P1 完全决定了曲线的形状。移动 P1：贝塞尔曲线总是包含在控制点形成的多边形中. 这个多边形因此被称为控制多边形，或者贝塞尔多边形. 这个属性也适用于任意数量的控制点，这使得它们在使用软件时的操作非常直观.控制多边形还具有以下特性：包含曲线的面积最小，称为凸包.2.3. 高阶贝塞尔曲线[1] - 三阶[2] - 四阶[3] - 五阶从三阶开始贝塞尔曲线就开始显得复杂了，高阶的通用公式如下：$$ P_i^k = (1 - t) P_i^{k-1} + t P_{i+1}^{k-1} $$其中，$k = 1, 2, ..., n$，$i = 0, 1, ..., n-k$.https://pages.mtu.edu/~shene/COURSES/cs3621/NOTES/spline/Bezier/de-casteljau.html任意数量的控制点贝塞尔曲线的 Python 实现：#!/usr/bin/python3 #!--*-- coding: utf-8 --*-- import numpy as np import matplotlib.pyplot as plt from math import factorial def comb(n, k): return factorial(n) // (factorial(k) * factorial(n-k)) def get_bezier_curve(points): n = len(points) - 1 return lambda t: sum(comb(n, i)*t**i * (1-t)**(n-i)*points[i] for i in range(n+1)) def evaluate_bezier(points, total): bezier = get_bezier_curve(points) new_points = np.array([bezier(t) for t in np.linspace(0, 1, total)]) return new_points[:, 0], new_points[:, 1] points = np.array([[0, 0], [-1, 3], [4, 3], [6, 0], [7, 2.5]]) x, y = points[:, 0], points[:, 1] bx, by = evaluate_bezier(points, 50) # plt.plot(bx, by, 'b-') plt.plot(x, y, 'r.') plt.show()如图：3. 插值贝塞尔曲线的一个有趣应用是绘制一条通过一组预定义点的平滑曲线。之所以有趣，是因为 $P(t)$ 的公式产生点，并且不是 $y = f(x)$ 的形式，因此一个 $x$ 有多个 $y$. 例如，可以这样画：贝塞尔插值实现如：From: Bézier Interpolation - Create smooth shapes using Bézier curves - 2020.05.09#!/usr/bin/python3 #!--*-- coding: utf-8 --*-- import numpy as np import matplotlib.pyplot as plt # find the a & b points def get_bezier_coef(points): # since the formulas work given that we have n+1 points # then n must be this: n = len(points) - 1 # build coefficents matrix C = 4 * np.identity(n) np.fill_diagonal(C[1:], 1) np.fill_diagonal(C[:, 1:], 1) C[0, 0] = 2 C[n - 1, n - 1] = 7 C[n - 1, n - 2] = 2 # build points vector P = [2 * (2 * points[i] + points[i + 1]) for i in range(n)] P[0] = points[0] + 2 * points[1] P[n - 1] = 8 * points[n - 1] + points[n] # solve system, find a & b A = np.linalg.solve(C, P) B = [0] * n for i in range(n - 1): B[i] = 2 * points[i + 1] - A[i + 1] B[n - 1] = (A[n - 1] + points[n]) / 2 return A, B # returns the general Bezier cubic formula given 4 control points def get_cubic(a, b, c, d): return lambda t: np.power(1 - t, 3) * a + 3 * np.power(1 - t, 2) * t * b + 3 * (1 - t) * np.power(t, 2) * c + np.power(t, 3) * d # return one cubic curve for each consecutive points def get_bezier_cubic(points): A, B = get_bezier_coef(points) return [ get_cubic(points[i], A[i], B[i], points[i + 1]) for i in range(len(points) - 1) ] # evalute each cubic curve on the range [0, 1] sliced in n points def evaluate_bezier(points, n): curves = get_bezier_cubic(points) return np.array([fun(t) for fun in curves for t in np.linspace(0, 1, n)]) # # generate 5 (or any number that you want) random points that we want to fit (or set them youreself) points = np.random.rand(5, 2) # fit the points with Bezier interpolation # use 50 points between each consecutive points to draw the curve path = evaluate_bezier(points, 50) # extract x & y coordinates of points x, y = points[:,0], points[:,1] px, py = path[:,0], path[:,1] # plot plt.figure(figsize=(11, 8)) plt.plot(px, py, 'b-') plt.plot(x, y, 'ro') plt.show()4. 三阶贝塞尔曲线插值平滑From: python基于三阶贝塞尔曲线的数据平滑算法 - 2019.12.27很多文章在谈及曲线平滑的时候，习惯使用拟合的概念，可能不太不恰当的. 平滑后的曲线，一定经过原始的数据点，而拟合曲线，则不一定要经过原始数据点.一般而言，需要平滑的数据分为两种：时间序列的单值数据和时间序列的二维数据. 对于前者，并非一定要用贝塞尔算法，仅用样条插值就可以轻松实现平滑；而对于后者，不管是 numpy 还是 scipy 提供的那些插值算法，就都不适用了.基于三阶贝塞尔曲线，实现时间序列的单值数据和时间序列的二维数据的平滑算法，可满足大多数的平滑需求。4.1. 算法描述如果把三阶贝塞尔曲线的 P0 和 P3 视为原始数据，只要找到 P1 和 P2 两个点(称其为控制点)，就可以根据三阶贝塞尔曲线公式，计算出 P0 和 P3 之间平滑曲线上的任意点。现在，平滑问题变成了如何计算两个原始数据点之间的控制点的问题。步骤如下：[1] - 绿色直线连接相邻的原始数据点，计算出个线段的中点，红色直线连接相邻的中点[2] - 根据相邻两条绿色直线长度之比，分割其中点之间红色连线，标记分割点[3] - 平移红色连线，使其分割点与相对的原始数据点重合[4] - 调整平移后红色连线的端点与原始数据点的距离，通常缩减40%-80%4.2. Python 实现# !/usr/bin/python3 # !--*-- coding: utf-8 --*-- import numpy as np import matplotlib.pyplot as plt def bezier_curve(p0, p1, p2, p3, inserted): """ 三阶贝塞尔曲线 p0, p1, p2, p3 - 点坐标，tuple、list或numpy.ndarray类型 inserted - p0和p3之间插值的数量 """ assert isinstance(p0, (tuple, list, np.ndarray)) assert isinstance(p0, (tuple, list, np.ndarray)) assert isinstance(p0, (tuple, list, np.ndarray)) assert isinstance(p0, (tuple, list, np.ndarray)) if isinstance(p0, (tuple, list)): p0 = np.array(p0) if isinstance(p1, (tuple, list)): p1 = np.array(p1) if isinstance(p2, (tuple, list)): p2 = np.array(p2) if isinstance(p3, (tuple, list)): p3 = np.array(p3) points = list() for t in np.linspace(0, 1, inserted + 2): points.append(p0 * np.power((1 - t), 3) + 3 * p1 * t * np.power((1 - t), 2) + 3 * p2 * (1 - t) * np.power(t, 2) + p3 * np.power( t, 3)) return np.vstack(points) def smoothing_base_bezier(date_x, date_y, k=0.5, inserted=10, closed=False): """ 基于三阶贝塞尔曲线的数据平滑算法 date_x - x维度数据集，list或numpy.ndarray类型 date_y - y维度数据集，list或numpy.ndarray类型 k - 调整平滑曲线形状的因子，取值一般在0.2~0.6之间。默认值为0.5 inserted - 两个原始数据点之间插值的数量。默认值为10 closed - 曲线是否封闭，如是，则首尾相连。默认曲线不封闭 """ assert isinstance(date_x, (list, np.ndarray)) assert isinstance(date_y, (list, np.ndarray)) if isinstance(date_x, list) and isinstance(date_y, list): assert len(date_x) == len(date_y), u'x数据集和y数据集长度不匹配' date_x = np.array(date_x) date_y = np.array(date_y) elif isinstance(date_x, np.ndarray) and isinstance(date_y, np.ndarray): assert date_x.shape == date_y.shape, u'x数据集和y数据集长度不匹配' else: raise Exception(u'x数据集或y数据集类型错误') # 第1步：生成原始数据折线中点集 mid_points = list() for i in range(1, date_x.shape[0]): mid_points.append({ 'start': (date_x[i - 1], date_y[i - 1]), 'end': (date_x[i], date_y[i]), 'mid': ((date_x[i] + date_x[i - 1]) / 2.0, (date_y[i] + date_y[i - 1]) / 2.0) }) if closed: mid_points.append({ 'start': (date_x[-1], date_y[-1]), 'end': (date_x[0], date_y[0]), 'mid': ((date_x[0] + date_x[-1]) / 2.0, (date_y[0] + date_y[-1]) / 2.0) }) # 第2步：找出中点连线及其分割点 split_points = list() for i in range(len(mid_points)): if i < (len(mid_points) - 1): j = i + 1 elif closed: j = 0 else: continue x00, y00 = mid_points[i]['start'] x01, y01 = mid_points[i]['end'] x10, y10 = mid_points[j]['start'] x11, y11 = mid_points[j]['end'] d0 = np.sqrt(np.power((x00 - x01), 2) + np.power((y00 - y01), 2)) d1 = np.sqrt(np.power((x10 - x11), 2) + np.power((y10 - y11), 2)) k_split = 1.0 * d0 / (d0 + d1) mx0, my0 = mid_points[i]['mid'] mx1, my1 = mid_points[j]['mid'] split_points.append({ 'start': (mx0, my0), 'end': (mx1, my1), 'split': (mx0 + (mx1 - mx0) * k_split, my0 + (my1 - my0) * k_split) }) # 第3步：平移中点连线，调整端点，生成控制点 crt_points = list() for i in range(len(split_points)): vx, vy = mid_points[i]['end'] # 当前顶点的坐标 dx = vx - split_points[i]['split'][0] # 平移线段x偏移量 dy = vy - split_points[i]['split'][1] # 平移线段y偏移量 sx, sy = split_points[i]['start'][0] + dx, split_points[i]['start'][1] + dy # 平移后线段起点坐标 ex, ey = split_points[i]['end'][0] + dx, split_points[i]['end'][1] + dy # 平移后线段终点坐标 cp0 = sx + (vx - sx) * k, sy + (vy - sy) * k # 控制点坐标 cp1 = ex + (vx - ex) * k, ey + (vy - ey) * k # 控制点坐标 if crt_points: crt_points[-1].insert(2, cp0) else: crt_points.append([mid_points[0]['start'], cp0, mid_points[0]['end']]) if closed: if i < (len(mid_points) - 1): crt_points.append([mid_points[i + 1]['start'], cp1, mid_points[i + 1]['end']]) else: crt_points[0].insert(1, cp1) else: if i < (len(mid_points) - 2): crt_points.append([mid_points[i + 1]['start'], cp1, mid_points[i + 1]['end']]) else: crt_points.append([mid_points[i + 1]['start'], cp1, mid_points[i + 1]['end'], mid_points[i + 1]['end']]) crt_points[0].insert(1, mid_points[0]['start']) # 第4步：应用贝塞尔曲线方程插值 out = list() for item in crt_points: group = bezier_curve(item[0], item[1], item[2], item[3], inserted) out.append(group[:-1]) out.append(group[-1:]) out = np.vstack(out) return out.T[0], out.T[1] if __name__ == '__main__': x = np.array([2, 4, 4, 3, 2]) y = np.array([2, 2, 4, 3, 4]) plt.plot(x, y, 'ro') x_curve, y_curve = smoothing_base_bezier(x, y, k=0.3, closed=True) plt.plot(x_curve, y_curve, label='$k=0.3$') x_curve, y_curve = smoothing_base_bezier(x, y, k=0.4, closed=True) plt.plot(x_curve, y_curve, label='$k=0.4$') x_curve, y_curve = smoothing_base_bezier(x, y, k=0.5, closed=True) plt.plot(x_curve, y_curve, label='$k=0.5$') x_curve, y_curve = smoothing_base_bezier(x, y, k=0.6, closed=True) plt.plot(x_curve, y_curve, label='$k=0.6$') plt.legend(loc='best') plt.show()下图为平滑效果。左侧是封闭曲线，两个原始数据点之间插值数量为默认值10；右侧为同样数据不封闭的效果，k值默认0.5.材料[1] - 关于贝塞尔曲线的公式推导和python代码实现 - 2020.09.27 - 知乎[2] - 了解贝塞尔曲线的数学和Python实现示例 - 2020.05.03[3] - Bézier Interpolation - Create smooth shapes using Bézier curves - 2020.05.09[4] - python基于三阶贝塞尔曲线的数据平滑算法 - 2019.12.27[5] - Interpolation with Bezier Curves - A very simple method of smoothing polygons[6] - 贝塞尔曲线 - 百度百科
1. 颜色背景实现如下：#!/usr/bin/python3 #!--*-- coding: utf-8 --*-- import numpy as np from PIL import Image import matplotlib.pyplot as plt img_file = 'test.jpg' png_file = './test.png' img_pil = Image.open(img_file) png_pil = Image.open(png_file) alpha = png_pil.split()[-1] ## pillow #im = Image.new('RGBA', png_pil.size, (255, 0, 255)) #im.paste(im, mask=im) # opencv bg = Image.new("RGB", png_pil.size, (255, 0, 255)) fg = np.array(img_pil, np.float32) alpha= np.expand_dims(np.array(alpha) / 255,axis=2) im = alpha * fg + (1 - alpha) * bg im = im.astype(np.uint8) # plt.figure(figsize=(12, 10)) plt.subplot(131) plt.imshow(img_pil) plt.subplot(132) plt.imshow(alpha) plt.subplot(133) plt.imshow(im) plt.show()如，左原图，中alpha图，右新背景图.2. 图片背景#!/usr/bin/python3 #!--*-- coding: utf-8 --*-- from PIL import Image img_file = 'test.jpg' png_file = './test.png' # bg = Image.open(bg_file) png_pil = Image.open(png_file) if png_pil.size != bg.size: png_pil = png_pil.resize(bg.size) r, g, b, alpha = png_pil.split() fg = Image.merge('RGB', (r, g, b)) out = Image.composite(fg, bg, png_pil)
论文： Background Matting: The World is Your Green Screen - CVPR2020Projects：Background Matting: The World is Your Green ScreenGithub: senguptaumd/Background-MattingBlog: Background Matting: The World is Your Green Screen - MediumBackground Matting 示例：图：手持智能手机相机，拍摄同一场景的两张照片 - 一张有目标(如，人物)和一张无目标；采用对抗损失(adversarial loss) 的神经网络来重建 alpha matte 和前景(foreground)；最后合并结果到新的背景(backgroud).1. Matting 问题Matting 是将图像分解为前景和背景的处理，以便于将前景合并到新的背景上. 这也是绿幕(green screen) 的关键技术，被广泛应用于视频生产、图形和APPs 中.Matting 问题的建模，将图像的每个像素表示为前景和背景的组合：$$ C = F * \alpha + B * (1 - \alpha) $$Matting 问题的求解是，给定图像 C，计算每个像素的前景F、背景B 和透明度alpha.由于图像有 RGB 三个通道，该公式需要基于 3 个已知量来求解 7 个未知量，问题是难以求解的.1.1. Matting 与 Segmentation一种可行方式是，采用分割(segmentation) 将前景从图像中分离出来. 但是，尽管分割技术近年取得了很大进步，其并没有解决 matting 公式里的问题.分割是对图像中每个像素分配一个二值标签(0, 1)，以表示背景和前景，而不是求解连续的 alpha 值. Segmention 和 Matting 的区别可以如下图简单例示：图：Segmentation 采用 DeepLabv3+ 得到. 图中，边缘区域，尤其是头发，是由 0-1 之间的真实 alpha 值的. 因此，分割的二值性导致前景边界周围出现严格边界，留下可见的不好效果. 解决了部分透明度和前景问题，对第二帧的合成更有利.1.2. Matting 与背景提取Is This Method Like Background Subtraction?Matting 问题和另一个接近的问题是，是否和背景提取类似呢？首先，如果使用任意背景组合这么简单，电影产业就不会每年花费大量美元投资在绿幕(green screens) 上.图：背景提取并不能对于任意拍摄背景取得较好效果.此外，背景提取并不能解决 alpha 值的问题，与分割一样存在相同的硬边缘问题. 且，当背景和前景比较相似、背景存在运动时，背景提取也不能很好的解决.2. 使用任意拍摄的背景Using A Casually Captured Background由于 matting 是比 segmentation 更难得问题，即使是采用深度学习技术，往往需要再利用其他信息来解决 matting 无约束问题.比如，很多算法需要采用 trimap 或手工标注的已知前景、背景和未知区域. 尽管其对于单张图片是可行的，但视频的标注是极其费时的，因此并不是视频抠图的一个可行研究方向.而这里，采用拍摄的背景作为真实背景的一种估计. 其使得前景和 alpha 值得求解更容易. 称之为“随意拍摄(casually captured)”背景，因为其包含轻微的抖动、颜色差一、轻微阴影、或者与前景相似的颜色.图：拍摄过程. 当目标离开场景后，再拍摄的一张背景有助于算法.图：拍摄的输入、背景和新背景组合后的结果.需要注意的是，这张图片是很有挑战的，因为其前景颜色和背景颜色非常相似(尤其是的头发周围). 而且，其是采用手持设备拍摄的，会包含轻微的背景抖动.“We call it a casually captured background because it can contain slight movements, color differences, slight shadows, or similar colors as the foreground.”2.1. 拍摄建议尽管本文算法能够处理一定的背景干扰，但如果背景是固定的且最好是室内，其效果会更好. 例如，算法对于由于目标所导致的存在高度可见阴影、背景移动(如水流、车流、树木) 或大曝光变化的情况时，是效果不好的. 图：失败情况. 人物在移动的喷泉前面.此外，还推荐在视频结束时，让人物离开场景；然后从连续的视频中拉出那一帧画面. 许多智能手机对于视频模式切换到照片模式时，都会有不同的变焦(zoom)和曝光(exposure)设置. 当采用手机拍摄时，也应该启用自动曝光(auto-exposure)锁定.图：理想拍摄场景. 背景在室内、非移动、目标无阴影.拍摄建议总结：[1] - 选择最固定的背景；[2] - 不要和背景站的太近，避免投下阴影.[3] - 开启手机的自动曝光(auto-exposure) 和自动对焦(auto-focus)锁定.3. 网络细节主要思路：进行拍照时额外拍一张对应的没有前景目标的背景图片，以辅助视频抠图.网络包括监督学习步骤和无监督精细化(refinement).图：白底区域-监督学习网络；浅色区域-无监督学习网络.3.1. 监督学习网络图：监督学习网络. 采用多个输入，输出 alpha matte 和预测的前景. 监督网络的训练，采用 Adobe Composition-1k 数据集，包含 450 张仔细标注 GT alpha mattes. 网络的训练采用完全监督的方式，输出层采用逐像素 loss.网络包含多个输入：图像(I)，背景(B'，需要将其预处理进行对齐)，软分割(S， 采用分割算法生成，如DeepLabv3+)，运动线索(M，时序运动信息，当前帧的前两帧和后两帧，需要数据处理为灰度图以忽略颜色信息，专注于运动信息).其中，提出的 Context Switching Block (对4个输入分别用编码器生成特征图；Selector块和Combinator 都由 1x1 Conv+BN+Relu组成) 对于不好的输入，仍具有鲁棒性.From: 关于CVPR2020论文 Background Matting的一些理解3.2. 基于 GANs 的无监督精细网络监督学习的问题在于，adobe 数据集近包含 450 张 GT，其是不足以训练较好的网络的. 而，获取更多的数据是极其有难度的，因为需要手工标注图像的 alpha matte.为了解决这个问题，这里采用 GAN 精细化处理. 采用监督网络所输出的 alpha matte，并将其组合到新的背景. 然后，判别网络(Discriminator) 尝试判断其是 real 或 fake 图片. 相应地，生成网络(Generator) 学习地更新 alpha matte；因此，最终的组合结果会越来越真实，以欺骗判别网络.图：无监督 GAN 精细网络.训练过程主要分为两步：[1] - 在Adobe数据集上进行有监督的训练，得到GAdobe；[2] - 在未标注的真实数据上进行自监督的对抗训练，得到Greal；无监督 GAN 处理的重要性在于，其不需要任何标注的训练数据. 判别网络的训练是在很多 real 图像上进行的，也是非常容易得到的.3.3. 图像对齐test_pre_process.py: 采用 AKAZE 算法图像对齐.import numpy as np import cv2 MAX_FEATURES = 500 GOOD_MATCH_PERCENT = 0.15 def alignImages(im1, im2,masksDL): # Convert images to grayscale im1Gray = cv2.cvtColor(im1, cv2.COLOR_BGR2GRAY) im2Gray = cv2.cvtColor(im2, cv2.COLOR_BGR2GRAY) akaze = cv2.AKAZE_create() keypoints1, descriptors1 = akaze.detectAndCompute(im1, None) keypoints2, descriptors2 = akaze.detectAndCompute(im2, None) # Match features. matcher = cv2.DescriptorMatcher_create(cv2.DESCRIPTOR_MATCHER_BRUTEFORCE) matches = matcher.match(descriptors1, descriptors2, None) # Sort matches by score matches.sort(key=lambda x: x.distance, reverse=False) # Remove not so good matches numGoodMatches = int(len(matches) * GOOD_MATCH_PERCENT) matches = matches[:numGoodMatches] # Extract location of good matches points1 = np.zeros((len(matches), 2), dtype=np.float32) points2 = np.zeros((len(matches), 2), dtype=np.float32) for i, match in enumerate(matches): points1[i, :] = keypoints1[match.queryIdx].pt points2[i, :] = keypoints2[match.trainIdx].pt # Find homography h, mask = cv2.findHomography(points1, points2, cv2.RANSAC) # Use homography height, width, channels = im2.shape im1Reg = cv2.warpPerspective(im1, h, (width, height)) # copy image in the empty region, unless it is a foreground. Then copy background mask_rep=(np.sum(im1Reg.astype('float32'),axis=2)==0) im1Reg[mask_rep,0]=im2[mask_rep,0] im1Reg[mask_rep,1]=im2[mask_rep,1] im1Reg[mask_rep,2]=im2[mask_rep,2] mask_rep1=np.logical_and(mask_rep , masksDL[...,0]==255) im1Reg[mask_rep1,0]=im1[mask_rep1,0] im1Reg[mask_rep1,1]=im1[mask_rep1,1] im1Reg[mask_rep1,2]=im1[mask_rep1,2] return im1Reg参考材料[1] S. Sengupta, V. Jayaram, B. Curless, S. Seitz, and I. Kemelmacher-Shlizerman, Background Matting: The World is Your Green Screen (2020), CVPR 2020[2] L.C. Chen, Y. Zhu, G. Papandreou, F. Schroff, and H. Adam, Encoder-Decoder with Atrous Separable Convolution for Semantic Image Segmentation (2018), ECCV 2018[3] Y.Y. Chuang, B. Curless, D. H. Salesin, and R. Szeliski, A Bayesian Approach to Digital Matting (2001), CVPR 2001[4] Q. Hou and F. Liu. Context-Aware Image Matting for Simultaneous Foreground and Alpha Estimation (2019), ICCV 2019[5] H. Lu, Y. Dai, C. Shen, and S. Xu, Indices Matter: Learning to Index for Deep Image Matting (2019), ICCV 2019[6] - 关于CVPR2020论文 Background Matting的一些理解
原文：Whiten/customize the background of images using Semantic Segmentation and Alpha Blending - 2019.12.26作者：Shilpi Bhattacharyya图：语义分割和 alpha blending 技术用于改变图像的背景.图：使用语义分割和 alpha blending 的处理流程. 将背景区域像素值设为 0.完整实现：https://github.com/shilpibhattacharyya/customizing_image_backgrounds是否有使用 Photoshop 进行图像背景白化(whiten)和改变时抓狂的时候呢？有的话，请进入语义分割.1. 语义分割语义分割是一种将图像中每个像素进行归类的图像分析任务. 类似于人的方式，当人在看一幅图像时，知道图像的每个像素属于哪一类.更多关于语义分割的资料，可参考：Image Segmentation - 2018.11.05这里采用 DeepLabV3 进行语义分割，alpha blending 进行将原图融合到新背景.图：输入图像采用 Deeplab 模型语义分割处理. 期望输出是，覆盖在输入图像上的语义标签. 语义分割模型关注与对图像中每个像素分配语义标签，如sky, person, car 等.DeepLab 是 Google Brain 提出的语义分割结构，其是一种先进的语义分割深度学习模型. 其系列主要有：DeepLabv1 - 采用 atrous conv 来显式地控制网络中所计算的特征图的分辨率；DeepLabv2 - 采用 atrous spatial pyramid pooling(ASPP) 来鲁棒的分割物体，其在不同的采样率和有效接受野进行多尺度处理(using atrous spatial pyramid pooling (ASPP) to robustly segment objects at multiple scales with filters at multiple sampling rates and effective fields-of-views)；DeepLabv3 - 采用图像层(image-level)特征来增强 ASPP 模块，以捕获更大范围的特征. 此外，还包含 batch normalization 参数来提升模型训练. 具体来说，在训练和验证阶段，采用 atrous conv 提取不同输出步长(output strides)的输出特征，有效的增强了在步长=16 的 BN 训练，且验证阶段在步长=8时获得了更好的效果.google - deeplab2. 语义分割实现2.1. 分割输出解码函数：将 2D 分割图像转换为 RGB 图像，其中每一个标签被映射到对应的颜色.import numpy as np def decode_segmap(image, nc=21): label_colors = np.array([(0, 0, 0), # 0=background # 1=aeroplane, 2=bicycle, 3=bird, 4=boat, 5=bottle (128, 0, 0), (0, 128, 0), (255, 255, 255), (0, 0, 128), (128, 0, 128), # 6=bus, 7=car, 8=cat, 9=chair, 10=cow (0, 128, 128), (128, 128, 128), (64, 0, 0), (192, 0, 0), (64, 128, 0), # 11=dining table, 12=dog, 13=horse, 14=motorbike, 15=person (192, 128, 0), (64, 0, 128), (192, 0, 128), (64, 128, 128), (255, 255, 255), # 16=potted plant, 17=sheep, 18=sofa, 19=train, 20=tv/monitor (0, 64, 0), (128, 64, 0), (0, 192, 0), (128, 192, 0), (0, 64, 128)]) r = np.zeros_like(image).astype(np.uint8) g = np.zeros_like(image).astype(np.uint8) b = np.zeros_like(image).astype(np.uint8) for l in range(0, nc): idx = image == l r[idx] = label_colors[l, 0] g[idx] = label_colors[l, 1] b[idx] = label_colors[l, 2] rgb = np.stack([r, g, b], axis=2) return rgb2.2. 图像预处理与Forwardimport matplotlib.pyplot as plt import torchvision.transforms as T import torch from PIL import Image def segment(net, path, show_orig=True, dev='cuda'): img = Image.open(path) if show_orig: plt.imshow(img); plt.axis('off'); plt.show() # Comment the Resize and CenterCrop for better inference results trf = T.Compose([T.ToTensor(), T.Normalize(mean = [0.485, 0.456, 0.406], std = [0.229, 0.224, 0.225])]) inp = trf(img).unsqueeze(0).to(dev) out = net.to(dev)(inp)['out'] om = torch.argmax(out.squeeze(), dim=0).detach().cpu().numpy() rgb = decode_segmap(om) return rgb2.3. 加载 deeplabv3_resnet101 模型from torchvision import models dlab = models.segmentation.deeplabv3_resnet101(pretrained=1).eval() rgb=segment(dlab, '/path/to/test.jpg') ## If there are multiple labeled objects in the image, use the below code to have only the target as the foreground rgb[rgb!=255]=03. 背景融合(Alpha blending)实现Alpha blending 是将前景图像和透明图覆盖到背景图像上的处理. 透明图通常是图像的第四个通道(如，透明 PNG)，但其也可以是单独的图像. 透明 mask 图也一般被称为 alpha mask 或 alpha matte. Alpha blending 背后的数据原理是比较直接的，对于图像的每个像素，需要采用 alpha mask 将前景图像颜色和背景图像颜色组合在一起.前景图像是原始图像，灰度 alpha mask 是有语义分割得到的结果，背景图像是需要更换的新的定制图像.#alpha blending to customize the background of the image # Read the images foreground = cv2.imread("/path/to/foreground.jpg") background = cv2.imread("/path/to/background.jpg", cv2.IMREAD_COLOR) background = cv2.resize(background, (rgb.shape[1],rgb.shape[0]), interpolation = cv2.INTER_AREA) alpha = rgb # 2.3 # Convert uint8 to float foreground = foreground.astype(float) background = background.astype(float) # Normalize the alpha mask to keep intensity between 0 and 1 alpha = alpha.astype(float)/255 # Multiply the foreground with the alpha matte foreground = cv2.multiply(alpha, foreground) # Multiply the background with ( 1 - alpha ) background = cv2.multiply(1.0 - alpha, background) # Add the masked foreground and background. outImage = cv2.add(foreground, background) # Display image cv2_imshow(outImage) cv2.waitKey(0) # Save/download image cv2.imwrite('org_plus_cust_bkg_img.png' , outImage)如：图：语义分割和 alpha blending 的背景替换结果图：语义分割和 alpha blending 的处理过程4. 背景白化(Whiten)#whiten the background of the image mask_out=cv2.subtract(rgb,img) mask_out=cv2.subtract(rgb,mask_out) mask_out[rgb == 0] = 255 # Display the result numpy_horizontal = np.hstack((img, mask_out)) numpy_horizontal_concat = np.concatenate((img, mask_out), axis=1) cv2_imshow(numpy_horizontal_concat) cv2.waitKey(0) # Save/download the resulting image cv2.imwrite('org_plus_white_bkg_image.jpeg',numpy_horizontal_concat)如：图：图像背景白化结果5. 背景模糊(Blur)# Read the images foreground = cv2.imread("/path/to/girl8.jpg") # Create a Gaussian blur of kernel size 7 for the background image blurredImage = cv2.GaussianBlur(foreground, (7,7), 0) # Convert uint8 to float foreground = foreground.astype(float) blurredImage = blurredImage.astype(float) # Create a binary mask of the RGB output map using the threshold value 0 th, alpha = cv2.threshold(np.array(rgb),0, 255, cv2.THRESH_BINARY) # Apply a slight blur to the mask to soften edges alpha = cv2.GaussianBlur(alpha, (7,7),0) # Normalize the alpha mask to keep intensity between 0 and 1 alpha = alpha.astype(float)/255 # Multiply the foreground with the alpha matte foreground = cv2.multiply(alpha, foreground) # Multiply the background with ( 1 - alpha ) background = cv2.multiply(1.0 - alpha, blurredImage) # Add the masked foreground and background outImage = cv2.add(foreground, background) # Return a normalized output image for display outImage= outImage numpy_horizontal = np.hstack((img, outImage)) numpy_horizontal_concat = np.concatenate((img, outImage), axis=1) # Display image cv2_imshow(numpy_horizontal_concat) cv2.waitKey(0) # Save/download the resulting image cv2.imwrite('res.png' , numpy_horizontal_concat)如：图：图像背景融合结果6. 背景灰度化(Grayscale)# Load the foreground input image foreground = cv2.imread("/path/to/woman-with-car.jpg") # Resize image to match shape of R-band in RGB output map foreground = cv2.resize(foreground, (rgb.shape[1],rgb.shape[0]), interpolation = cv2.INTER_AREA) # Create a background image by copying foreground and converting into grayscale background = cv2.cvtColor(foreground, cv2.COLOR_BGR2GRAY) # convert single channel grayscale image to 3-channel grayscale image background = cv2.cvtColor(background, cv2.COLOR_GRAY2RGB) # Convert uint8 to float foreground = foreground.astype(float) background = background.astype(float) # Create a binary mask of the RGB output map using the threshold value 0 th, alpha = cv2.threshold(np.array(rgb),0,255, cv2.THRESH_BINARY) # Apply a slight blur to the mask to soften edges alpha = cv2.GaussianBlur(alpha, (7,7),0) # Normalize the alpha mask to keep intensity between 0 and 1 alpha = alpha.astype(float)/255 # Multiply the foreground with the alpha matte foreground = cv2.multiply(alpha, foreground) # Multiply the background with ( 1 - alpha ) background = cv2.multiply(1.0 - alpha, background) # Add the masked foreground and background outImage = cv2.add(foreground, background) numpy_horizontal = np.hstack((img, outImage)) numpy_horizontal_concat = np.concatenate((img, outImage), axis=1) # Display image cv2_imshow(numpy_horizontal_concat) cv2.waitKey(0) # Save image cv2.imwrite('res.png' , numpy_horizontal_concat)如：图：背景灰度化结果图：背景灰度化处理过程参考Alpha Blending — https://www.learnopencv.com/alpha-blending-using-opencv-cpp-python/Semantic Segmentation — https://www.learnopencv.com/pytorch-for-beginners-semantic-segmentation-using-torchvision/DeepLab — https://github.com/tensorflow/models/tree/master/research/deeplabhttps://www.learnopencv.com/applications-of-foreground-background-separation-with-semantic-segmentation/
原文：达摩院视觉算法黑科技之透明抠图 -2020.07.13出处：阿里云视觉智能开放平台论文：Segmenting Transparent Objects in the Wild - 2020Github - Segment_Transparent_Objects简介： 透明抠图问题作为抠图问题的一种，其采用的方法和模型构建与通常的模型有所不同，透明抠图需要将环境光、折射率的影响纳入计算，而一般的折射光图又很难获得，因此透明抠图的模型在过去一直难以建立，或者说很难达到令人满意的效果，达摩院视觉算法团队通过双分支解码器（Object Mask获取，Opacity预测），颜色纠正模块，对图像实现高精度透明抠图.1. 透明抠图 vs 非透明抠图物体的抠图问题可以定义为求解以下的公式，即给定图像I，求解前景颜色F、背景颜色B和Alpha matte的线性组合:$$ I = \alpha F + (1 - \alpha)B, \alpha \in [0, 1] $$对于透明物体而言，它展现在观察者眼中的颜色是由其前景颜色，背景颜色以及环境光线经过前景物体自我反射折射混合而成的，因此，它的公式会更加复杂一些：$$ I = \alpha F + (1 - \alpha)B + \Phi, \alpha \in [0, 1] $$$$ \Phi = \iint R(\omega \rightarrow p) E(\omega) d \omega dp $$Φ表示的是环境光的影响, 它是所有光线E(w)与反射率方程R乘积在所有点上的二重积分，求解很复杂[1]，导致实现精确的透明抠图是一个非常困难的问题. 因此，现有透明抠图研究的目标也是实现视觉感受“真实”的抠图而已，并非追求完全真实的抠图结果.2. 透明抠图-现有研究SOTA的Matting算法在同时提供原图和对应trimap的情况下，可以的实现对半透物体的处理（如下图为GCA-Matting[2]的效果），但tripmap在实际的图像的处理中难以获取，限制了这类算法其在业务中的使用. TOM-Net[3] 将透明抠图问题视为折射流的估计问题，网络支持对单图输入,经过三分支的编解码器网络，分别预测图像的Object Mask，attenuative mask，flow mask(折射流图)，并可以通过折射流信息进一步在新的背景进行合成. 该方法的局限性在于其假设物体必须全部为无色透明物体，并且在训练过程中需要折射流图作为label, 而折射流图在真实世界是非常难以获取的，因此该方法的训练数只能依赖于图形学合成，与真实透明图像的分布无法一致（图像的语义合理性存疑，例如玻璃杯在山前悬浮）. 经过我们在实际数据上的测试，该方法在实际图像的表现并不理想. Segmenting Transparent Objects in the Wild[4]提出了基于语义分支和边缘分支结构的真实世界透明物体分割网络，通过边界注意力模块（Boundary Attention Modeule）增强对透明物体的分割精度，并发布了目前数量最大的透明物体分割标注数据集Trans10K. 然而，文章提出的算法和发布的数据集都是处理到语义分割层面，并没有对物体的透明度做进一步处理. 3. 透明抠图-问题简化考虑到透明抠图问题本身难以求解，而且数据构建也非常困难，在实际的应用场景中，为保证同时保证算法的泛化能力和抠图效果，我们对问题进行了简化， 我们假设所需处理的物体的透明部分是无色的，且所在环境的背景颜色分布相对均匀. 在这样的条件下，背景的自发光或反射光的颜色可以认为是全局一致的颜色，不会出现多种颜色叠加的情况，Φ的估计就只是和背景颜色相关了. 特别的，如果预知背景的颜色，可以通过将其作为先验引入Φ中，对结果进行背景杂色的抑制及去除.4. 透明抠图-模型设计我们的模型输入为单张图像，首先提取其深层特征. 然后进行语义级别的分割和提取，力求完整准确地获得物体所在图像区域（Mask）. 同步进行对图像物体不透明度(Opacity)的预测. 而不透明度(Opacity)的预测由于在训练的时候没有进行语义的约束，容易存在非主体区域的噪声影响，因此，将两者进行融合可以将透明信息约束在主体范围内. 最后，对于已知背景颜色先验的场景（如已知是绿幕），我们可以引入颜色纠正模块，实现对背景透出的杂色进行去除. 对于背景颜色未知，但饱和度低的场景，抠图结果也依然可用. 5. 透明抠图-结果与应用在已知背景颜色先验的情况下，可以通过颜色纠正模块对背景透出的杂色进行去除（左到右：实拍图，Opacity, 直接抠图结果，色偏纠正结果）对于背景颜色未知，但饱和度低的场景，抠图结果也依然可用，例如下面几组展示图. 目前在车辆分割算法上，我们已经实现了基于透明抠图的思路用于改善半透车窗区域的效果，使得车辆经过抠图，能够更加自然和谐与新背景进行融合. 目前车辆分割已经上线阿里云视觉智能开放平台（vision.aliyun.com），欢迎体验试用. 6. 总结与展望目前的透明抠图算法，面对更为多样的真实场景下物体，仍然具有以下不足，需要进一步探索解决：1.真实透明物体图像数据量严重不足，且标注困难；2.主体不透明度图（Opacity）的预测容易受到图像中噪声的影响；3.在得到Opacity图后，在未知背景先验，且背景颜色饱和度高的情况下，如何实现将背景色的去除（如下图的杯子整体泛蓝);后续会考虑进一步提取背景的特征，将背景先验知识引入到透明度的估计上，增加RGB偏移输出信息，尝试对前景物体的颜色进行纠正.参考文献[1] Environment matting and compositing[2] Natural Image Matting via Guided Contextual Attention[3] Learning Transparent Object Matting[4] Segmenting Transparent Objects in the Wild
论文：U^2-Net: Going Deeper with Nested U-Structure for Salient Object Detection - 2020作者：Xuebin Qin， Zichen Zhang， Chenyang Huang， Masood Dehghan， Osmar R.Zaiane， MartinJagersand团队：University of Alberta, Edmonton, CanadaGithub - U-2-NetU2Net，即U^2-Net(U square Net).1. U2Net 网络结构U2Net 网络结构如下图，其类似于编码-解码(Encoder-Decoder)结构的 U-Net. U2Net 每个 stage 由新提出的 RSU模块(residual U-block) 组成. 例如，En_1 即为基于 RSU 构建的.其优势在于：[1] - 提出 RSU 模块，融合不同尺寸接受野的特征，以捕获更多不同尺度的上下文信息(contextual information).[2] - 基于 RSU 模块的 池化(pooling) 操作，在不显著增加计算成本的前提下，增加了整个网络结构的深度(depth).1.1. RSU 模块RSU, ReSidual Ublock, 用于捕获 intra-stage 的 multi-scales 特征. 其结构如图：图: (a) Plain CNN, PLN. (b) Residual-like block, RES. (c) Dense-like block, DSE. (d) Inception-like block, INC. (e) Residual U-block, RSU.RSU 模块和 residual block 对比：注：PPM, pyramdi pooling modules.1.2. 损失函数类似于论文 Holistically-nested edge detection(HED)，U2Net 训练的损失函数定义为：$$ \mathcal{L} = \sum_{m=1}^M w_{side}^{(m)} \mathcal{l}_{side}^{(m)} + w_{fuse}\mathcal{l}_{fuse} $$其中，M=6, 为U2Net 的 Sup1, Sup2, ..., Sup6 stage. $\mathcal{l}_{side}^{(m)}$ 为 Sup1, Sup2, ..., Sup6 输出的显著图(saliency map) $S_{side}^{(m)}$ 的损失函数.$\mathcal{l} _{fuse}$ 为最终融合输出的显著图$S_{fuse}$ 的损失函数.$w_{fuse}$ 为每个损失函数的权重.对于每个 $\mathcal{l}$，采用标准的二值交叉熵损失函数：$$ \mathcal{l} = - \sum_{(r, c)}^{(H, W)}[P_{G_{(r, c)}} log P_{S(r, c)} + (1 - P_{G(r, c)}) log(1 - P_{S(r, c)})] $$其中，$(r, c)$ 为像素坐标；(H, W) 为图像尺寸，height 和 width.$P_{G(r, c)}$ 和 $P_{S(r, c)}$ 分别表示 GT 像素值和预测的显著概率图(saliency probability map).2. U2Net 输出结果可视化U2Net的输出结果，如图：3. U2Net 实现Github - U-2-Net3.1. U2Net 和 U2NetPu2net.py:import torch import torch.nn as nn from torchvision import models import torch.nn.functional as F class REBNCONV(nn.Module): def __init__(self,in_ch=3,out_ch=3,dirate=1): super(REBNCONV,self).__init__() self.conv_s1 = nn.Conv2d(in_ch,out_ch,3,padding=1*dirate,dilation=1*dirate) self.bn_s1 = nn.BatchNorm2d(out_ch) self.relu_s1 = nn.ReLU(inplace=True) def forward(self,x): hx = x xout = self.relu_s1(self.bn_s1(self.conv_s1(hx))) return xout ## upsample tensor 'src' to have the same spatial size with tensor 'tar' def _upsample_like(src,tar): src = F.upsample(src,size=tar.shape[2:],mode='bilinear') return src ### RSU-7 ### class RSU7(nn.Module):#UNet07DRES(nn.Module): def __init__(self, in_ch=3, mid_ch=12, out_ch=3): super(RSU7,self).__init__() self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1) self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1) self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True) self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1) self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True) self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1) self.pool3 = nn.MaxPool2d(2,stride=2,ceil_mode=True) self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1) self.pool4 = nn.MaxPool2d(2,stride=2,ceil_mode=True) self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=1) self.pool5 = nn.MaxPool2d(2,stride=2,ceil_mode=True) self.rebnconv6 = REBNCONV(mid_ch,mid_ch,dirate=1) self.rebnconv7 = REBNCONV(mid_ch,mid_ch,dirate=2) self.rebnconv6d = REBNCONV(mid_ch*2,mid_ch,dirate=1) self.rebnconv5d = REBNCONV(mid_ch*2,mid_ch,dirate=1) self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1) self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1) self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1) self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1) def forward(self,x): hx = x hxin = self.rebnconvin(hx) hx1 = self.rebnconv1(hxin) hx = self.pool1(hx1) hx2 = self.rebnconv2(hx) hx = self.pool2(hx2) hx3 = self.rebnconv3(hx) hx = self.pool3(hx3) hx4 = self.rebnconv4(hx) hx = self.pool4(hx4) hx5 = self.rebnconv5(hx) hx = self.pool5(hx5) hx6 = self.rebnconv6(hx) hx7 = self.rebnconv7(hx6) hx6d = self.rebnconv6d(torch.cat((hx7,hx6),1)) hx6dup = _upsample_like(hx6d,hx5) hx5d = self.rebnconv5d(torch.cat((hx6dup,hx5),1)) hx5dup = _upsample_like(hx5d,hx4) hx4d = self.rebnconv4d(torch.cat((hx5dup,hx4),1)) hx4dup = _upsample_like(hx4d,hx3) hx3d = self.rebnconv3d(torch.cat((hx4dup,hx3),1)) hx3dup = _upsample_like(hx3d,hx2) hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1)) hx2dup = _upsample_like(hx2d,hx1) hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1)) return hx1d + hxin ### RSU-6 ### class RSU6(nn.Module):#UNet06DRES(nn.Module): def __init__(self, in_ch=3, mid_ch=12, out_ch=3): super(RSU6,self).__init__() self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1) self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1) self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True) self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1) self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True) self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1) self.pool3 = nn.MaxPool2d(2,stride=2,ceil_mode=True) self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1) self.pool4 = nn.MaxPool2d(2,stride=2,ceil_mode=True) self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=1) self.rebnconv6 = REBNCONV(mid_ch,mid_ch,dirate=2) self.rebnconv5d = REBNCONV(mid_ch*2,mid_ch,dirate=1) self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1) self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1) self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1) self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1) def forward(self,x): hx = x hxin = self.rebnconvin(hx) hx1 = self.rebnconv1(hxin) hx = self.pool1(hx1) hx2 = self.rebnconv2(hx) hx = self.pool2(hx2) hx3 = self.rebnconv3(hx) hx = self.pool3(hx3) hx4 = self.rebnconv4(hx) hx = self.pool4(hx4) hx5 = self.rebnconv5(hx) hx6 = self.rebnconv6(hx5) hx5d = self.rebnconv5d(torch.cat((hx6,hx5),1)) hx5dup = _upsample_like(hx5d,hx4) hx4d = self.rebnconv4d(torch.cat((hx5dup,hx4),1)) hx4dup = _upsample_like(hx4d,hx3) hx3d = self.rebnconv3d(torch.cat((hx4dup,hx3),1)) hx3dup = _upsample_like(hx3d,hx2) hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1)) hx2dup = _upsample_like(hx2d,hx1) hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1)) return hx1d + hxin ### RSU-5 ### class RSU5(nn.Module):#UNet05DRES(nn.Module): def __init__(self, in_ch=3, mid_ch=12, out_ch=3): super(RSU5,self).__init__() self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1) self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1) self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True) self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1) self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True) self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1) self.pool3 = nn.MaxPool2d(2,stride=2,ceil_mode=True) self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1) self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=2) self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1) self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1) self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1) self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1) def forward(self,x): hx = x hxin = self.rebnconvin(hx) hx1 = self.rebnconv1(hxin) hx = self.pool1(hx1) hx2 = self.rebnconv2(hx) hx = self.pool2(hx2) hx3 = self.rebnconv3(hx) hx = self.pool3(hx3) hx4 = self.rebnconv4(hx) hx5 = self.rebnconv5(hx4) hx4d = self.rebnconv4d(torch.cat((hx5,hx4),1)) hx4dup = _upsample_like(hx4d,hx3) hx3d = self.rebnconv3d(torch.cat((hx4dup,hx3),1)) hx3dup = _upsample_like(hx3d,hx2) hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1)) hx2dup = _upsample_like(hx2d,hx1) hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1)) return hx1d + hxin ### RSU-4 ### class RSU4(nn.Module):#UNet04DRES(nn.Module): def __init__(self, in_ch=3, mid_ch=12, out_ch=3): super(RSU4,self).__init__() self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1) self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1) self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True) self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1) self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True) self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1) self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=2) self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1) self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1) self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1) def forward(self,x): hx = x hxin = self.rebnconvin(hx) hx1 = self.rebnconv1(hxin) hx = self.pool1(hx1) hx2 = self.rebnconv2(hx) hx = self.pool2(hx2) hx3 = self.rebnconv3(hx) hx4 = self.rebnconv4(hx3) hx3d = self.rebnconv3d(torch.cat((hx4,hx3),1)) hx3dup = _upsample_like(hx3d,hx2) hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1)) hx2dup = _upsample_like(hx2d,hx1) hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1)) return hx1d + hxin ### RSU-4F ### class RSU4F(nn.Module):#UNet04FRES(nn.Module): def __init__(self, in_ch=3, mid_ch=12, out_ch=3): super(RSU4F,self).__init__() self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1) self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1) self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=2) self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=4) self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=8) self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=4) self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=2) self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1) def forward(self,x): hx = x hxin = self.rebnconvin(hx) hx1 = self.rebnconv1(hxin) hx2 = self.rebnconv2(hx1) hx3 = self.rebnconv3(hx2) hx4 = self.rebnconv4(hx3) hx3d = self.rebnconv3d(torch.cat((hx4,hx3),1)) hx2d = self.rebnconv2d(torch.cat((hx3d,hx2),1)) hx1d = self.rebnconv1d(torch.cat((hx2d,hx1),1)) return hx1d + hxin ##### U^2-Net #### class U2NET(nn.Module): def __init__(self,in_ch=3,out_ch=1): super(U2NET,self).__init__() self.stage1 = RSU7(in_ch,32,64) self.pool12 = nn.MaxPool2d(2,stride=2,ceil_mode=True) self.stage2 = RSU6(64,32,128) self.pool23 = nn.MaxPool2d(2,stride=2,ceil_mode=True) self.stage3 = RSU5(128,64,256) self.pool34 = nn.MaxPool2d(2,stride=2,ceil_mode=True) self.stage4 = RSU4(256,128,512) self.pool45 = nn.MaxPool2d(2,stride=2,ceil_mode=True) self.stage5 = RSU4F(512,256,512) self.pool56 = nn.MaxPool2d(2,stride=2,ceil_mode=True) self.stage6 = RSU4F(512,256,512) # decoder self.stage5d = RSU4F(1024,256,512) self.stage4d = RSU4(1024,128,256) self.stage3d = RSU5(512,64,128) self.stage2d = RSU6(256,32,64) self.stage1d = RSU7(128,16,64) self.side1 = nn.Conv2d(64,out_ch,3,padding=1) self.side2 = nn.Conv2d(64,out_ch,3,padding=1) self.side3 = nn.Conv2d(128,out_ch,3,padding=1) self.side4 = nn.Conv2d(256,out_ch,3,padding=1) self.side5 = nn.Conv2d(512,out_ch,3,padding=1) self.side6 = nn.Conv2d(512,out_ch,3,padding=1) self.outconv = nn.Conv2d(6,out_ch,1) def forward(self,x): hx = x #stage 1 hx1 = self.stage1(hx) hx = self.pool12(hx1) #stage 2 hx2 = self.stage2(hx) hx = self.pool23(hx2) #stage 3 hx3 = self.stage3(hx) hx = self.pool34(hx3) #stage 4 hx4 = self.stage4(hx) hx = self.pool45(hx4) #stage 5 hx5 = self.stage5(hx) hx = self.pool56(hx5) #stage 6 hx6 = self.stage6(hx) hx6up = _upsample_like(hx6,hx5) #-------------------- decoder -------------------- hx5d = self.stage5d(torch.cat((hx6up,hx5),1)) hx5dup = _upsample_like(hx5d,hx4) hx4d = self.stage4d(torch.cat((hx5dup,hx4),1)) hx4dup = _upsample_like(hx4d,hx3) hx3d = self.stage3d(torch.cat((hx4dup,hx3),1)) hx3dup = _upsample_like(hx3d,hx2) hx2d = self.stage2d(torch.cat((hx3dup,hx2),1)) hx2dup = _upsample_like(hx2d,hx1) hx1d = self.stage1d(torch.cat((hx2dup,hx1),1)) #side output d1 = self.side1(hx1d) d2 = self.side2(hx2d) d2 = _upsample_like(d2,d1) d3 = self.side3(hx3d) d3 = _upsample_like(d3,d1) d4 = self.side4(hx4d) d4 = _upsample_like(d4,d1) d5 = self.side5(hx5d) d5 = _upsample_like(d5,d1) d6 = self.side6(hx6) d6 = _upsample_like(d6,d1) d0 = self.outconv(torch.cat((d1,d2,d3,d4,d5,d6),1)) return F.sigmoid(d0), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6) ### U^2-Net small ### class U2NETP(nn.Module): def __init__(self,in_ch=3,out_ch=1): super(U2NETP,self).__init__() self.stage1 = RSU7(in_ch,16,64) self.pool12 = nn.MaxPool2d(2,stride=2,ceil_mode=True) self.stage2 = RSU6(64,16,64) self.pool23 = nn.MaxPool2d(2,stride=2,ceil_mode=True) self.stage3 = RSU5(64,16,64) self.pool34 = nn.MaxPool2d(2,stride=2,ceil_mode=True) self.stage4 = RSU4(64,16,64) self.pool45 = nn.MaxPool2d(2,stride=2,ceil_mode=True) self.stage5 = RSU4F(64,16,64) self.pool56 = nn.MaxPool2d(2,stride=2,ceil_mode=True) self.stage6 = RSU4F(64,16,64) # decoder self.stage5d = RSU4F(128,16,64) self.stage4d = RSU4(128,16,64) self.stage3d = RSU5(128,16,64) self.stage2d = RSU6(128,16,64) self.stage1d = RSU7(128,16,64) self.side1 = nn.Conv2d(64,out_ch,3,padding=1) self.side2 = nn.Conv2d(64,out_ch,3,padding=1) self.side3 = nn.Conv2d(64,out_ch,3,padding=1) self.side4 = nn.Conv2d(64,out_ch,3,padding=1) self.side5 = nn.Conv2d(64,out_ch,3,padding=1) self.side6 = nn.Conv2d(64,out_ch,3,padding=1) self.outconv = nn.Conv2d(6,out_ch,1) def forward(self,x): hx = x #stage 1 hx1 = self.stage1(hx) hx = self.pool12(hx1) #stage 2 hx2 = self.stage2(hx) hx = self.pool23(hx2) #stage 3 hx3 = self.stage3(hx) hx = self.pool34(hx3) #stage 4 hx4 = self.stage4(hx) hx = self.pool45(hx4) #stage 5 hx5 = self.stage5(hx) hx = self.pool56(hx5) #stage 6 hx6 = self.stage6(hx) hx6up = _upsample_like(hx6,hx5) #decoder hx5d = self.stage5d(torch.cat((hx6up,hx5),1)) hx5dup = _upsample_like(hx5d,hx4) hx4d = self.stage4d(torch.cat((hx5dup,hx4),1)) hx4dup = _upsample_like(hx4d,hx3) hx3d = self.stage3d(torch.cat((hx4dup,hx3),1)) hx3dup = _upsample_like(hx3d,hx2) hx2d = self.stage2d(torch.cat((hx3dup,hx2),1)) hx2dup = _upsample_like(hx2d,hx1) hx1d = self.stage1d(torch.cat((hx2dup,hx1),1)) #side output d1 = self.side1(hx1d) d2 = self.side2(hx2d) d2 = _upsample_like(d2,d1) d3 = self.side3(hx3d) d3 = _upsample_like(d3,d1) d4 = self.side4(hx4d) d4 = _upsample_like(d4,d1) d5 = self.side5(hx5d) d5 = _upsample_like(d5,d1) d6 = self.side6(hx6) d6 = _upsample_like(d6,d1) d0 = self.outconv(torch.cat((d1,d2,d3,d4,d5,d6),1)) return F.sigmoid(d0), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)3.2. DataLoaderdata_loader.py:# data loader from __future__ import print_function, division import glob import torch from skimage import io, transform, color import numpy as np import random import math import matplotlib.pyplot as plt from torch.utils.data import Dataset, DataLoader from torchvision import transforms, utils from PIL import Image #==== dataset load==== class RescaleT(object): def __init__(self,output_size): assert isinstance(output_size,(int,tuple)) self.output_size = output_size def __call__(self,sample): imidx, image, label = sample['imidx'], sample['image'],sample['label'] h, w = image.shape[:2] if isinstance(self.output_size,int): if h > w: new_h, new_w = self.output_size*h/w,self.output_size else: new_h, new_w = self.output_size,self.output_size*w/h else: new_h, new_w = self.output_size new_h, new_w = int(new_h), int(new_w) # #resize the image to new_h x new_w and convert image from range [0,255] to [0,1] # img = transform.resize(image,(new_h,new_w),mode='constant') # lbl = transform.resize(label,(new_h,new_w),mode='constant', order=0, preserve_range=True) img = transform.resize(image,(self.output_size,self.output_size),mode='constant') lbl = transform.resize(label,(self.output_size,self.output_size),mode='constant', order=0, preserve_range=True) return {'imidx':imidx, 'image':img,'label':lbl} class Rescale(object): def __init__(self,output_size): assert isinstance(output_size,(int,tuple)) self.output_size = output_size def __call__(self,sample): imidx, image, label = sample['imidx'], sample['image'],sample['label'] if random.random() >= 0.5: image = image[::-1] label = label[::-1] h, w = image.shape[:2] if isinstance(self.output_size,int): if h > w: new_h, new_w = self.output_size*h/w,self.output_size else: new_h, new_w = self.output_size,self.output_size*w/h else: new_h, new_w = self.output_size new_h, new_w = int(new_h), int(new_w) # #resize the image to new_h x new_w and convert image from range [0,255] to [0,1] img = transform.resize(image,(new_h,new_w),mode='constant') lbl = transform.resize(label,(new_h,new_w),mode='constant', order=0, preserve_range=True) return {'imidx':imidx, 'image':img,'label':lbl} class RandomCrop(object): def __init__(self,output_size): assert isinstance(output_size, (int, tuple)) if isinstance(output_size, int): self.output_size = (output_size, output_size) else: assert len(output_size) == 2 self.output_size = output_size def __call__(self,sample): imidx, image, label = sample['imidx'], sample['image'], sample['label'] if random.random() >= 0.5: image = image[::-1] label = label[::-1] h, w = image.shape[:2] new_h, new_w = self.output_size top = np.random.randint(0, h - new_h) left = np.random.randint(0, w - new_w) image = image[top: top + new_h, left: left + new_w] label = label[top: top + new_h, left: left + new_w] return {'imidx':imidx,'image':image, 'label':label} class ToTensor(object): """Convert ndarrays in sample to Tensors.""" def __call__(self, sample): imidx, image, label = sample['imidx'], sample['image'], sample['label'] tmpImg = np.zeros((image.shape[0],image.shape[1],3)) tmpLbl = np.zeros(label.shape) image = image/np.max(image) if(np.max(label)<1e-6): label = label else: label = label/np.max(label) if image.shape[2]==1: tmpImg[:,:,0] = (image[:,:,0]-0.485)/0.229 tmpImg[:,:,1] = (image[:,:,0]-0.485)/0.229 tmpImg[:,:,2] = (image[:,:,0]-0.485)/0.229 else: tmpImg[:,:,0] = (image[:,:,0]-0.485)/0.229 tmpImg[:,:,1] = (image[:,:,1]-0.456)/0.224 tmpImg[:,:,2] = (image[:,:,2]-0.406)/0.225 tmpLbl[:,:,0] = label[:,:,0] # change the r,g,b to b,r,g from [0,255] to [0,1] #transforms.Normalize(mean = (0.485, 0.456, 0.406), std = (0.229, 0.224, 0.225)) tmpImg = tmpImg.transpose((2, 0, 1)) tmpLbl = label.transpose((2, 0, 1)) return {'imidx':torch.from_numpy(imidx), 'image': torch.from_numpy(tmpImg), 'label': torch.from_numpy(tmpLbl)} class ToTensorLab(object): """Convert ndarrays in sample to Tensors.""" def __init__(self,flag=0): self.flag = flag def __call__(self, sample): imidx, image, label =sample['imidx'], sample['image'], sample['label'] tmpLbl = np.zeros(label.shape) if(np.max(label)<1e-6): label = label else: label = label/np.max(label) # change the color space if self.flag == 2: # with rgb and Lab colors tmpImg = np.zeros((image.shape[0],image.shape[1],6)) tmpImgt = np.zeros((image.shape[0],image.shape[1],3)) if image.shape[2]==1: tmpImgt[:,:,0] = image[:,:,0] tmpImgt[:,:,1] = image[:,:,0] tmpImgt[:,:,2] = image[:,:,0] else: tmpImgt = image tmpImgtl = color.rgb2lab(tmpImgt) # nomalize image to range [0,1] tmpImg[:,:,0] = (tmpImgt[:,:,0]-np.min(tmpImgt[:,:,0]))/(np.max(tmpImgt[:,:,0])-np.min(tmpImgt[:,:,0])) tmpImg[:,:,1] = (tmpImgt[:,:,1]-np.min(tmpImgt[:,:,1]))/(np.max(tmpImgt[:,:,1])-np.min(tmpImgt[:,:,1])) tmpImg[:,:,2] = (tmpImgt[:,:,2]-np.min(tmpImgt[:,:,2]))/(np.max(tmpImgt[:,:,2])-np.min(tmpImgt[:,:,2])) tmpImg[:,:,3] = (tmpImgtl[:,:,0]-np.min(tmpImgtl[:,:,0]))/(np.max(tmpImgtl[:,:,0])-np.min(tmpImgtl[:,:,0])) tmpImg[:,:,4] = (tmpImgtl[:,:,1]-np.min(tmpImgtl[:,:,1]))/(np.max(tmpImgtl[:,:,1])-np.min(tmpImgtl[:,:,1])) tmpImg[:,:,5] = (tmpImgtl[:,:,2]-np.min(tmpImgtl[:,:,2]))/(np.max(tmpImgtl[:,:,2])-np.min(tmpImgtl[:,:,2])) # tmpImg = tmpImg/(np.max(tmpImg)-np.min(tmpImg)) tmpImg[:,:,0] = (tmpImg[:,:,0]-np.mean(tmpImg[:,:,0]))/np.std(tmpImg[:,:,0]) tmpImg[:,:,1] = (tmpImg[:,:,1]-np.mean(tmpImg[:,:,1]))/np.std(tmpImg[:,:,1]) tmpImg[:,:,2] = (tmpImg[:,:,2]-np.mean(tmpImg[:,:,2]))/np.std(tmpImg[:,:,2]) tmpImg[:,:,3] = (tmpImg[:,:,3]-np.mean(tmpImg[:,:,3]))/np.std(tmpImg[:,:,3]) tmpImg[:,:,4] = (tmpImg[:,:,4]-np.mean(tmpImg[:,:,4]))/np.std(tmpImg[:,:,4]) tmpImg[:,:,5] = (tmpImg[:,:,5]-np.mean(tmpImg[:,:,5]))/np.std(tmpImg[:,:,5]) elif self.flag == 1: #with Lab color tmpImg = np.zeros((image.shape[0],image.shape[1],3)) if image.shape[2]==1: tmpImg[:,:,0] = image[:,:,0] tmpImg[:,:,1] = image[:,:,0] tmpImg[:,:,2] = image[:,:,0] else: tmpImg = image tmpImg = color.rgb2lab(tmpImg) # tmpImg = tmpImg/(np.max(tmpImg)-np.min(tmpImg)) tmpImg[:,:,0] = (tmpImg[:,:,0]-np.min(tmpImg[:,:,0]))/(np.max(tmpImg[:,:,0])-np.min(tmpImg[:,:,0])) tmpImg[:,:,1] = (tmpImg[:,:,1]-np.min(tmpImg[:,:,1]))/(np.max(tmpImg[:,:,1])-np.min(tmpImg[:,:,1])) tmpImg[:,:,2] = (tmpImg[:,:,2]-np.min(tmpImg[:,:,2]))/(np.max(tmpImg[:,:,2])-np.min(tmpImg[:,:,2])) tmpImg[:,:,0] = (tmpImg[:,:,0]-np.mean(tmpImg[:,:,0]))/np.std(tmpImg[:,:,0]) tmpImg[:,:,1] = (tmpImg[:,:,1]-np.mean(tmpImg[:,:,1]))/np.std(tmpImg[:,:,1]) tmpImg[:,:,2] = (tmpImg[:,:,2]-np.mean(tmpImg[:,:,2]))/np.std(tmpImg[:,:,2]) else: # with rgb color tmpImg = np.zeros((image.shape[0],image.shape[1],3)) image = image/np.max(image) if image.shape[2]==1: tmpImg[:,:,0] = (image[:,:,0]-0.485)/0.229 tmpImg[:,:,1] = (image[:,:,0]-0.485)/0.229 tmpImg[:,:,2] = (image[:,:,0]-0.485)/0.229 else: tmpImg[:,:,0] = (image[:,:,0]-0.485)/0.229 tmpImg[:,:,1] = (image[:,:,1]-0.456)/0.224 tmpImg[:,:,2] = (image[:,:,2]-0.406)/0.225 tmpLbl[:,:,0] = label[:,:,0] # change the r,g,b to b,r,g from [0,255] to [0,1] #transforms.Normalize(mean = (0.485, 0.456, 0.406), std = (0.229, 0.224, 0.225)) tmpImg = tmpImg.transpose((2, 0, 1)) tmpLbl = label.transpose((2, 0, 1)) return {'imidx':torch.from_numpy(imidx), 'image': torch.from_numpy(tmpImg), 'label': torch.from_numpy(tmpLbl)} class SalObjDataset(Dataset): def __init__(self,img_name_list,lbl_name_list,transform=None): # self.root_dir = root_dir # self.image_name_list = glob.glob(image_dir+'*.png') # self.label_name_list = glob.glob(label_dir+'*.png') self.image_name_list = img_name_list self.label_name_list = lbl_name_list self.transform = transform def __len__(self): return len(self.image_name_list) def __getitem__(self,idx): # image = Image.open(self.image_name_list[idx])#io.imread(self.image_name_list[idx]) # label = Image.open(self.label_name_list[idx])#io.imread(self.label_name_list[idx]) image = io.imread(self.image_name_list[idx]) imname = self.image_name_list[idx] imidx = np.array([idx]) if(0==len(self.label_name_list)): label_3 = np.zeros(image.shape) else: label_3 = io.imread(self.label_name_list[idx]) label = np.zeros(label_3.shape[0:2]) if(3==len(label_3.shape)): label = label_3[:,:,0] elif(2==len(label_3.shape)): label = label_3 if(3==len(image.shape) and 2==len(label.shape)): label = label[:,:,np.newaxis] elif(2==len(image.shape) and 2==len(label.shape)): image = image[:,:,np.newaxis] label = label[:,:,np.newaxis] sample = {'imidx':imidx, 'image':image, 'label':label} if self.transform: sample = self.transform(sample) return sample3.3. Trainu2net_train.py:import torch import torchvision from torch.autograd import Variable import torch.nn as nn import torch.nn.functional as F from torch.utils.data import Dataset, DataLoader from torchvision import transforms, utils import torch.optim as optim import torchvision.transforms as standard_transforms import numpy as np import glob from data_loader import Rescale from data_loader import RescaleT from data_loader import RandomCrop from data_loader import ToTensor from data_loader import ToTensorLab from data_loader import SalObjDataset from model import U2NET from model import U2NETP # ------- 1. define loss function -------- bce_loss = nn.BCELoss(size_average=True) def muti_bce_loss_fusion(d0, d1, d2, d3, d4, d5, d6, labels_v): loss0 = bce_loss(d0,labels_v) loss1 = bce_loss(d1,labels_v) loss2 = bce_loss(d2,labels_v) loss3 = bce_loss(d3,labels_v) loss4 = bce_loss(d4,labels_v) loss5 = bce_loss(d5,labels_v) loss6 = bce_loss(d6,labels_v) loss = loss0 + loss1 + loss2 + loss3 + loss4 + loss5 + loss6 print("l0: %3f, l1: %3f, l2: %3f, l3: %3f, l4: %3f, l5: %3f, l6: %3f\n"%(loss0.data[0],loss1.data[0],loss2.data[0],loss3.data[0],loss4.data[0],loss5.data[0],loss6.data[0])) return loss0, loss # ------- 2. set the directory of training dataset -------- model_name = 'u2net' #'u2netp' data_dir = os.path.join(os.getcwd(), 'train_data' + os.sep) tra_image_dir = os.path.join('DUTS', 'DUTS-TR', 'DUTS-TR', 'im_aug' + os.sep) tra_label_dir = os.path.join('DUTS', 'DUTS-TR', 'DUTS-TR', 'gt_aug' + os.sep) image_ext = '.jpg' label_ext = '.png' model_dir = os.path.join(os.getcwd(), 'saved_models', model_name + os.sep) epoch_num = 100000 batch_size_train = 12 batch_size_val = 1 train_num = 0 val_num = 0 tra_img_name_list = glob.glob(data_dir + tra_image_dir + '*' + image_ext) tra_lbl_name_list = [] for img_path in tra_img_name_list: img_name = img_path.split(os.sep)[-1] aaa = img_name.split(".") bbb = aaa[0:-1] imidx = bbb[0] for i in range(1,len(bbb)): imidx = imidx + "." + bbb[i] # tra_lbl_name_list.append(data_dir + tra_label_dir + imidx + label_ext) print("---") print("train images: ", len(tra_img_name_list)) print("train labels: ", len(tra_lbl_name_list)) print("---") train_num = len(tra_img_name_list) salobj_dataset = SalObjDataset( img_name_list=tra_img_name_list, lbl_name_list=tra_lbl_name_list, transform=transforms.Compose([ RescaleT(320), RandomCrop(288), ToTensorLab(flag=0)])) salobj_dataloader = DataLoader(salobj_dataset, batch_size=batch_size_train, shuffle=True, num_workers=1) # ------- 3. define model -------- # define the net if(model_name=='u2net'): net = U2NET(3, 1) elif(model_name=='u2netp'): net = U2NETP(3,1) if torch.cuda.is_available(): net.cuda() # ------- 4. define optimizer -------- print("---define optimizer...") optimizer = optim.Adam(net.parameters(), lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0) # ------- 5. training process -------- print("---start training...") ite_num = 0 running_loss = 0.0 running_tar_loss = 0.0 ite_num4val = 0 save_frq = 2000 # save the model every 2000 iterations for epoch in range(0, epoch_num): net.train() for i, data in enumerate(salobj_dataloader): ite_num = ite_num + 1 ite_num4val = ite_num4val + 1 inputs, labels = data['image'], data['label'] inputs = inputs.type(torch.FloatTensor) labels = labels.type(torch.FloatTensor) # wrap them in Variable if torch.cuda.is_available(): inputs_v, labels_v = Variable(inputs.cuda(), requires_grad=False), Variable(labels.cuda(), requires_grad=False) else: inputs_v, labels_v = Variable(inputs, requires_grad=False), Variable(labels, requires_grad=False) # y zero the parameter gradients optimizer.zero_grad() # forward + backward + optimize d0, d1, d2, d3, d4, d5, d6 = net(inputs_v) loss2, loss = muti_bce_loss_fusion(d0, d1, d2, d3, d4, d5, d6, labels_v) loss.backward() optimizer.step() # # print statistics running_loss += loss.data[0] running_tar_loss += loss2.data[0] # del temporary outputs and loss del d0, d1, d2, d3, d4, d5, d6, loss2, loss print("[epoch: %3d/%3d, batch: %5d/%5d, ite: %d] train loss: %3f, tar: %3f " % ( epoch + 1, epoch_num, (i + 1) * batch_size_train, train_num, ite_num, running_loss / ite_num4val, running_tar_loss / ite_num4val)) if ite_num % save_frq == 0: torch.save(net.state_dict(), model_dir + model_name+"_bce_itr_%d_train_%3f_tar_%3f.pth" % (ite_num, running_loss / ite_num4val, running_tar_loss / ite_num4val)) running_loss = 0.0 running_tar_loss = 0.0 net.train() # resume train ite_num4val = 0 if __name__ == "__main__": main()3.4. Testu2net_test.py:import os from skimage import io, transform import torch import torchvision from torch.autograd import Variable import torch.nn as nn import torch.nn.functional as F from torch.utils.data import Dataset, DataLoader from torchvision import transforms#, utils # import torch.optim as optim import numpy as np from PIL import Image import glob from data_loader import RescaleT from data_loader import ToTensor from data_loader import ToTensorLab from data_loader import SalObjDataset from model import U2NET # full size version 173.6 MB from model import U2NETP # small version u2net 4.7 MB # normalize the predicted SOD probability map def normPRED(d): # QA. ma = torch.max(d) mi = torch.min(d) dn = (d-mi)/(ma-mi) return dn def save_output(image_name,pred,d_dir): predict = pred predict = predict.squeeze() predict_np = predict.cpu().data.numpy() im = Image.fromarray(predict_np*255).convert('RGB') img_name = image_name.split(os.sep)[-1] image = io.imread(image_name) imo = im.resize((image.shape[1],image.shape[0]),resample=Image.BILINEAR) pb_np = np.array(imo) aaa = img_name.split(".") bbb = aaa[0:-1] imidx = bbb[0] for i in range(1,len(bbb)): imidx = imidx + "." + bbb[i] imo.save(d_dir+imidx+'.png') def main(): # --------- 1. get image path and name --------- model_name='u2net' #u2netp image_dir = os.path.join(os.getcwd(), 'test_data', 'test_images') prediction_dir = os.path.join(os.getcwd(), 'test_data', model_name + '_results' + os.sep) model_dir = os.path.join(os.getcwd(), 'saved_models', model_name, model_name + '.pth') img_name_list = glob.glob(image_dir + os.sep + '*') print(img_name_list) # --------- 2. dataloader --------- #1. dataloader test_salobj_dataset = SalObjDataset(img_name_list = img_name_list, lbl_name_list = [], transform=transforms.Compose([RescaleT(320), ToTensorLab(flag=0)]) ) test_salobj_dataloader = DataLoader(test_salobj_dataset, batch_size=1, shuffle=False, num_workers=1) # --------- 3. model define --------- if(model_name=='u2net'): print("...load U2NET---173.6 MB") net = U2NET(3,1) elif(model_name=='u2netp'): print("...load U2NEP---4.7 MB") net = U2NETP(3,1) net.load_state_dict(torch.load(model_dir)) if torch.cuda.is_available(): net.cuda() net.eval() # --------- 4. inference for each image --------- for i_test, data_test in enumerate(test_salobj_dataloader): print("inferencing:",img_name_list[i_test].split(os.sep)[-1]) inputs_test = data_test['image'] inputs_test = inputs_test.type(torch.FloatTensor) if torch.cuda.is_available(): inputs_test = Variable(inputs_test.cuda()) else: inputs_test = Variable(inputs_test) d1,d2,d3,d4,d5,d6,d7= net(inputs_test) # normalization pred = d1[:,0,:,:] pred = normPRED(pred) # save results to test_results folder if not os.path.exists(prediction_dir): os.makedirs(prediction_dir, exist_ok=True) save_output(img_name_list[i_test],pred,prediction_dir) del d1,d2,d3,d4,d5,d6,d7 if __name__ == "__main__": main()4. 相关[1] - 使用u^2net打造属于自己的remove-the-background - 2020.05.15 - 知乎[2] - 发现一个好用的显著性检测算法u^2 Net 2020 - 2020.06.23 - 知乎[3] - DeepLearing—CV系列（十九）——图像分割之U^2-Net（效果极好）的Pytorch实现 - 2020.06.29 - CSDN