Spaces:
Runtime error
Runtime error
| #modified from Github repo: https://github.com/JizhiziLi/P3M | |
| #added inference code for other networks | |
| import torch | |
| import cv2 | |
| import argparse | |
| import numpy as np | |
| from tqdm import tqdm | |
| from PIL import Image | |
| from skimage.transform import resize | |
| from torchvision import transforms,models | |
| import os | |
| from models import * | |
| import torch.nn.functional as F | |
| import torch | |
| import torch.nn as nn | |
| import math | |
| from torch.autograd import Variable | |
| import torch.nn.functional as fnn | |
| import glob | |
| import tqdm | |
| from torch.autograd import Variable | |
| from typing import Type, Any, Callable, Union, List, Optional | |
| import logging | |
| import time | |
| from omegaconf import OmegaConf | |
| config = OmegaConf.load(os.path.join(os.path.dirname( | |
| os.path.abspath(__file__)), "config/base.yaml")) | |
| device = "cuda" | |
| def conv3x3(in_planes, out_planes, stride=1): | |
| "3x3 convolution with padding" | |
| return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, | |
| padding=1, bias=False) | |
| class TFI(nn.Module): | |
| expansion = 1 | |
| def __init__(self, planes,stride=1): | |
| super(TFI, self).__init__() | |
| middle_planes = int(planes/2) | |
| self.transform = conv1x1(planes, middle_planes) | |
| self.conv1 = conv3x3(middle_planes*3, planes, stride) | |
| self.bn1 = nn.BatchNorm2d(planes) | |
| self.relu = nn.ReLU(inplace=True) | |
| self.stride = stride | |
| def forward(self, input_s_guidance, input_m_decoder, input_m_encoder): | |
| input_s_guidance_transform = self.transform(input_s_guidance) | |
| input_m_decoder_transform = self.transform(input_m_decoder) | |
| input_m_encoder_transform = self.transform(input_m_encoder) | |
| x = torch.cat((input_s_guidance_transform,input_m_decoder_transform,input_m_encoder_transform),1) | |
| out = self.conv1(x) | |
| out = self.bn1(out) | |
| out = self.relu(out) | |
| return out | |
| class SBFI(nn.Module): | |
| def __init__(self, planes,stride=1): | |
| super(SBFI, self).__init__() | |
| self.stride = stride | |
| self.transform1 = conv1x1(planes, int(planes/2)) | |
| self.transform2 = conv1x1(64, int(planes/2)) | |
| self.maxpool = nn.MaxPool2d(2, stride=stride) | |
| self.conv1 = conv3x3(planes, planes, 1) | |
| self.bn1 = nn.BatchNorm2d(planes) | |
| self.relu = nn.ReLU(inplace=True) | |
| def forward(self, input_m_decoder,e0): | |
| input_m_decoder_transform = self.transform1(input_m_decoder) | |
| e0_maxpool = self.maxpool(e0) | |
| e0_transform = self.transform2(e0_maxpool) | |
| x = torch.cat((input_m_decoder_transform,e0_transform),1) | |
| out = self.conv1(x) | |
| out = self.bn1(out) | |
| out = self.relu(out) | |
| out = out+input_m_decoder | |
| return out | |
| class DBFI(nn.Module): | |
| def __init__(self, planes,stride=1): | |
| super(DBFI, self).__init__() | |
| self.stride = stride | |
| self.transform1 = conv1x1(planes, int(planes/2)) | |
| self.transform2 = conv1x1(512, int(planes/2)) | |
| self.upsample = nn.Upsample(scale_factor=stride, mode='bilinear') | |
| self.conv1 = conv3x3(planes, planes, 1) | |
| self.bn1 = nn.BatchNorm2d(planes) | |
| self.relu = nn.ReLU(inplace=True) | |
| self.conv2 = conv3x3(planes, 3, 1) | |
| self.upsample2 = nn.Upsample(scale_factor=int(32/stride), mode='bilinear') | |
| def forward(self, input_s_decoder,e4): | |
| input_s_decoder_transform = self.transform1(input_s_decoder) | |
| e4_transform = self.transform2(e4) | |
| e4_upsample = self.upsample(e4_transform) | |
| x = torch.cat((input_s_decoder_transform,e4_upsample),1) | |
| out = self.conv1(x) | |
| out = self.bn1(out) | |
| out = self.relu(out) | |
| out = out+input_s_decoder | |
| out_side = self.conv2(out) | |
| out_side = self.upsample2(out_side) | |
| return out, out_side | |
| class P3mNet(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| self.resnet = resnet34_mp() | |
| ############################ | |
| ### Encoder part - RESNETMP | |
| ############################ | |
| self.encoder0 = nn.Sequential( | |
| self.resnet.conv1, | |
| self.resnet.bn1, | |
| self.resnet.relu, | |
| ) | |
| self.mp0 = self.resnet.maxpool1 | |
| self.encoder1 = nn.Sequential( | |
| self.resnet.layer1) | |
| self.mp1 = self.resnet.maxpool2 | |
| self.encoder2 = self.resnet.layer2 | |
| self.mp2 = self.resnet.maxpool3 | |
| self.encoder3 = self.resnet.layer3 | |
| self.mp3 = self.resnet.maxpool4 | |
| self.encoder4 = self.resnet.layer4 | |
| self.mp4 = self.resnet.maxpool5 | |
| self.tfi_3 = TFI(256) | |
| self.tfi_2 = TFI(128) | |
| self.tfi_1 = TFI(64) | |
| self.tfi_0 = TFI(64) | |
| self.sbfi_2 = SBFI(128, 8) | |
| self.sbfi_1 = SBFI(64, 4) | |
| self.sbfi_0 = SBFI(64, 2) | |
| self.dbfi_2 = DBFI(128, 4) | |
| self.dbfi_1 = DBFI(64, 8) | |
| self.dbfi_0 = DBFI(64, 16) | |
| ########################## | |
| ### Decoder part - GLOBAL | |
| ########################## | |
| self.decoder4_g = nn.Sequential( | |
| nn.Conv2d(512,512,3,padding=1), | |
| nn.BatchNorm2d(512), | |
| nn.ReLU(inplace=True), | |
| nn.Conv2d(512,512,3,padding=1), | |
| nn.BatchNorm2d(512), | |
| nn.ReLU(inplace=True), | |
| nn.Conv2d(512,256,3,padding=1), | |
| nn.BatchNorm2d(256), | |
| nn.ReLU(inplace=True), | |
| nn.Upsample(scale_factor=2, mode='bilinear') ) | |
| self.decoder3_g = nn.Sequential( | |
| nn.Conv2d(256,256,3,padding=1), | |
| nn.BatchNorm2d(256), | |
| nn.ReLU(inplace=True), | |
| nn.Conv2d(256,256,3,padding=1), | |
| nn.BatchNorm2d(256), | |
| nn.ReLU(inplace=True), | |
| nn.Conv2d(256,128,3,padding=1), | |
| nn.BatchNorm2d(128), | |
| nn.ReLU(inplace=True), | |
| nn.Upsample(scale_factor=2, mode='bilinear') ) | |
| self.decoder2_g = nn.Sequential( | |
| nn.Conv2d(128,128,3,padding=1), | |
| nn.BatchNorm2d(128), | |
| nn.ReLU(inplace=True), | |
| nn.Conv2d(128,128,3,padding=1), | |
| nn.BatchNorm2d(128), | |
| nn.ReLU(inplace=True), | |
| nn.Conv2d(128,64,3,padding=1), | |
| nn.BatchNorm2d(64), | |
| nn.ReLU(inplace=True), | |
| nn.Upsample(scale_factor=2, mode='bilinear')) | |
| self.decoder1_g = nn.Sequential( | |
| nn.Conv2d(64,64,3,padding=1), | |
| nn.BatchNorm2d(64), | |
| nn.ReLU(inplace=True), | |
| nn.Conv2d(64,64,3,padding=1), | |
| nn.BatchNorm2d(64), | |
| nn.ReLU(inplace=True), | |
| nn.Conv2d(64,64,3,padding=1), | |
| nn.BatchNorm2d(64), | |
| nn.ReLU(inplace=True), | |
| nn.Upsample(scale_factor=2, mode='bilinear')) | |
| self.decoder0_g = nn.Sequential( | |
| nn.Conv2d(64,64,3,padding=1), | |
| nn.BatchNorm2d(64), | |
| nn.ReLU(inplace=True), | |
| nn.Conv2d(64,64,3,padding=1), | |
| nn.BatchNorm2d(64), | |
| nn.ReLU(inplace=True), | |
| nn.Conv2d(64,3,3,padding=1), | |
| nn.Upsample(scale_factor=2, mode='bilinear')) | |
| ########################## | |
| ### Decoder part - LOCAL | |
| ########################## | |
| self.decoder4_l = nn.Sequential( | |
| nn.Conv2d(512,512,3,padding=1), | |
| nn.BatchNorm2d(512), | |
| nn.ReLU(inplace=True), | |
| nn.Conv2d(512,512,3,padding=1), | |
| nn.BatchNorm2d(512), | |
| nn.ReLU(inplace=True), | |
| nn.Conv2d(512,256,3,padding=1), | |
| nn.BatchNorm2d(256), | |
| nn.ReLU(inplace=True)) | |
| self.decoder3_l = nn.Sequential( | |
| nn.Conv2d(256,256,3,padding=1), | |
| nn.BatchNorm2d(256), | |
| nn.ReLU(inplace=True), | |
| nn.Conv2d(256,256,3,padding=1), | |
| nn.BatchNorm2d(256), | |
| nn.ReLU(inplace=True), | |
| nn.Conv2d(256,128,3,padding=1), | |
| nn.BatchNorm2d(128), | |
| nn.ReLU(inplace=True)) | |
| self.decoder2_l = nn.Sequential( | |
| nn.Conv2d(128,128,3,padding=1), | |
| nn.BatchNorm2d(128), | |
| nn.ReLU(inplace=True), | |
| nn.Conv2d(128,128,3,padding=1), | |
| nn.BatchNorm2d(128), | |
| nn.ReLU(inplace=True), | |
| nn.Conv2d(128,64,3,padding=1), | |
| nn.BatchNorm2d(64), | |
| nn.ReLU(inplace=True)) | |
| self.decoder1_l = nn.Sequential( | |
| nn.Conv2d(64,64,3,padding=1), | |
| nn.BatchNorm2d(64), | |
| nn.ReLU(inplace=True), | |
| nn.Conv2d(64,64,3,padding=1), | |
| nn.BatchNorm2d(64), | |
| nn.ReLU(inplace=True), | |
| nn.Conv2d(64,64,3,padding=1), | |
| nn.BatchNorm2d(64), | |
| nn.ReLU(inplace=True)) | |
| self.decoder0_l = nn.Sequential( | |
| nn.Conv2d(64,64,3,padding=1), | |
| nn.BatchNorm2d(64), | |
| nn.ReLU(inplace=True), | |
| nn.Conv2d(64,64,3,padding=1), | |
| nn.BatchNorm2d(64), | |
| nn.ReLU(inplace=True)) | |
| self.decoder_final_l = nn.Conv2d(64,1,3,padding=1) | |
| def forward(self, input): | |
| ########################## | |
| ### Encoder part - RESNET | |
| ########################## | |
| e0 = self.encoder0(input) | |
| e0p, id0 = self.mp0(e0) | |
| e1p, id1 = self.mp1(e0p) | |
| e1 = self.encoder1(e1p) | |
| e2p, id2 = self.mp2(e1) | |
| e2 = self.encoder2(e2p) | |
| e3p, id3 = self.mp3(e2) | |
| e3 = self.encoder3(e3p) | |
| e4p, id4 = self.mp4(e3) | |
| e4 = self.encoder4(e4p) | |
| ########################### | |
| ### Decoder part - Global | |
| ########################### | |
| d4_g = self.decoder4_g(e4) | |
| d3_g = self.decoder3_g(d4_g) | |
| d2_g, global_sigmoid_side2 = self.dbfi_2(d3_g, e4) | |
| d2_g = self.decoder2_g(d2_g) | |
| d1_g, global_sigmoid_side1 = self.dbfi_1(d2_g, e4) | |
| d1_g = self.decoder1_g(d1_g) | |
| d0_g, global_sigmoid_side0 = self.dbfi_0(d1_g, e4) | |
| d0_g = self.decoder0_g(d0_g) | |
| global_sigmoid = d0_g | |
| ########################### | |
| ### Decoder part - Local | |
| ########################### | |
| d4_l = self.decoder4_l(e4) | |
| d4_l = F.max_unpool2d(d4_l, id4, kernel_size=2, stride=2) | |
| d3_l = self.tfi_3(d4_g, d4_l, e3) | |
| d3_l = self.decoder3_l(d3_l) | |
| d3_l = F.max_unpool2d(d3_l, id3, kernel_size=2, stride=2) | |
| d2_l = self.tfi_2(d3_g, d3_l, e2) | |
| d2_l = self.sbfi_2(d2_l, e0) | |
| d2_l = self.decoder2_l(d2_l) | |
| d2_l = F.max_unpool2d(d2_l, id2, kernel_size=2, stride=2) | |
| d1_l = self.tfi_1(d2_g, d2_l, e1) | |
| d1_l = self.sbfi_1(d1_l, e0) | |
| d1_l = self.decoder1_l(d1_l) | |
| d1_l = F.max_unpool2d(d1_l, id1, kernel_size=2, stride=2) | |
| d0_l = self.tfi_0(d1_g, d1_l, e0p) | |
| d0_l = self.sbfi_0(d0_l, e0) | |
| d0_l = self.decoder0_l(d0_l) | |
| d0_l = F.max_unpool2d(d0_l, id0, kernel_size=2, stride=2) | |
| d0_l = self.decoder_final_l(d0_l) | |
| local_sigmoid = F.sigmoid(d0_l) | |
| ########################## | |
| ### Fusion net - G/L | |
| ########################## | |
| fusion_sigmoid = get_masked_local_from_global(global_sigmoid, local_sigmoid) | |
| return global_sigmoid, local_sigmoid, fusion_sigmoid, global_sigmoid_side2, global_sigmoid_side1, global_sigmoid_side0 | |
| def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1): | |
| """3x3 convolution with padding""" | |
| return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, | |
| padding=dilation, groups=groups, bias=False, dilation=dilation) | |
| def conv1x1(in_planes, out_planes, stride=1): | |
| """1x1 convolution""" | |
| return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False) | |
| class BasicBlock(nn.Module): | |
| expansion: int = 1 | |
| def __init__( | |
| self, | |
| inplanes: int, | |
| planes: int, | |
| stride: int = 1, | |
| downsample: Optional[nn.Module] = None, | |
| groups: int = 1, | |
| base_width: int = 64, | |
| dilation: int = 1, | |
| norm_layer: Optional[Callable[..., nn.Module]] = None | |
| ) -> None: | |
| super(BasicBlock, self).__init__() | |
| if norm_layer is None: | |
| norm_layer = nn.BatchNorm2d | |
| if groups != 1 or base_width != 64: | |
| raise ValueError('BasicBlock only supports groups=1 and base_width=64') | |
| if dilation > 1: | |
| raise NotImplementedError("Dilation > 1 not supported in BasicBlock") | |
| self.conv1 = conv3x3(inplanes, planes, stride) | |
| self.bn1 = norm_layer(planes) | |
| self.relu = nn.ReLU(inplace=True) | |
| self.conv2 = conv3x3(planes, planes) | |
| self.bn2 = norm_layer(planes) | |
| self.downsample = downsample | |
| self.stride = stride | |
| def forward(self, x): | |
| identity = x | |
| out = self.conv1(x) | |
| out = self.bn1(out) | |
| out = self.relu(out) | |
| out = self.conv2(out) | |
| out = self.bn2(out) | |
| if self.downsample is not None: | |
| identity = self.downsample(x) | |
| out += identity | |
| out = self.relu(out) | |
| return out | |
| class Bottleneck(nn.Module): | |
| expansion = 4 | |
| __constants__ = ['downsample'] | |
| def __init__(self, inplanes, planes,stride=1, downsample=None, groups=1, | |
| base_width=64, dilation=1, norm_layer=None): | |
| super(Bottleneck, self).__init__() | |
| if norm_layer is None: | |
| norm_layer = nn.BatchNorm2d | |
| width = int(planes * (base_width / 64.)) * groups | |
| self.conv1 = conv1x1(inplanes, width) | |
| self.bn1 = norm_layer(width) | |
| self.conv2 = conv3x3(width, width, stride, groups, dilation) | |
| self.bn2 = norm_layer(width) | |
| self.conv3 = conv1x1(width, planes * self.expansion) | |
| self.bn3 = norm_layer(planes * self.expansion) | |
| self.relu = nn.ReLU(inplace=True) | |
| self.downsample = downsample | |
| self.stride = stride | |
| def forward(self, x): | |
| identity = x | |
| out = self.conv1(x) | |
| out = self.bn1(out) | |
| out = self.relu(out) | |
| out = self.conv2(out) | |
| out = self.bn2(out) | |
| out = self.relu(out) | |
| out = self.attention(out) | |
| out = self.conv3(out) | |
| out = self.bn3(out) | |
| if self.downsample is not None: | |
| identity = self.downsample(x) | |
| out += identity | |
| out = self.relu(out) | |
| return out | |
| class ResNet(nn.Module): | |
| def __init__(self, block, layers, zero_init_residual=False, | |
| groups=1, width_per_group=64, replace_stride_with_dilation=None, | |
| norm_layer=None): | |
| super(ResNet, self).__init__() | |
| if norm_layer is None: | |
| norm_layer = nn.BatchNorm2d | |
| self._norm_layer = norm_layer | |
| self.inplanes = 64 | |
| self.dilation = 1 | |
| if replace_stride_with_dilation is None: | |
| replace_stride_with_dilation = [False, False, False] | |
| if len(replace_stride_with_dilation) != 3: | |
| raise ValueError("replace_stride_with_dilation should be None " | |
| "or a 3-element tuple, got {}".format(replace_stride_with_dilation)) | |
| self.groups = groups | |
| self.base_width = width_per_group | |
| self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=1, padding=3, | |
| bias=False) | |
| self.bn1 = norm_layer(self.inplanes) | |
| self.relu = nn.ReLU(inplace=True) | |
| self.maxpool1 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1, return_indices=True) | |
| self.maxpool2 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1, return_indices=True) | |
| self.maxpool3 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1, return_indices=True) | |
| self.maxpool4 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1, return_indices=True) | |
| self.maxpool5 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1, return_indices=True) | |
| #pdb.set_trace() | |
| self.layer1 = self._make_layer(block, 64, layers[0]) | |
| self.layer2 = self._make_layer(block, 128, layers[1], stride=1, | |
| dilate=replace_stride_with_dilation[0]) | |
| self.layer3 = self._make_layer(block, 256, layers[2], stride=1, | |
| dilate=replace_stride_with_dilation[1]) | |
| self.layer4 = self._make_layer(block, 512, layers[3], stride=1, | |
| dilate=replace_stride_with_dilation[2]) | |
| self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) | |
| self.fc = nn.Linear(512 * block.expansion, 1000) | |
| for m in self.modules(): | |
| if isinstance(m, nn.Conv2d): | |
| nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') | |
| elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): | |
| nn.init.constant_(m.weight, 1) | |
| nn.init.constant_(m.bias, 0) | |
| if zero_init_residual: | |
| for m in self.modules(): | |
| if isinstance(m, Bottleneck): | |
| nn.init.constant_(m.bn3.weight, 0) | |
| elif isinstance(m, BasicBlock): | |
| nn.init.constant_(m.bn2.weight, 0) | |
| def _make_layer(self, block, planes, blocks, stride=1, dilate=False): | |
| norm_layer = self._norm_layer | |
| downsample = None | |
| previous_dilation = self.dilation | |
| if dilate: | |
| self.dilation *= stride | |
| stride = 1 | |
| if stride != 1 or self.inplanes != planes * block.expansion: | |
| downsample = nn.Sequential( | |
| conv1x1(self.inplanes, planes * block.expansion, stride), | |
| norm_layer(planes * block.expansion), | |
| ) | |
| layers = [] | |
| layers.append(block(self.inplanes, planes,stride, downsample, self.groups, | |
| self.base_width, previous_dilation, norm_layer)) | |
| self.inplanes = planes * block.expansion | |
| for _ in range(1, blocks): | |
| layers.append(block(self.inplanes, planes,groups=self.groups, | |
| base_width=self.base_width, dilation=self.dilation, | |
| norm_layer=norm_layer)) | |
| return nn.Sequential(*layers) | |
| def _forward_impl(self, x): | |
| x1 = self.conv1(x) | |
| x1 = self.bn1(x1) | |
| x1 = self.relu(x1) | |
| x1, idx1 = self.maxpool1(x1) | |
| x2, idx2 = self.maxpool2(x1) | |
| x2 = self.layer1(x2) | |
| x3, idx3 = self.maxpool3(x2) | |
| x3 = self.layer2(x3) | |
| x4, idx4 = self.maxpool4(x3) | |
| x4 = self.layer3(x4) | |
| x5, idx5 = self.maxpool5(x4) | |
| x5 = self.layer4(x5) | |
| x_cls = self.avgpool(x5) | |
| x_cls = torch.flatten(x_cls, 1) | |
| x_cls = self.fc(x_cls) | |
| return x_cls | |
| def forward(self, x): | |
| return self._forward_impl(x) | |
| def resnet34_mp(**kwargs): | |
| r"""ResNet-34 model from | |
| `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>` | |
| """ | |
| model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs) | |
| checkpoint = torch.load("checkpoints/r34mp_pretrained_imagenet.pth.tar") | |
| model.load_state_dict(checkpoint) | |
| return model | |
| ############################## | |
| ### Training loses for P3M-NET | |
| ############################## | |
| def get_crossentropy_loss(gt,pre): | |
| gt_copy = gt.clone() | |
| gt_copy[gt_copy==0] = 0 | |
| gt_copy[gt_copy==255] = 2 | |
| gt_copy[gt_copy>2] = 1 | |
| gt_copy = gt_copy.long() | |
| gt_copy = gt_copy[:,0,:,:] | |
| criterion = nn.CrossEntropyLoss() | |
| entropy_loss = criterion(pre, gt_copy) | |
| return entropy_loss | |
| def get_alpha_loss(predict, alpha, trimap): | |
| weighted = torch.zeros(trimap.shape).cuda() | |
| weighted[trimap == 128] = 1. | |
| alpha_f = alpha / 255. | |
| alpha_f = alpha_f.cuda() | |
| diff = predict - alpha_f | |
| diff = diff * weighted | |
| alpha_loss = torch.sqrt(diff ** 2 + 1e-12) | |
| alpha_loss_weighted = alpha_loss.sum() / (weighted.sum() + 1.) | |
| return alpha_loss_weighted | |
| def get_alpha_loss_whole_img(predict, alpha): | |
| weighted = torch.ones(alpha.shape).cuda() | |
| alpha_f = alpha / 255. | |
| alpha_f = alpha_f.cuda() | |
| diff = predict - alpha_f | |
| alpha_loss = torch.sqrt(diff ** 2 + 1e-12) | |
| alpha_loss = alpha_loss.sum()/(weighted.sum()) | |
| return alpha_loss | |
| ## Laplacian loss is refer to | |
| ## https://gist.github.com/MarcoForte/a07c40a2b721739bb5c5987671aa5270 | |
| def build_gauss_kernel(size=5, sigma=1.0, n_channels=1, cuda=False): | |
| if size % 2 != 1: | |
| raise ValueError("kernel size must be uneven") | |
| grid = np.float32(np.mgrid[0:size,0:size].T) | |
| gaussian = lambda x: np.exp((x - size//2)**2/(-2*sigma**2))**2 | |
| kernel = np.sum(gaussian(grid), axis=2) | |
| kernel /= np.sum(kernel) | |
| kernel = np.tile(kernel, (n_channels, 1, 1)) | |
| kernel = torch.FloatTensor(kernel[:, None, :, :]).cuda() | |
| return Variable(kernel, requires_grad=False) | |
| def conv_gauss(img, kernel): | |
| """ convolve img with a gaussian kernel that has been built with build_gauss_kernel """ | |
| n_channels, _, kw, kh = kernel.shape | |
| img = fnn.pad(img, (kw//2, kh//2, kw//2, kh//2), mode='replicate') | |
| return fnn.conv2d(img, kernel, groups=n_channels) | |
| def laplacian_pyramid(img, kernel, max_levels=5): | |
| current = img | |
| pyr = [] | |
| for level in range(max_levels): | |
| filtered = conv_gauss(current, kernel) | |
| diff = current - filtered | |
| pyr.append(diff) | |
| current = fnn.avg_pool2d(filtered, 2) | |
| pyr.append(current) | |
| return pyr | |
| def get_laplacian_loss(predict, alpha, trimap): | |
| weighted = torch.zeros(trimap.shape).cuda() | |
| weighted[trimap == 128] = 1. | |
| alpha_f = alpha / 255. | |
| alpha_f = alpha_f.cuda() | |
| alpha_f = alpha_f.clone()*weighted | |
| predict = predict.clone()*weighted | |
| gauss_kernel = build_gauss_kernel(size=5, sigma=1.0, n_channels=1, cuda=True) | |
| pyr_alpha = laplacian_pyramid(alpha_f, gauss_kernel, 5) | |
| pyr_predict = laplacian_pyramid(predict, gauss_kernel, 5) | |
| laplacian_loss_weighted = sum(fnn.l1_loss(a, b) for a, b in zip(pyr_alpha, pyr_predict)) | |
| return laplacian_loss_weighted | |
| def get_laplacian_loss_whole_img(predict, alpha): | |
| alpha_f = alpha / 255. | |
| alpha_f = alpha_f.cuda() | |
| gauss_kernel = build_gauss_kernel(size=5, sigma=1.0, n_channels=1, cuda=True) | |
| pyr_alpha = laplacian_pyramid(alpha_f, gauss_kernel, 5) | |
| pyr_predict = laplacian_pyramid(predict, gauss_kernel, 5) | |
| laplacian_loss = sum(fnn.l1_loss(a, b) for a, b in zip(pyr_alpha, pyr_predict)) | |
| return laplacian_loss | |
| def get_composition_loss_whole_img(img, alpha, fg, bg, predict): | |
| weighted = torch.ones(alpha.shape).cuda() | |
| predict_3 = torch.cat((predict, predict, predict), 1) | |
| comp = predict_3 * fg + (1. - predict_3) * bg | |
| comp_loss = torch.sqrt((comp - img) ** 2 + 1e-12) | |
| comp_loss = comp_loss.sum()/(weighted.sum()) | |
| return comp_loss | |
| ############################## | |
| ### Test loss for matting | |
| ############################## | |
| def calculate_sad_mse_mad(predict_old,alpha,trimap): | |
| predict = np.copy(predict_old) | |
| pixel = float((trimap == 128).sum()) | |
| predict[trimap == 255] = 1. | |
| predict[trimap == 0 ] = 0. | |
| sad_diff = np.sum(np.abs(predict - alpha))/1000 | |
| if pixel==0: | |
| pixel = trimap.shape[0]*trimap.shape[1]-float((trimap==255).sum())-float((trimap==0).sum()) | |
| mse_diff = np.sum((predict - alpha) ** 2)/pixel | |
| mad_diff = np.sum(np.abs(predict - alpha))/pixel | |
| return sad_diff, mse_diff, mad_diff | |
| def calculate_sad_mse_mad_whole_img(predict, alpha): | |
| pixel = predict.shape[0]*predict.shape[1] | |
| sad_diff = np.sum(np.abs(predict - alpha))/1000 | |
| mse_diff = np.sum((predict - alpha) ** 2)/pixel | |
| mad_diff = np.sum(np.abs(predict - alpha))/pixel | |
| return sad_diff, mse_diff, mad_diff | |
| def calculate_sad_fgbg(predict, alpha, trimap): | |
| sad_diff = np.abs(predict-alpha) | |
| weight_fg = np.zeros(predict.shape) | |
| weight_bg = np.zeros(predict.shape) | |
| weight_trimap = np.zeros(predict.shape) | |
| weight_fg[trimap==255] = 1. | |
| weight_bg[trimap==0 ] = 1. | |
| weight_trimap[trimap==128 ] = 1. | |
| sad_fg = np.sum(sad_diff*weight_fg)/1000 | |
| sad_bg = np.sum(sad_diff*weight_bg)/1000 | |
| sad_trimap = np.sum(sad_diff*weight_trimap)/1000 | |
| return sad_fg, sad_bg | |
| def compute_gradient_whole_image(pd, gt): | |
| from scipy.ndimage import gaussian_filter | |
| pd_x = gaussian_filter(pd, sigma=1.4, order=[1, 0], output=np.float32) | |
| pd_y = gaussian_filter(pd, sigma=1.4, order=[0, 1], output=np.float32) | |
| gt_x = gaussian_filter(gt, sigma=1.4, order=[1, 0], output=np.float32) | |
| gt_y = gaussian_filter(gt, sigma=1.4, order=[0, 1], output=np.float32) | |
| pd_mag = np.sqrt(pd_x**2 + pd_y**2) | |
| gt_mag = np.sqrt(gt_x**2 + gt_y**2) | |
| error_map = np.square(pd_mag - gt_mag) | |
| loss = np.sum(error_map) / 10 | |
| return loss | |
| def compute_connectivity_loss_whole_image(pd, gt, step=0.1): | |
| from scipy.ndimage import morphology | |
| from skimage.measure import label, regionprops | |
| h, w = pd.shape | |
| thresh_steps = np.arange(0, 1.1, step) | |
| l_map = -1 * np.ones((h, w), dtype=np.float32) | |
| lambda_map = np.ones((h, w), dtype=np.float32) | |
| for i in range(1, thresh_steps.size): | |
| pd_th = pd >= thresh_steps[i] | |
| gt_th = gt >= thresh_steps[i] | |
| label_image = label(pd_th & gt_th, connectivity=1) | |
| cc = regionprops(label_image) | |
| size_vec = np.array([c.area for c in cc]) | |
| if len(size_vec) == 0: | |
| continue | |
| max_id = np.argmax(size_vec) | |
| coords = cc[max_id].coords | |
| omega = np.zeros((h, w), dtype=np.float32) | |
| omega[coords[:, 0], coords[:, 1]] = 1 | |
| flag = (l_map == -1) & (omega == 0) | |
| l_map[flag == 1] = thresh_steps[i-1] | |
| dist_maps = morphology.distance_transform_edt(omega==0) | |
| dist_maps = dist_maps / dist_maps.max() | |
| l_map[l_map == -1] = 1 | |
| d_pd = pd - l_map | |
| d_gt = gt - l_map | |
| phi_pd = 1 - d_pd * (d_pd >= 0.15).astype(np.float32) | |
| phi_gt = 1 - d_gt * (d_gt >= 0.15).astype(np.float32) | |
| loss = np.sum(np.abs(phi_pd - phi_gt)) / 1000 | |
| return loss | |
| def gen_trimap_from_segmap_e2e(segmap): | |
| trimap = np.argmax(segmap, axis=1)[0] | |
| trimap = trimap.astype(np.int64) | |
| trimap[trimap==1]=128 | |
| trimap[trimap==2]=255 | |
| return trimap.astype(np.uint8) | |
| def get_masked_local_from_global(global_sigmoid, local_sigmoid): | |
| values, index = torch.max(global_sigmoid,1) | |
| index = index[:,None,:,:].float() | |
| ### index <===> [0, 1, 2] | |
| ### bg_mask <===> [1, 0, 0] | |
| bg_mask = index.clone() | |
| bg_mask[bg_mask==2]=1 | |
| bg_mask = 1- bg_mask | |
| ### trimap_mask <===> [0, 1, 0] | |
| trimap_mask = index.clone() | |
| trimap_mask[trimap_mask==2]=0 | |
| ### fg_mask <===> [0, 0, 1] | |
| fg_mask = index.clone() | |
| fg_mask[fg_mask==1]=0 | |
| fg_mask[fg_mask==2]=1 | |
| fusion_sigmoid = local_sigmoid*trimap_mask+fg_mask | |
| return fusion_sigmoid | |
| def get_masked_local_from_global_test(global_result, local_result): | |
| weighted_global = np.ones(global_result.shape) | |
| weighted_global[global_result==255] = 0 | |
| weighted_global[global_result==0] = 0 | |
| fusion_result = global_result*(1.-weighted_global)/255+local_result*weighted_global | |
| return fusion_result | |
| def inference_once( model, scale_img, scale_trimap=None): | |
| pred_list = [] | |
| tensor_img = torch.from_numpy(scale_img[:, :, :]).permute(2, 0, 1).cuda() | |
| input_t = tensor_img | |
| input_t = input_t/255.0 | |
| normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], | |
| std=[0.229, 0.224, 0.225]) | |
| input_t = normalize(input_t) | |
| input_t = input_t.unsqueeze(0).float() | |
| # pred_global, pred_local, pred_fusion = model(input_t)[:3] | |
| pred_fusion = model(input_t)[:3] | |
| pred_global = pred_fusion | |
| pred_local = pred_fusion | |
| pred_global = pred_global.data.cpu().numpy() | |
| pred_global = gen_trimap_from_segmap_e2e(pred_global) | |
| pred_local = pred_local.data.cpu().numpy()[0,0,:,:] | |
| pred_fusion = pred_fusion.data.cpu().numpy()[0,0,:,:] | |
| return pred_global, pred_local, pred_fusion | |
| # def inference_img( test_choice,model, img): | |
| # h, w, c = img.shape | |
| # new_h = min(config['datasets'].MAX_SIZE_H, h - (h % 32)) | |
| # new_w = min(config['datasets'].MAX_SIZE_W, w - (w % 32)) | |
| # if test_choice=='HYBRID': | |
| # global_ratio = 1/2 | |
| # local_ratio = 1 | |
| # resize_h = int(h*global_ratio) | |
| # resize_w = int(w*global_ratio) | |
| # new_h = min(config['datasets'].MAX_SIZE_H, resize_h - (resize_h % 32)) | |
| # new_w = min(config['datasets'].MAX_SIZE_W, resize_w - (resize_w % 32)) | |
| # scale_img = resize(img,(new_h,new_w))*255.0 | |
| # pred_coutour_1, pred_retouching_1, pred_fusion_1 = inference_once( model, scale_img) | |
| # pred_coutour_1 = resize(pred_coutour_1,(h,w))*255.0 | |
| # resize_h = int(h*local_ratio) | |
| # resize_w = int(w*local_ratio) | |
| # new_h = min(config['datasets'].MAX_SIZE_H, resize_h - (resize_h % 32)) | |
| # new_w = min(config['datasets'].MAX_SIZE_W, resize_w - (resize_w % 32)) | |
| # scale_img = resize(img,(new_h,new_w))*255.0 | |
| # pred_coutour_2, pred_retouching_2, pred_fusion_2 = inference_once( model, scale_img) | |
| # pred_retouching_2 = resize(pred_retouching_2,(h,w)) | |
| # pred_fusion = get_masked_local_from_global_test(pred_coutour_1, pred_retouching_2) | |
| # return pred_fusion | |
| # else: | |
| # resize_h = int(h/2) | |
| # resize_w = int(w/2) | |
| # new_h = min(config['datasets'].MAX_SIZE_H, resize_h - (resize_h % 32)) | |
| # new_w = min(config['datasets'].MAX_SIZE_W, resize_w - (resize_w % 32)) | |
| # scale_img = resize(img,(new_h,new_w))*255.0 | |
| # pred_global, pred_local, pred_fusion = inference_once( model, scale_img) | |
| # pred_local = resize(pred_local,(h,w)) | |
| # pred_global = resize(pred_global,(h,w))*255.0 | |
| # pred_fusion = resize(pred_fusion,(h,w)) | |
| # return pred_fusion | |
| def inference_img(model, img): | |
| h,w,_ = img.shape | |
| # print(img.shape) | |
| if h%8!=0 or w%8!=0: | |
| img=cv2.copyMakeBorder(img, 8-h%8, 0, 8-w%8, 0, cv2.BORDER_REFLECT) | |
| # print(img.shape) | |
| tensor_img = torch.from_numpy(img).permute(2, 0, 1).cuda() | |
| input_t = tensor_img | |
| input_t = input_t/255.0 | |
| normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], | |
| std=[0.229, 0.224, 0.225]) | |
| input_t = normalize(input_t) | |
| input_t = input_t.unsqueeze(0).float() | |
| with torch.no_grad(): | |
| out=model(input_t) | |
| # print("out",out.shape) | |
| result = out[0][:,-h:,-w:].cpu().numpy() | |
| # print(result.shape) | |
| return result[0] | |
| def test_am2k(model): | |
| ############################ | |
| # Some initial setting for paths | |
| ############################ | |
| ORIGINAL_PATH = config['datasets']['am2k']['validation_original'] | |
| MASK_PATH = config['datasets']['am2k']['validation_mask'] | |
| TRIMAP_PATH = config['datasets']['am2k']['validation_trimap'] | |
| img_paths = glob.glob(ORIGINAL_PATH+"/*.jpg") | |
| ############################ | |
| # Start testing | |
| ############################ | |
| sad_diffs = 0. | |
| mse_diffs = 0. | |
| mad_diffs = 0. | |
| grad_diffs = 0. | |
| conn_diffs = 0. | |
| sad_trimap_diffs = 0. | |
| mse_trimap_diffs = 0. | |
| mad_trimap_diffs = 0. | |
| sad_fg_diffs = 0. | |
| sad_bg_diffs = 0. | |
| total_number = len(img_paths) | |
| log("===============================") | |
| log(f'====> Start Testing\n\t--Dataset: AM2k\n\t-\n\t--Number: {total_number}') | |
| for img_path in tqdm.tqdm(img_paths): | |
| img_name=(img_path.split("/")[-1])[:-4] | |
| alpha_path = MASK_PATH+img_name+'.png' | |
| trimap_path = TRIMAP_PATH+img_name+'.png' | |
| pil_img = Image.open(img_path) | |
| img = np.array(pil_img) | |
| trimap = np.array(Image.open(trimap_path)) | |
| alpha = np.array(Image.open(alpha_path))/255. | |
| img = img[:,:,:3] if img.ndim>2 else img | |
| trimap = trimap[:,:,0] if trimap.ndim>2 else trimap | |
| alpha = alpha[:,:,0] if alpha.ndim>2 else alpha | |
| with torch.no_grad(): | |
| torch.cuda.empty_cache() | |
| predict = inference_img( model, img) | |
| sad_trimap_diff, mse_trimap_diff, mad_trimap_diff = calculate_sad_mse_mad(predict, alpha, trimap) | |
| sad_diff, mse_diff, mad_diff = calculate_sad_mse_mad_whole_img(predict, alpha) | |
| sad_fg_diff, sad_bg_diff = calculate_sad_fgbg(predict, alpha, trimap) | |
| conn_diff = compute_connectivity_loss_whole_image(predict, alpha) | |
| grad_diff = compute_gradient_whole_image(predict, alpha) | |
| log(f"[{img_paths.index(img_path)}/{total_number}]\nImage:{img_name}\nsad:{sad_diff}\nmse:{mse_diff}\nmad:{mad_diff}\nsad_trimap:{sad_trimap_diff}\nmse_trimap:{mse_trimap_diff}\nmad_trimap:{mad_trimap_diff}\nsad_fg:{sad_fg_diff}\nsad_bg:{sad_bg_diff}\nconn:{conn_diff}\ngrad:{grad_diff}\n-----------") | |
| sad_diffs += sad_diff | |
| mse_diffs += mse_diff | |
| mad_diffs += mad_diff | |
| mse_trimap_diffs += mse_trimap_diff | |
| sad_trimap_diffs += sad_trimap_diff | |
| mad_trimap_diffs += mad_trimap_diff | |
| sad_fg_diffs += sad_fg_diff | |
| sad_bg_diffs += sad_bg_diff | |
| conn_diffs += conn_diff | |
| grad_diffs += grad_diff | |
| Image.fromarray(np.uint8(predict*255)).save(f"test/{img_name}.png") | |
| log("===============================") | |
| log(f"Testing numbers: {total_number}") | |
| log("SAD: {}".format(sad_diffs / total_number)) | |
| log("MSE: {}".format(mse_diffs / total_number)) | |
| log("MAD: {}".format(mad_diffs / total_number)) | |
| log("GRAD: {}".format(grad_diffs / total_number)) | |
| log("CONN: {}".format(conn_diffs / total_number)) | |
| log("SAD TRIMAP: {}".format(sad_trimap_diffs / total_number)) | |
| log("MSE TRIMAP: {}".format(mse_trimap_diffs / total_number)) | |
| log("MAD TRIMAP: {}".format(mad_trimap_diffs / total_number)) | |
| log("SAD FG: {}".format(sad_fg_diffs / total_number)) | |
| log("SAD BG: {}".format(sad_bg_diffs / total_number)) | |
| return sad_diffs/total_number,mse_diffs/total_number,grad_diffs/total_number | |
| def test_p3m10k(model,dataset_choice, max_image=-1): | |
| ############################ | |
| # Some initial setting for paths | |
| ############################ | |
| if dataset_choice == 'P3M_500_P': | |
| val_option = 'VAL500P' | |
| else: | |
| val_option = 'VAL500NP' | |
| ORIGINAL_PATH = config['datasets']['p3m10k']+"/validation/"+config['datasets']['p3m10k_test'][val_option]['ORIGINAL_PATH'] | |
| MASK_PATH = config['datasets']['p3m10k']+"/validation/"+config['datasets']['p3m10k_test'][val_option]['MASK_PATH'] | |
| TRIMAP_PATH = config['datasets']['p3m10k']+"/validation/"+config['datasets']['p3m10k_test'][val_option]['TRIMAP_PATH'] | |
| ############################ | |
| # Start testing | |
| ############################ | |
| sad_diffs = 0. | |
| mse_diffs = 0. | |
| mad_diffs = 0. | |
| sad_trimap_diffs = 0. | |
| mse_trimap_diffs = 0. | |
| mad_trimap_diffs = 0. | |
| sad_fg_diffs = 0. | |
| sad_bg_diffs = 0. | |
| conn_diffs = 0. | |
| grad_diffs = 0. | |
| model.eval() | |
| img_paths = glob.glob(ORIGINAL_PATH+"/*.jpg") | |
| if (max_image>1): | |
| img_paths = img_paths[:max_image] | |
| total_number = len(img_paths) | |
| log("===============================") | |
| log(f'====> Start Testing\n\t----Test: {dataset_choice}\n\t--Number: {total_number}') | |
| for img_path in tqdm.tqdm(img_paths): | |
| img_name=(img_path.split("/")[-1])[:-4] | |
| alpha_path = MASK_PATH+img_name+'.png' | |
| trimap_path = TRIMAP_PATH+img_name+'.png' | |
| pil_img = Image.open(img_path) | |
| img = np.array(pil_img) | |
| trimap = np.array(Image.open(trimap_path)) | |
| alpha = np.array(Image.open(alpha_path))/255. | |
| img = img[:,:,:3] if img.ndim>2 else img | |
| trimap = trimap[:,:,0] if trimap.ndim>2 else trimap | |
| alpha = alpha[:,:,0] if alpha.ndim>2 else alpha | |
| with torch.no_grad(): | |
| torch.cuda.empty_cache() | |
| start = time.time() | |
| predict = inference_img( model, img) #HYBRID show less accuracy | |
| # tensorimg=transforms.ToTensor()(pil_img) | |
| # input_img=transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])(tensorimg) | |
| # predict = model(input_img.unsqueeze(0).to(device))[0][0].detach().cpu().numpy() | |
| # if predict.shape!=(pil_img.height,pil_img.width): | |
| # print("resize for ",img_path) | |
| # predict = resize(predict,(pil_img.height,pil_img.width)) | |
| sad_trimap_diff, mse_trimap_diff, mad_trimap_diff = calculate_sad_mse_mad(predict, alpha, trimap) | |
| sad_diff, mse_diff, mad_diff = calculate_sad_mse_mad_whole_img(predict, alpha) | |
| sad_fg_diff, sad_bg_diff = calculate_sad_fgbg(predict, alpha, trimap) | |
| conn_diff = compute_connectivity_loss_whole_image(predict, alpha) | |
| grad_diff = compute_gradient_whole_image(predict, alpha) | |
| log(f"[{img_paths.index(img_path)}/{total_number}]\nImage:{img_name}\nsad:{sad_diff}\nmse:{mse_diff}\nmad:{mad_diff}\nconn:{conn_diff}\ngrad:{grad_diff}\n-----------") | |
| sad_diffs += sad_diff | |
| mse_diffs += mse_diff | |
| mad_diffs += mad_diff | |
| mse_trimap_diffs += mse_trimap_diff | |
| sad_trimap_diffs += sad_trimap_diff | |
| mad_trimap_diffs += mad_trimap_diff | |
| sad_fg_diffs += sad_fg_diff | |
| sad_bg_diffs += sad_bg_diff | |
| conn_diffs += conn_diff | |
| grad_diffs += grad_diff | |
| Image.fromarray(np.uint8(predict*255)).save(f"test/{img_name}.png") | |
| log("===============================") | |
| log(f"Testing numbers: {total_number}") | |
| log("SAD: {}".format(sad_diffs / total_number)) | |
| log("MSE: {}".format(mse_diffs / total_number)) | |
| log("MAD: {}".format(mad_diffs / total_number)) | |
| log("SAD TRIMAP: {}".format(sad_trimap_diffs / total_number)) | |
| log("MSE TRIMAP: {}".format(mse_trimap_diffs / total_number)) | |
| log("MAD TRIMAP: {}".format(mad_trimap_diffs / total_number)) | |
| log("SAD FG: {}".format(sad_fg_diffs / total_number)) | |
| log("SAD BG: {}".format(sad_bg_diffs / total_number)) | |
| log("CONN: {}".format(conn_diffs / total_number)) | |
| log("GRAD: {}".format(grad_diffs / total_number)) | |
| return sad_diffs/total_number,mse_diffs/total_number,grad_diffs/total_number | |
| def log(str): | |
| print(str) | |
| logging.info(str) | |
| if __name__ == '__main__': | |
| print('*********************************') | |
| config = OmegaConf.load(os.path.join(os.path.dirname( | |
| os.path.abspath(__file__)), "config/base.yaml")) | |
| config=OmegaConf.merge(config,OmegaConf.from_cli()) | |
| print(config) | |
| model = MaskForm() | |
| model = model.to(device) | |
| checkpoint = f"{config.checkpoint_dir}/{config.checkpoint}" | |
| state_dict = torch.load(checkpoint, map_location=f'{device}') | |
| print("loaded",checkpoint) | |
| model.load_state_dict(state_dict) | |
| model.eval() | |
| logging.basicConfig(filename=f'report/{config.checkpoint.replace("/","--")}.report', encoding='utf-8',filemode='w', level=logging.INFO) | |
| # ckpt = torch.load("checkpoints/p3mnet_pretrained_on_p3m10k.pth") | |
| # model.load_state_dict(ckpt['state_dict'], strict=True) | |
| # model = model.cuda() | |
| if config.dataset_to_use =="AM2K": | |
| test_am2k(model) | |
| else: | |
| for dataset_choice in ['P3M_500_P','P3M_500_NP']: | |
| test_p3m10k(model,dataset_choice) | |