191 lines
6.5 KiB
Python
191 lines
6.5 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
|
|
class ConvLayer(nn.Module):
|
|
|
|
def __init__(self, in_ch, out_ch):
|
|
super(ConvLayer, self).__init__()
|
|
|
|
self.conv = nn.Sequential(
|
|
nn.ReflectionPad2d(1),
|
|
nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=0),
|
|
nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True))
|
|
|
|
def forward(self, x):
|
|
x = self.conv(x)
|
|
return x
|
|
|
|
|
|
class SASA(nn.Module):
|
|
|
|
def __init__(self, in_dim):
|
|
super(SASA, self).__init__()
|
|
self.chanel_in = in_dim
|
|
|
|
self.query_conv = nn.Conv2d(
|
|
in_channels=in_dim, out_channels=in_dim // 8, kernel_size=1)
|
|
self.key_conv = nn.Conv2d(
|
|
in_channels=in_dim, out_channels=in_dim // 8, kernel_size=1)
|
|
self.value_conv = nn.Conv2d(
|
|
in_channels=in_dim, out_channels=in_dim, kernel_size=1)
|
|
self.mag_conv = nn.Conv2d(
|
|
in_channels=5, out_channels=in_dim // 32, kernel_size=1)
|
|
|
|
self.gamma = nn.Parameter(torch.zeros(1))
|
|
|
|
self.softmax = nn.Softmax(dim=-1) #
|
|
self.sigmoid = nn.Sigmoid()
|
|
|
|
def structure_encoder(self, paf_mag, target_height, target_width):
|
|
torso_mask = torch.sum(paf_mag[:, 1:3, :, :], dim=1, keepdim=True)
|
|
torso_mask = torch.clamp(torso_mask, 0, 1)
|
|
|
|
arms_mask = torch.sum(paf_mag[:, 4:8, :, :], dim=1, keepdim=True)
|
|
arms_mask = torch.clamp(arms_mask, 0, 1)
|
|
|
|
legs_mask = torch.sum(paf_mag[:, 8:12, :, :], dim=1, keepdim=True)
|
|
legs_mask = torch.clamp(legs_mask, 0, 1)
|
|
|
|
fg_mask = paf_mag[:, 12, :, :].unsqueeze(1)
|
|
bg_mask = 1 - fg_mask
|
|
Y = torch.cat((arms_mask, torso_mask, legs_mask, fg_mask, bg_mask),
|
|
dim=1)
|
|
Y = F.interpolate(Y, size=(target_height, target_width), mode='area')
|
|
return Y
|
|
|
|
def forward(self, X, PAF_mag):
|
|
"""extract self-attention features.
|
|
Args:
|
|
X : input feature maps( B x C x H x W)
|
|
PAF_mag : ( B x C x H x W), 1 denotes connectivity, 0 denotes non-connectivity
|
|
|
|
Returns:
|
|
out : self attention value + input feature
|
|
Y: B X N X N (N is Width*Height)
|
|
"""
|
|
|
|
m_batchsize, C, height, width = X.size()
|
|
|
|
Y = self.structure_encoder(PAF_mag, height, width)
|
|
|
|
connectivity_mask_vec = self.mag_conv(Y).view(m_batchsize, -1,
|
|
width * height)
|
|
affinity = torch.bmm(
|
|
connectivity_mask_vec.permute(0, 2, 1), connectivity_mask_vec)
|
|
affinity_centered = affinity - torch.mean(affinity)
|
|
affinity_sigmoid = self.sigmoid(affinity_centered)
|
|
|
|
proj_query = self.query_conv(X).view(m_batchsize, -1,
|
|
width * height).permute(0, 2, 1)
|
|
proj_key = self.key_conv(X).view(m_batchsize, -1, width * height)
|
|
selfatten_map = torch.bmm(proj_query, proj_key)
|
|
selfatten_centered = selfatten_map - torch.mean(
|
|
selfatten_map) # centering
|
|
selfatten_sigmoid = self.sigmoid(selfatten_centered)
|
|
|
|
SASA_map = selfatten_sigmoid * affinity_sigmoid
|
|
|
|
proj_value = self.value_conv(X).view(m_batchsize, -1, width * height)
|
|
|
|
out = torch.bmm(proj_value, SASA_map.permute(0, 2, 1))
|
|
out = out.view(m_batchsize, C, height, width)
|
|
|
|
out = self.gamma * out + X
|
|
return out, Y
|
|
|
|
|
|
class FlowGenerator(nn.Module):
|
|
|
|
def __init__(self, n_channels, deep_supervision=False):
|
|
super(FlowGenerator, self).__init__()
|
|
self.deep_supervision = deep_supervision
|
|
|
|
self.Encoder = nn.Sequential(
|
|
ConvLayer(n_channels, 64),
|
|
ConvLayer(64, 64),
|
|
nn.MaxPool2d(2),
|
|
ConvLayer(64, 128),
|
|
ConvLayer(128, 128),
|
|
nn.MaxPool2d(2),
|
|
ConvLayer(128, 256),
|
|
ConvLayer(256, 256),
|
|
nn.MaxPool2d(2),
|
|
ConvLayer(256, 512),
|
|
ConvLayer(512, 512),
|
|
nn.MaxPool2d(2),
|
|
ConvLayer(512, 1024),
|
|
ConvLayer(1024, 1024),
|
|
ConvLayer(1024, 1024),
|
|
ConvLayer(1024, 1024),
|
|
ConvLayer(1024, 1024),
|
|
)
|
|
|
|
self.SASA = SASA(in_dim=1024)
|
|
|
|
self.Decoder = nn.Sequential(
|
|
ConvLayer(1024, 1024),
|
|
nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
|
|
ConvLayer(1024, 512),
|
|
ConvLayer(512, 512),
|
|
nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
|
|
ConvLayer(512, 256),
|
|
ConvLayer(256, 256),
|
|
ConvLayer(256, 128),
|
|
ConvLayer(128, 64),
|
|
ConvLayer(64, 32),
|
|
nn.Conv2d(32, 2, kernel_size=1, padding=0),
|
|
nn.Tanh(),
|
|
nn.Upsample(scale_factor=4, mode='bilinear', align_corners=True),
|
|
)
|
|
|
|
dilation_ksize = 17
|
|
self.dilation = torch.nn.MaxPool2d(
|
|
kernel_size=dilation_ksize,
|
|
stride=1,
|
|
padding=int((dilation_ksize - 1) / 2))
|
|
|
|
def warp(self, x, flow, mode='bilinear', padding_mode='zeros', coff=0.2):
|
|
n, c, h, w = x.size()
|
|
yv, xv = torch.meshgrid([torch.arange(h), torch.arange(w)])
|
|
xv = xv.float() / (w - 1) * 2.0 - 1
|
|
yv = yv.float() / (h - 1) * 2.0 - 1
|
|
grid = torch.cat((xv.unsqueeze(-1), yv.unsqueeze(-1)), -1).unsqueeze(0)
|
|
grid = grid.to(flow.device)
|
|
grid_x = grid + 2 * flow * coff
|
|
warp_x = F.grid_sample(x, grid_x, mode=mode, padding_mode=padding_mode)
|
|
return warp_x
|
|
|
|
def forward(self, img, skeleton_map, coef=0.2):
|
|
"""extract self-attention features.
|
|
Args:
|
|
img : input numpy image
|
|
skeleton_map : skeleton map of input image
|
|
coef: warp degree
|
|
|
|
Returns:
|
|
warp_x : warped image
|
|
flow: predicted flow
|
|
"""
|
|
|
|
img_concat = torch.cat((img, skeleton_map), dim=1)
|
|
X = self.Encoder(img_concat)
|
|
|
|
_, _, height, width = X.size()
|
|
|
|
# directly get PAF magnitude from skeleton maps via dilation
|
|
PAF_mag = self.dilation((skeleton_map + 1.0) * 0.5)
|
|
|
|
out, Y = self.SASA(X, PAF_mag)
|
|
flow = self.Decoder(out)
|
|
|
|
flow = flow.permute(0, 2, 3, 1) # [n, 2, h, w] ==> [n, h, w, 2]
|
|
|
|
warp_x = self.warp(img, flow, coff=coef)
|
|
warp_x = torch.clamp(warp_x, min=-1.0, max=1.0)
|
|
|
|
return warp_x, flow
|