mirror of https://github.com/ZHKKKe/MODNet.git
300 lines
12 KiB
Python
300 lines
12 KiB
Python
import math
|
|
import scipy
|
|
import numpy as np
|
|
from scipy.ndimage import grey_dilation, grey_erosion
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
|
|
__all__ = [
|
|
'supervised_training_iter',
|
|
'soc_adaptation_iter',
|
|
]
|
|
|
|
|
|
# ----------------------------------------------------------------------------------
|
|
# Tool Classes/Functions
|
|
# ----------------------------------------------------------------------------------
|
|
|
|
class GaussianBlurLayer(nn.Module):
|
|
""" Add Gaussian Blur to a 4D tensors
|
|
This layer takes a 4D tensor of {N, C, H, W} as input.
|
|
The Gaussian blur will be performed in given channel number (C) splitly.
|
|
"""
|
|
|
|
def __init__(self, channels, kernel_size):
|
|
"""
|
|
Arguments:
|
|
channels (int): Channel for input tensor
|
|
kernel_size (int): Size of the kernel used in blurring
|
|
"""
|
|
|
|
super(GaussianBlurLayer, self).__init__()
|
|
self.channels = channels
|
|
self.kernel_size = kernel_size
|
|
assert self.kernel_size % 2 != 0
|
|
|
|
self.op = nn.Sequential(
|
|
nn.ReflectionPad2d(math.floor(self.kernel_size / 2)),
|
|
nn.Conv2d(channels, channels, self.kernel_size,
|
|
stride=1, padding=0, bias=None, groups=channels)
|
|
)
|
|
|
|
self._init_kernel()
|
|
|
|
def forward(self, x):
|
|
"""
|
|
Arguments:
|
|
x (torch.Tensor): input 4D tensor
|
|
Returns:
|
|
torch.Tensor: Blurred version of the input
|
|
"""
|
|
|
|
if not len(list(x.shape)) == 4:
|
|
print('\'GaussianBlurLayer\' requires a 4D tensor as input\n')
|
|
exit()
|
|
elif not x.shape[1] == self.channels:
|
|
print('In \'GaussianBlurLayer\', the required channel ({0}) is'
|
|
'not the same as input ({1})\n'.format(self.channels, x.shape[1]))
|
|
exit()
|
|
|
|
return self.op(x)
|
|
|
|
def _init_kernel(self):
|
|
sigma = 0.3 * ((self.kernel_size - 1) * 0.5 - 1) + 0.8
|
|
|
|
n = np.zeros((self.kernel_size, self.kernel_size))
|
|
i = math.floor(self.kernel_size / 2)
|
|
n[i, i] = 1
|
|
kernel = scipy.ndimage.gaussian_filter(n, sigma)
|
|
|
|
for name, param in self.named_parameters():
|
|
param.data.copy_(torch.from_numpy(kernel))
|
|
|
|
# ----------------------------------------------------------------------------------
|
|
|
|
|
|
# ----------------------------------------------------------------------------------
|
|
# MODNet Training Functions
|
|
# ----------------------------------------------------------------------------------
|
|
|
|
blurer = GaussianBlurLayer(1, 3).cuda()
|
|
|
|
|
|
def supervised_training_iter(
|
|
modnet, optimizer, image, trimap, gt_matte,
|
|
semantic_scale=10.0, detail_scale=10.0, matte_scale=1.0):
|
|
""" Supervised training iteration of MODNet
|
|
This function trains MODNet for one iteration in a labeled dataset.
|
|
|
|
Arguments:
|
|
modnet (torch.nn.Module): instance of MODNet
|
|
optimizer (torch.optim.Optimizer): optimizer for supervised training
|
|
image (torch.autograd.Variable): input RGB image
|
|
its pixel values should be normalized
|
|
trimap (torch.autograd.Variable): trimap used to calculate the losses
|
|
its pixel values can be 0, 0.5, or 1
|
|
(foreground=1, background=0, unknown=0.5)
|
|
gt_matte (torch.autograd.Variable): ground truth alpha matte
|
|
its pixel values are between [0, 1]
|
|
semantic_scale (float): scale of the semantic loss
|
|
NOTE: please adjust according to your dataset
|
|
detail_scale (float): scale of the detail loss
|
|
NOTE: please adjust according to your dataset
|
|
matte_scale (float): scale of the matte loss
|
|
NOTE: please adjust according to your dataset
|
|
|
|
Returns:
|
|
semantic_loss (torch.Tensor): loss of the semantic estimation [Low-Resolution (LR) Branch]
|
|
detail_loss (torch.Tensor): loss of the detail prediction [High-Resolution (HR) Branch]
|
|
matte_loss (torch.Tensor): loss of the semantic-detail fusion [Fusion Branch]
|
|
|
|
Example:
|
|
import torch
|
|
from src.models.modnet import MODNet
|
|
from src.trainer import supervised_training_iter
|
|
|
|
bs = 16 # batch size
|
|
lr = 0.01 # learn rate
|
|
epochs = 40 # total epochs
|
|
|
|
modnet = torch.nn.DataParallel(MODNet()).cuda()
|
|
optimizer = torch.optim.SGD(modnet.parameters(), lr=lr, momentum=0.9)
|
|
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=int(0.25 * epochs), gamma=0.1)
|
|
|
|
dataloader = CREATE_YOUR_DATALOADER(bs) # NOTE: please finish this function
|
|
|
|
for epoch in range(0, epochs):
|
|
for idx, (image, trimap, gt_matte) in enumerate(dataloader):
|
|
semantic_loss, detail_loss, matte_loss = \
|
|
supervised_training_iter(modnet, optimizer, image, trimap, gt_matte)
|
|
lr_scheduler.step()
|
|
"""
|
|
|
|
global blurer
|
|
|
|
# set the model to train mode and clear the optimizer
|
|
modnet.train()
|
|
optimizer.zero_grad()
|
|
|
|
# forward the model
|
|
pred_semantic, pred_detail, pred_matte = modnet(image, False)
|
|
|
|
# calculate the boundary mask from the trimap
|
|
boundaries = (trimap < 0.5) + (trimap > 0.5)
|
|
|
|
# calculate the semantic loss
|
|
gt_semantic = F.interpolate(gt_matte, scale_factor=1/16, mode='bilinear')
|
|
gt_semantic = blurer(gt_semantic)
|
|
semantic_loss = torch.mean(F.mse_loss(pred_semantic, gt_semantic))
|
|
semantic_loss = semantic_scale * semantic_loss
|
|
|
|
# calculate the detail loss
|
|
pred_boundary_detail = torch.where(boundaries, trimap, pred_detail)
|
|
gt_detail = torch.where(boundaries, trimap, gt_matte)
|
|
detail_loss = torch.mean(F.l1_loss(pred_boundary_detail, gt_detail))
|
|
detail_loss = detail_scale * detail_loss
|
|
|
|
# calculate the matte loss
|
|
pred_boundary_matte = torch.where(boundaries, trimap, pred_matte)
|
|
matte_l1_loss = F.l1_loss(pred_matte, gt_matte) + 4.0 * F.l1_loss(pred_boundary_matte, gt_matte)
|
|
matte_compositional_loss = F.l1_loss(image * pred_matte, image * gt_matte) \
|
|
+ 4.0 * F.l1_loss(image * pred_boundary_matte, image * gt_matte)
|
|
matte_loss = torch.mean(matte_l1_loss + matte_compositional_loss)
|
|
matte_loss = matte_scale * matte_loss
|
|
|
|
# calculate the final loss, backward the loss, and update the model
|
|
loss = semantic_loss + detail_loss + matte_loss
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
# for test
|
|
return semantic_loss, detail_loss, matte_loss
|
|
|
|
|
|
def soc_adaptation_iter(
|
|
modnet, backup_modnet, optimizer, image,
|
|
soc_semantic_scale=100.0, soc_detail_scale=1.0):
|
|
""" Self-Supervised sub-objective consistency (SOC) adaptation iteration of MODNet
|
|
This function fine-tunes MODNet for one iteration in an unlabeled dataset.
|
|
Note that SOC can only fine-tune a converged MODNet, i.e., MODNet that has been
|
|
trained in a labeled dataset.
|
|
|
|
Arguments:
|
|
modnet (torch.nn.Module): instance of MODNet
|
|
backup_modnet (torch.nn.Module): backup of the trained MODNet
|
|
optimizer (torch.optim.Optimizer): optimizer for self-supervised SOC
|
|
image (torch.autograd.Variable): input RGB image
|
|
its pixel values should be normalized
|
|
soc_semantic_scale (float): scale of the SOC semantic loss
|
|
NOTE: please adjust according to your dataset
|
|
soc_detail_scale (float): scale of the SOC detail loss
|
|
NOTE: please adjust according to your dataset
|
|
|
|
Returns:
|
|
soc_semantic_loss (torch.Tensor): loss of the semantic SOC
|
|
soc_detail_loss (torch.Tensor): loss of the detail SOC
|
|
|
|
Example:
|
|
import copy
|
|
import torch
|
|
from src.models.modnet import MODNet
|
|
from src.trainer import soc_adaptation_iter
|
|
|
|
bs = 1 # batch size
|
|
lr = 0.00001 # learn rate
|
|
epochs = 10 # total epochs
|
|
|
|
modnet = torch.nn.DataParallel(MODNet()).cuda()
|
|
modnet = LOAD_TRAINED_CKPT() # NOTE: please finish this function
|
|
|
|
optimizer = torch.optim.Adam(modnet.parameters(), lr=lr, betas=(0.9, 0.99))
|
|
dataloader = CREATE_YOUR_DATALOADER(bs) # NOTE: please finish this function
|
|
|
|
for epoch in range(0, epochs):
|
|
backup_modnet = copy.deepcopy(modnet)
|
|
for idx, (image) in enumerate(dataloader):
|
|
soc_semantic_loss, soc_detail_loss = \
|
|
soc_adaptation_iter(modnet, backup_modnet, optimizer, image)
|
|
"""
|
|
|
|
global blurer
|
|
|
|
# set the backup model to eval mode
|
|
backup_modnet.eval()
|
|
|
|
# set the main model to train mode and freeze its norm layers
|
|
modnet.train()
|
|
modnet.module.freeze_norm()
|
|
|
|
# clear the optimizer
|
|
optimizer.zero_grad()
|
|
|
|
# forward the main model
|
|
pred_semantic, pred_detail, pred_matte = modnet(image, False)
|
|
|
|
# forward the backup model
|
|
with torch.no_grad():
|
|
_, pred_backup_detail, pred_backup_matte = backup_modnet(image, False)
|
|
|
|
# calculate the boundary mask from `pred_matte` and `pred_semantic`
|
|
pred_matte_fg = (pred_matte.detach() > 0.1).float()
|
|
pred_semantic_fg = (pred_semantic.detach() > 0.1).float()
|
|
pred_semantic_fg = F.interpolate(pred_semantic_fg, scale_factor=16, mode='bilinear')
|
|
pred_fg = pred_matte_fg * pred_semantic_fg
|
|
|
|
n, c, h, w = pred_matte.shape
|
|
np_pred_fg = pred_fg.data.cpu().numpy()
|
|
np_boundaries = np.zeros([n, c, h, w])
|
|
for sdx in range(0, n):
|
|
sample_np_boundaries = np_boundaries[sdx, 0, ...]
|
|
sample_np_pred_fg = np_pred_fg[sdx, 0, ...]
|
|
|
|
side = int((h + w) / 2 * 0.05)
|
|
dilated = grey_dilation(sample_np_pred_fg, size=(side, side))
|
|
eroded = grey_erosion(sample_np_pred_fg, size=(side, side))
|
|
|
|
sample_np_boundaries[np.where(dilated - eroded != 0)] = 1
|
|
np_boundaries[sdx, 0, ...] = sample_np_boundaries
|
|
|
|
boundaries = torch.tensor(np_boundaries).float().cuda()
|
|
|
|
# sub-objectives consistency between `pred_semantic` and `pred_matte`
|
|
# generate pseudo ground truth for `pred_semantic`
|
|
downsampled_pred_matte = blurer(F.interpolate(pred_matte, scale_factor=1/16, mode='bilinear'))
|
|
pseudo_gt_semantic = downsampled_pred_matte.detach()
|
|
pseudo_gt_semantic = pseudo_gt_semantic * (pseudo_gt_semantic > 0.01).float()
|
|
|
|
# generate pseudo ground truth for `pred_matte`
|
|
pseudo_gt_matte = pred_semantic.detach()
|
|
pseudo_gt_matte = pseudo_gt_matte * (pseudo_gt_matte > 0.01).float()
|
|
|
|
# calculate the SOC semantic loss
|
|
soc_semantic_loss = F.mse_loss(pred_semantic, pseudo_gt_semantic) + F.mse_loss(downsampled_pred_matte, pseudo_gt_matte)
|
|
soc_semantic_loss = soc_semantic_scale * torch.mean(soc_semantic_loss)
|
|
|
|
# NOTE: using the formulas in our paper to calculate the following losses has similar results
|
|
# sub-objectives consistency between `pred_detail` and `pred_backup_detail` (on boundaries only)
|
|
backup_detail_loss = boundaries * F.l1_loss(pred_detail, pred_backup_detail, reduction='none')
|
|
backup_detail_loss = torch.sum(backup_detail_loss, dim=(1,2,3)) / torch.sum(boundaries, dim=(1,2,3))
|
|
backup_detail_loss = torch.mean(backup_detail_loss)
|
|
|
|
# sub-objectives consistency between pred_matte` and `pred_backup_matte` (on boundaries only)
|
|
backup_matte_loss = boundaries * F.l1_loss(pred_matte, pred_backup_matte, reduction='none')
|
|
backup_matte_loss = torch.sum(backup_matte_loss, dim=(1,2,3)) / torch.sum(boundaries, dim=(1,2,3))
|
|
backup_matte_loss = torch.mean(backup_matte_loss)
|
|
|
|
soc_detail_loss = soc_detail_scale * (backup_detail_loss + backup_matte_loss)
|
|
|
|
# calculate the final loss, backward the loss, and update the model
|
|
loss = soc_semantic_loss + soc_detail_loss
|
|
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
return soc_semantic_loss, soc_detail_loss
|
|
|
|
# ----------------------------------------------------------------------------------
|