初始化换发型项目：3个微服务代码 + 部署脚本
包含： - hair_service_sd: 换发型/换发色算法服务 (端口 8801) - photo_service: LoRA 训练调度服务 (端口 32678) - stable-diffusion-webui: SD WebUI 推理服务 (端口 57860) - kohya_ss_home: 训练环境代码 - meidaojia: 监控测试脚本 - setup.sh: 一键部署脚本 (conda环境恢复 + 配置生成 + 完整性检查) - start_all_services.sh: 启动3个服务 - configure.ini.template: 路径模板化 (BASE_DIR自动推导) - conda_envs/py310.yml: py310 环境定义 大文件 (weights/, models/, data/, conda_envs/*.tar.gz 等) 通过 .gitignore 排除， 由网盘单独上传。
This commit is contained in:
colomi
@@ -0,0 +1,75 @@
|
||||
import os
|
||||
import pickle
|
||||
|
||||
import torch
|
||||
from faceseg.u2net import U2NET
|
||||
from utils import landmark_processor
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
class FaceSeg:
|
||||
def __init__(self, gpu_id = 0):
|
||||
model = U2NET(in_ch=4, out_ch=1)
|
||||
weights = torch.load('weights/20210927_01.pth', map_location='cpu')
|
||||
model_dict = model.state_dict()
|
||||
pretrained_dict = {}
|
||||
for ix, (k, v) in enumerate(model_dict.items()):
|
||||
if k in weights and weights[k].data.shape == v.data.shape:
|
||||
pretrained_dict[k] = weights[k]
|
||||
else:
|
||||
print('ignore {}'.format(k))
|
||||
model_dict.update(pretrained_dict)
|
||||
model.load_state_dict(model_dict)
|
||||
print('update success')
|
||||
model.cuda(gpu_id)
|
||||
model.eval()
|
||||
self.model = model
|
||||
self.last_mask = None
|
||||
self.output_img_size = 320
|
||||
self.gpu_id = gpu_id
|
||||
|
||||
def inference(self, frame, pt1k, video_mode=False):
|
||||
image_to_face_mat = landmark_processor.get_transform_mat_full_face(pt1k, self.output_img_size)
|
||||
face_image = cv2.warpAffine(frame, image_to_face_mat, (self.output_img_size, self.output_img_size), flags=cv2.INTER_LANCZOS4)
|
||||
if face_image.dtype == np.uint8: face_image = face_image.astype(np.float32) / 255
|
||||
if video_mode and self.last_mask is not None:
|
||||
last_small_mask = cv2.warpAffine(self.last_mask, image_to_face_mat,
|
||||
(self.output_img_size, self.output_img_size), flags=cv2.INTER_LANCZOS4)[:,:,np.newaxis]
|
||||
input_img = np.concatenate([face_image, last_small_mask], axis=2)
|
||||
else:
|
||||
zero_mask = np.zeros((face_image.shape[1], face_image.shape[0], 1), dtype=np.float32)
|
||||
input_img = np.concatenate([face_image, zero_mask], axis=2)
|
||||
|
||||
face_image_tensor = input_img.transpose((2, 0, 1))[np.newaxis]
|
||||
face_image_tensor = torch.from_numpy(face_image_tensor).cuda(self.gpu_id)
|
||||
|
||||
mask = self.model.test(face_image_tensor)
|
||||
mask = mask[0].detach().cpu().numpy().transpose((1, 2, 0))
|
||||
|
||||
origin_mask = cv2.warpAffine(mask, image_to_face_mat, (frame.shape[1], frame.shape[0]),
|
||||
flags=cv2.WARP_INVERSE_MAP|cv2.INTER_LANCZOS4)[:, :, np.newaxis]
|
||||
if video_mode: self.last_mask = origin_mask.copy()
|
||||
|
||||
return origin_mask
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
face_segmentor = FaceSeg(gpu_id=0)
|
||||
testdata_dir = "/mnt/DataDisk/my_projects/faceswap_hq/train_data/example"
|
||||
for picname in os.listdir(testdata_dir):
|
||||
img_path = os.path.join(testdata_dir, picname)
|
||||
pkl_path = img_path[:-4]+".pkl"
|
||||
if not picname.endswith(".jpg"):
|
||||
continue
|
||||
if not os.path.exists(pkl_path):
|
||||
continue
|
||||
img = cv2.imread(img_path)
|
||||
with open(pkl_path, "rb") as fp:
|
||||
info = pickle.load(fp)
|
||||
pt1k = info["human_pt1k"]
|
||||
|
||||
face_seg_mask = face_segmentor.inference(img, pt1k, video_mode=False)
|
||||
cv2.imshow("face_seg_mask", face_seg_mask)
|
||||
cv2.imshow("img", img)
|
||||
cv2.waitKey()
|
||||
@@ -0,0 +1,299 @@
|
||||
from __future__ import division
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torch.nn.init as init
|
||||
import torch.utils.model_zoo as model_zoo
|
||||
from torchvision import models
|
||||
|
||||
# general libs
|
||||
import cv2
|
||||
import matplotlib.pyplot as plt
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
import math
|
||||
import time
|
||||
import tqdm
|
||||
import os
|
||||
import argparse
|
||||
import copy
|
||||
import sys
|
||||
|
||||
from utils.helpers import *
|
||||
|
||||
|
||||
class ResBlock(nn.Module):
|
||||
def __init__(self, backbone, indim, outdim=None, stride=1):
|
||||
super(ResBlock, self).__init__()
|
||||
self.backbone = backbone
|
||||
if outdim == None:
|
||||
outdim = indim
|
||||
if indim == outdim and stride == 1:
|
||||
self.downsample = None
|
||||
else:
|
||||
self.downsample = nn.Conv2d(indim, outdim, kernel_size=3, padding=1, stride=stride)
|
||||
|
||||
self.conv1 = nn.Conv2d(indim, outdim, kernel_size=3, padding=1, stride=stride)
|
||||
self.conv2 = nn.Conv2d(outdim, outdim, kernel_size=3, padding=1)
|
||||
|
||||
def forward(self, x):
|
||||
if self.backbone == 'resnest101':
|
||||
r = self.conv1(F.relu(x, inplace=True))
|
||||
r = self.conv2(F.relu(r, inplace=True))
|
||||
else:
|
||||
r = self.conv1(F.relu(x))
|
||||
r = self.conv2(F.relu(r))
|
||||
|
||||
if self.downsample is not None:
|
||||
x = self.downsample(x)
|
||||
|
||||
return x + r
|
||||
|
||||
|
||||
class Encoder_M(nn.Module):
|
||||
def __init__(self, backbone):
|
||||
super(Encoder_M, self).__init__()
|
||||
if backbone == 'resnest101':
|
||||
self.conv1_m = nn.Conv2d(1, 128, kernel_size=7, stride=2, padding=3, bias=False)
|
||||
self.conv1_o = nn.Conv2d(1, 128, kernel_size=7, stride=2, padding=3, bias=False)
|
||||
else:
|
||||
self.conv1_m = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
|
||||
self.conv1_o = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
|
||||
|
||||
if backbone == 'resnet50':
|
||||
resnet = models.resnet50(pretrained=True)
|
||||
elif backbone == 'resnet18':
|
||||
resnet = models.resnet18(pretrained=True)
|
||||
|
||||
self.conv1 = resnet.conv1
|
||||
self.bn1 = resnet.bn1
|
||||
self.relu = resnet.relu # 1/2, 64
|
||||
self.maxpool = resnet.maxpool
|
||||
|
||||
self.res2 = resnet.layer1 # 1/4, 256
|
||||
self.res3 = resnet.layer2 # 1/8, 512
|
||||
self.res4 = resnet.layer3 # 1/8, 1024
|
||||
|
||||
self.register_buffer('mean', torch.FloatTensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
|
||||
self.register_buffer('std', torch.FloatTensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))
|
||||
|
||||
def forward(self, in_f, in_m, in_o):
|
||||
f = (in_f - self.mean) / self.std
|
||||
m = torch.unsqueeze(in_m, dim=1).float() # add channel dim
|
||||
o = torch.unsqueeze(in_o, dim=1).float() # add channel dim
|
||||
|
||||
x = self.conv1(f) + self.conv1_m(m) + self.conv1_o(o)
|
||||
x = self.bn1(x)
|
||||
c1 = self.relu(x) # 1/2, 64
|
||||
x = self.maxpool(c1) # 1/4, 64
|
||||
r2 = self.res2(x) # 1/4, 256
|
||||
r3 = self.res3(r2) # 1/8, 512
|
||||
r4 = self.res4(r3) # 1/8, 1024
|
||||
return r4, r3, r2, c1, f
|
||||
|
||||
|
||||
class Encoder_Q(nn.Module):
|
||||
def __init__(self, backbone):
|
||||
super(Encoder_Q, self).__init__()
|
||||
|
||||
if backbone == 'resnet50':
|
||||
resnet = models.resnet50(pretrained=True)
|
||||
elif backbone == 'resnet18':
|
||||
resnet = models.resnet18(pretrained=True)
|
||||
|
||||
self.conv1 = resnet.conv1
|
||||
self.bn1 = resnet.bn1
|
||||
self.relu = resnet.relu # 1/2, 64
|
||||
self.maxpool = resnet.maxpool
|
||||
|
||||
self.res2 = resnet.layer1 # 1/4, 256
|
||||
self.res3 = resnet.layer2 # 1/8, 512
|
||||
self.res4 = resnet.layer3 # 1/8, 1024
|
||||
|
||||
self.register_buffer('mean', torch.FloatTensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
|
||||
self.register_buffer('std', torch.FloatTensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))
|
||||
|
||||
def forward(self, in_f):
|
||||
f = (in_f - self.mean) / self.std
|
||||
|
||||
x = self.conv1(f)
|
||||
x = self.bn1(x)
|
||||
c1 = self.relu(x) # 1/2, 64
|
||||
x = self.maxpool(c1) # 1/4, 64
|
||||
r2 = self.res2(x) # 1/4, 256
|
||||
r3 = self.res3(r2) # 1/8, 512
|
||||
r4 = self.res4(r3) # 1/8, 1024
|
||||
return r4, r3, r2, c1, f
|
||||
|
||||
|
||||
class Refine(nn.Module):
|
||||
def __init__(self, backbone, inplanes, planes, scale_factor=2):
|
||||
super(Refine, self).__init__()
|
||||
self.convFS = nn.Conv2d(inplanes, planes, kernel_size=(3, 3), padding=(1, 1), stride=1)
|
||||
self.ResFS = ResBlock(backbone, planes, planes)
|
||||
self.ResMM = ResBlock(backbone, planes, planes)
|
||||
self.scale_factor = scale_factor
|
||||
|
||||
def forward(self, f, pm):
|
||||
s = self.ResFS(self.convFS(f))
|
||||
m = s + F.interpolate(pm, scale_factor=self.scale_factor, mode='bilinear', align_corners=False)
|
||||
m = self.ResMM(m)
|
||||
return m
|
||||
|
||||
|
||||
class Decoder(nn.Module):
|
||||
def __init__(self, mdim, scale_rate, backbone):
|
||||
super(Decoder, self).__init__()
|
||||
self.backbone = backbone
|
||||
if backbone == 'resnest101':
|
||||
self.convFM = nn.Conv2d(256, mdim, kernel_size=(3, 3), padding=(1, 1), stride=1)
|
||||
else:
|
||||
self.convFM = nn.Conv2d(1024 // scale_rate, mdim, kernel_size=(3, 3), padding=(1, 1), stride=1)
|
||||
self.ResMM = ResBlock(backbone, mdim, mdim)
|
||||
self.RF3 = Refine(backbone, 512 // scale_rate, mdim) # 1/8 -> 1/4
|
||||
self.RF2 = Refine(backbone, 256 // scale_rate, mdim) # 1/4 -> 1
|
||||
|
||||
self.pred2 = nn.Conv2d(mdim, 2, kernel_size=(3, 3), padding=(1, 1), stride=1)
|
||||
|
||||
def forward(self, r4, r3, r2):
|
||||
m4 = self.ResMM(self.convFM(r4))
|
||||
m3 = self.RF3(r3, m4) # out: 1/8, 256
|
||||
m2 = self.RF2(r2, m3) # out: 1/4, 256
|
||||
|
||||
if self.backbone == 'resnest101':
|
||||
p2 = self.pred2(F.relu(m2, inplace=True))
|
||||
else:
|
||||
p2 = self.pred2(F.relu(m2))
|
||||
|
||||
p = F.interpolate(p2, scale_factor=4, mode='bilinear', align_corners=False)
|
||||
return p # , p2, p3, p4
|
||||
|
||||
|
||||
class Memory(nn.Module):
|
||||
def __init__(self):
|
||||
super(Memory, self).__init__()
|
||||
|
||||
def forward(self, m_in, m_out, q_in, q_out): # m_in: o,c,t,h,w
|
||||
B, D_e, T, H, W = m_in.size()
|
||||
_, D_o, _, _, _ = m_out.size()
|
||||
|
||||
mi = m_in.view(B, D_e, T * H * W)
|
||||
mi = torch.transpose(mi, 1, 2) # b, THW, emb
|
||||
|
||||
qi = q_in.view(B, D_e, H * W) # b, emb, HW
|
||||
|
||||
p = torch.bmm(mi, qi) # b, THW, HW
|
||||
p = p / math.sqrt(D_e)
|
||||
p = F.softmax(p, dim=1) # b, THW, HW
|
||||
|
||||
mo = m_out.view(B, D_o, T * H * W)
|
||||
mem = torch.bmm(mo, p) # Weighted-sum B, D_o, HW
|
||||
mem = mem.view(B, D_o, H, W)
|
||||
|
||||
mem_out = torch.cat([mem, q_out], dim=1)
|
||||
|
||||
return mem_out, p
|
||||
|
||||
|
||||
class KeyValue(nn.Module):
|
||||
# Not using location
|
||||
def __init__(self, indim, keydim, valdim):
|
||||
super(KeyValue, self).__init__()
|
||||
self.Key = nn.Conv2d(indim, keydim, kernel_size=(3, 3), padding=(1, 1), stride=1)
|
||||
self.Value = nn.Conv2d(indim, valdim, kernel_size=(3, 3), padding=(1, 1), stride=1)
|
||||
|
||||
def forward(self, x):
|
||||
return self.Key(x), self.Value(x)
|
||||
|
||||
|
||||
class STM(nn.Module):
|
||||
def __init__(self, backbone='resnet50'):
|
||||
super(STM, self).__init__()
|
||||
self.backbone = backbone
|
||||
assert backbone == 'resnet50' or backbone == 'resnet18' or backbone == 'resnest101'
|
||||
scale_rate = (1 if (backbone == 'resnet50' or backbone == 'resnest101') else 4)
|
||||
|
||||
self.Encoder_M = Encoder_M(backbone)
|
||||
self.Encoder_Q = Encoder_Q(backbone)
|
||||
|
||||
self.KV_M_r4 = KeyValue(1024 // scale_rate, keydim=128 // scale_rate, valdim=512 // scale_rate)
|
||||
self.KV_Q_r4 = KeyValue(1024 // scale_rate, keydim=128 // scale_rate, valdim=512 // scale_rate)
|
||||
|
||||
self.Memory = Memory()
|
||||
self.Decoder = Decoder(256, scale_rate, backbone)
|
||||
|
||||
def Pad_memory(self, mems, num_objects, K):
|
||||
pad_mems = []
|
||||
for mem in mems:
|
||||
pad_mem = ToCuda(torch.zeros(1, K, mem.size()[1], 1, mem.size()[2], mem.size()[3]))
|
||||
pad_mem[0, 1:num_objects + 1, :, 0] = mem
|
||||
pad_mems.append(pad_mem)
|
||||
return pad_mems
|
||||
|
||||
def memorize(self, frame, masks, num_objects):
|
||||
# memorize a frame
|
||||
num_objects = num_objects[0].item()
|
||||
_, K, H, W = masks.shape # B = 1
|
||||
|
||||
(frame, masks), pad = pad_divide_by([frame, masks], 16, (frame.size()[2], frame.size()[3]))
|
||||
|
||||
# make batch arg list
|
||||
B_list = {'f': [], 'm': [], 'o': []}
|
||||
for o in range(1, num_objects + 1): # 1 - no
|
||||
B_list['f'].append(frame)
|
||||
B_list['m'].append(masks[:, o])
|
||||
B_list['o'].append((torch.sum(masks[:, 1:o], dim=1) + \
|
||||
torch.sum(masks[:, o + 1:num_objects + 1], dim=1)).clamp(0, 1))
|
||||
|
||||
# make Batch
|
||||
B_ = {}
|
||||
for arg in B_list.keys():
|
||||
B_[arg] = torch.cat(B_list[arg], dim=0)
|
||||
|
||||
r4, _, _, _, _ = self.Encoder_M(B_['f'], B_['m'], B_['o'])
|
||||
k4, v4 = self.KV_M_r4(r4) # num_objects, 128 and 512, H/16, W/16
|
||||
k4, v4 = self.Pad_memory([k4, v4], num_objects=num_objects, K=K)
|
||||
return k4, v4
|
||||
|
||||
def Soft_aggregation(self, ps, K):
|
||||
num_objects, H, W = ps.shape
|
||||
em = ToCuda(torch.zeros(1, K, H, W))
|
||||
em[0, 0] = torch.prod(1 - ps, dim=0) # bg prob
|
||||
em[0, 1:num_objects + 1] = ps # obj prob
|
||||
em = torch.clamp(em, 1e-7, 1 - 1e-7)
|
||||
logit = torch.log((em / (1 - em)))
|
||||
return logit
|
||||
|
||||
def segment(self, frame, keys, values, num_objects):
|
||||
num_objects = num_objects[0].item()
|
||||
_, K, keydim, T, H, W = keys.shape # B = 1
|
||||
# pad
|
||||
[frame], pad = pad_divide_by([frame], 16, (frame.size()[2], frame.size()[3]))
|
||||
|
||||
r4, r3, r2, _, _ = self.Encoder_Q(frame)
|
||||
k4, v4 = self.KV_Q_r4(r4) # 1, dim, H/16, W/16
|
||||
|
||||
# expand to --- no, c, h, w
|
||||
k4e, v4e = k4.expand(num_objects, -1, -1, -1), v4.expand(num_objects, -1, -1, -1)
|
||||
r3e, r2e = r3.expand(num_objects, -1, -1, -1), r2.expand(num_objects, -1, -1, -1)
|
||||
|
||||
# memory select kv:(1, K, C, T, H, W)
|
||||
m4, viz = self.Memory(keys[0, 1:num_objects + 1], values[0, 1:num_objects + 1], k4e, v4e)
|
||||
logits = self.Decoder(m4, r3e, r2e)
|
||||
ps = F.softmax(logits, dim=1)[:, 1] # no, h, w
|
||||
# ps = indipendant possibility to belong to each object
|
||||
|
||||
logit = self.Soft_aggregation(ps, K) # 1, K, H, W
|
||||
|
||||
if pad[2] + pad[3] > 0:
|
||||
logit = logit[:, :, pad[2]:-pad[3], :]
|
||||
if pad[0] + pad[1] > 0:
|
||||
logit = logit[:, :, :, pad[0]:-pad[1]]
|
||||
|
||||
return logit
|
||||
|
||||
def forward(self, *args, **kwargs):
|
||||
if args[1].dim() > 4: # keys
|
||||
return self.segment(*args, **kwargs)
|
||||
else:
|
||||
return self.memorize(*args, **kwargs)
|
||||
@@ -0,0 +1,185 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
import numpy as np
|
||||
|
||||
class SequenceConv(nn.ModuleList):
|
||||
"""Sequence conv module.
|
||||
|
||||
Args:
|
||||
in_channels (int): input tensor channel.
|
||||
out_channels (int): output tensor channel.
|
||||
kernel_size (int): convolution kernel size.
|
||||
sequence_num (int): sequence length.
|
||||
conv_cfg (dict): convolution config dictionary.
|
||||
norm_cfg (dict): normalization config dictionary.
|
||||
act_cfg (dict): activation config dictionary.
|
||||
"""
|
||||
|
||||
def __init__(self, in_channels, out_channels, kernel_size, sequence_num):
|
||||
super(SequenceConv, self).__init__()
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
self.kernel_size = kernel_size
|
||||
self.sequence_num = sequence_num
|
||||
for _ in range(sequence_num):
|
||||
self.append(
|
||||
nn.Sequential(
|
||||
nn.Conv2d(self.in_channels, self.out_channels, self.kernel_size, 1, self.kernel_size // 2, bias=False),
|
||||
nn.BatchNorm2d(self.out_channels),
|
||||
nn.ReLU()
|
||||
)
|
||||
)
|
||||
|
||||
def forward(self, sequence_imgs):
|
||||
"""
|
||||
|
||||
Args:
|
||||
sequence_imgs (Tensor): TxBxCxHxW
|
||||
|
||||
Returns:
|
||||
sequence conv output: TxBxCxHxW
|
||||
"""
|
||||
sequence_outs = []
|
||||
assert sequence_imgs.shape[0] == self.sequence_num
|
||||
for i, sequence_conv in enumerate(self):
|
||||
sequence_out = sequence_conv(sequence_imgs[i, ...])
|
||||
sequence_out = sequence_out.unsqueeze(0)
|
||||
sequence_outs.append(sequence_out)
|
||||
|
||||
sequence_outs = torch.cat(sequence_outs, dim=0) # TxBxCxHxW
|
||||
return sequence_outs
|
||||
|
||||
class MemoryModule(nn.Module):
|
||||
"""Memory read module.
|
||||
Args:
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
matmul_norm=False):
|
||||
super(MemoryModule, self).__init__()
|
||||
self.matmul_norm = matmul_norm
|
||||
|
||||
def forward(self, memory_keys, memory_values, query_key, query_value):
|
||||
"""
|
||||
Memory Module forward.
|
||||
Args:
|
||||
memory_keys (Tensor): memory keys tensor, shape: TxBxCxHxW
|
||||
memory_values (Tensor): memory values tensor, shape: TxBxCxHxW
|
||||
query_key (Tensor): query keys tensor, shape: BxCxHxW
|
||||
query_value (Tensor): query values tensor, shape: BxCxHxW
|
||||
|
||||
Returns:
|
||||
Concat query and memory tensor.
|
||||
"""
|
||||
sequence_num, batch_size, key_channels, height, width = memory_keys.shape
|
||||
_, _, value_channels, _, _ = memory_values.shape
|
||||
assert query_key.shape[1] == key_channels and query_value.shape[1] == value_channels
|
||||
memory_keys = memory_keys.permute(1, 2, 0, 3, 4).contiguous() # BxCxTxHxW
|
||||
memory_keys = memory_keys.view(batch_size, key_channels, sequence_num * height * width) # BxCxT*H*W
|
||||
|
||||
query_key = query_key.view(batch_size, key_channels, height * width).permute(0, 2, 1).contiguous() # BxH*WxCk
|
||||
key_attention = torch.bmm(query_key, memory_keys) # BxH*WxT*H*W
|
||||
if self.matmul_norm:
|
||||
key_attention = (key_channels ** -.5) * key_attention
|
||||
key_attention = F.softmax(key_attention, dim=-1) # BxH*WxT*H*W
|
||||
|
||||
memory_values = memory_values.permute(1, 2, 0, 3, 4).contiguous() # BxCxTxHxW
|
||||
memory_values = memory_values.view(batch_size, value_channels, sequence_num * height * width)
|
||||
memory_values = memory_values.permute(0, 2, 1).contiguous() # BxT*H*WxC
|
||||
memory = torch.bmm(key_attention, memory_values) # BxH*WxC
|
||||
memory = memory.permute(0, 2, 1).contiguous() # BxCxH*W
|
||||
memory = memory.view(batch_size, value_channels, height, width) # BxCxHxW
|
||||
|
||||
query_memory = torch.cat([query_value, memory], dim=1)
|
||||
return query_memory
|
||||
#
|
||||
# class TMAHead(nn.Module):
|
||||
# """TMAHead decoder for video semantic segmentation."""
|
||||
#
|
||||
# def __init__(self, sequence_num, key_channels, value_channels, num_classes=2, dropout_ratio=0):
|
||||
# super(TMAHead, self).__init__()
|
||||
#
|
||||
# self.sequence_num = sequence_num
|
||||
# self.memory_key_conv = nn.Sequential(
|
||||
# SequenceConv(self.in_channels, key_channels, 1, sequence_num),
|
||||
# SequenceConv(key_channels, key_channels, 3, sequence_num)
|
||||
# )
|
||||
# self.memory_value_conv = nn.Sequential(
|
||||
# SequenceConv(self.in_channels, value_channels, 1, sequence_num),
|
||||
# SequenceConv(value_channels, value_channels, 3, sequence_num)
|
||||
# )
|
||||
# self.query_key_conv = nn.Sequential(
|
||||
# nn.Sequential(
|
||||
# nn.Conv2d(self.in_channels, key_channels, 1, 1, 0, bias=False),
|
||||
# nn.BatchNorm2d(key_channels),
|
||||
# nn.ReLU()
|
||||
# ),
|
||||
# nn.Sequential(
|
||||
# nn.Conv2d(key_channels, key_channels, 3, 1, 1, bias=False),
|
||||
# nn.BatchNorm2d(key_channels),
|
||||
# nn.ReLU()
|
||||
# ),
|
||||
# )
|
||||
#
|
||||
# self.query_value_conv = nn.Sequential(
|
||||
# nn.Sequential(
|
||||
# nn.Conv2d(self.in_channels, value_channels, 1, 1, 0, bias=False),
|
||||
# nn.BatchNorm2d(value_channels),
|
||||
# nn.ReLU()
|
||||
# ),
|
||||
# nn.Sequential(
|
||||
# nn.Conv2d(value_channels, value_channels, 3, 1, 1, bias=False),
|
||||
# nn.BatchNorm2d(value_channels),
|
||||
# nn.ReLU()
|
||||
# ),
|
||||
# )
|
||||
# self.memory_module = MemoryModule(matmul_norm=False)
|
||||
# self.bottleneck = nn.Sequential(
|
||||
# nn.Conv2d(value_channels * 2, self.channels, 3, 1, 1, bias=False),
|
||||
# nn.BatchNorm2d(value_channels),
|
||||
# nn.ReLU()
|
||||
# )
|
||||
#
|
||||
# self.conv_seg = nn.Conv2d(self.channels, num_classes, kernel_size=1)
|
||||
# if dropout_ratio > 0:
|
||||
# self.dropout = nn.Dropout2d(dropout_ratio)
|
||||
# else:
|
||||
# self.dropout = None
|
||||
#
|
||||
# def cls_seg(self, feat):
|
||||
# """Classify each pixel."""
|
||||
# if self.dropout is not None:
|
||||
# feat = self.dropout(feat)
|
||||
# output = self.conv_seg(feat)
|
||||
# return output
|
||||
#
|
||||
# def forward(self, inputs, sequence_imgs):
|
||||
# """
|
||||
# Forward fuction.
|
||||
# Args:
|
||||
# inputs (list[Tensor]): backbone multi-level outputs.
|
||||
# sequence_imgs (list[Tensor]): len(sequence_imgs) is equal to batch_size,
|
||||
# each element is a Tensor with shape of TxCxHxW.
|
||||
#
|
||||
# Returns:
|
||||
# decoder logits.
|
||||
# """
|
||||
# x = inputs
|
||||
# sequence_imgs = [y.unsqueeze(0) for y in sequence_imgs] # T, BxCxHxW
|
||||
# sequence_imgs = torch.cat(sequence_imgs, dim=0) # TxBxCxHxW
|
||||
# sequence_num, batch_size, channels, height, width = sequence_imgs.shape
|
||||
#
|
||||
# assert sequence_num == self.sequence_num
|
||||
# memory_keys = self.memory_key_conv(sequence_imgs)
|
||||
# memory_values = self.memory_value_conv(sequence_imgs)
|
||||
# query_key = self.query_key_conv(x) # BxCxHxW
|
||||
# query_value = self.query_value_conv(x) # BxCxHxW
|
||||
#
|
||||
# # memory read
|
||||
# output = self.memory_module(memory_keys, memory_values, query_key, query_value)
|
||||
# output = self.bottleneck(output)
|
||||
# output = self.cls_seg(output)
|
||||
#
|
||||
# return output
|
||||
@@ -0,0 +1,624 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
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
|
||||
|
||||
from faceseg.tma import SequenceConv, MemoryModule
|
||||
|
||||
##### 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,out_ch,1)
|
||||
|
||||
self.in_channels = 512
|
||||
key_channels = 128
|
||||
value_channels = 512
|
||||
self.sequence_num = sequence_num = 2
|
||||
self.memory_key_conv = nn.Sequential(
|
||||
SequenceConv(self.in_channels, key_channels, 1, sequence_num),
|
||||
SequenceConv(key_channels, key_channels, 3, sequence_num)
|
||||
)
|
||||
self.memory_value_conv = nn.Sequential(
|
||||
SequenceConv(self.in_channels, value_channels, 1, sequence_num),
|
||||
SequenceConv(value_channels, value_channels, 3, sequence_num)
|
||||
)
|
||||
self.query_key_conv = nn.Sequential(
|
||||
nn.Sequential(
|
||||
nn.Conv2d(self.in_channels, key_channels, 1, 1, 0, bias=False),
|
||||
nn.BatchNorm2d(key_channels),
|
||||
nn.ReLU()
|
||||
),
|
||||
nn.Sequential(
|
||||
nn.Conv2d(key_channels, key_channels, 3, 1, 1, bias=False),
|
||||
nn.BatchNorm2d(key_channels),
|
||||
nn.ReLU()
|
||||
),
|
||||
)
|
||||
self.query_value_conv = nn.Sequential(
|
||||
nn.Sequential(
|
||||
nn.Conv2d(self.in_channels, value_channels, 1, 1, 0, bias=False),
|
||||
nn.BatchNorm2d(value_channels),
|
||||
nn.ReLU()
|
||||
),
|
||||
nn.Sequential(
|
||||
nn.Conv2d(value_channels, value_channels, 3, 1, 1, bias=False),
|
||||
nn.BatchNorm2d(value_channels),
|
||||
nn.ReLU()
|
||||
),
|
||||
)
|
||||
self.memory_module = MemoryModule(matmul_norm=False)
|
||||
self.bottleneck = nn.Sequential(
|
||||
nn.Conv2d(value_channels * 2, self.in_channels, 3, 1, 1, bias=False),
|
||||
nn.BatchNorm2d(value_channels),
|
||||
nn.ReLU()
|
||||
)
|
||||
|
||||
self.is_train = True
|
||||
|
||||
def extract_feature(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)
|
||||
|
||||
return hx1, hx2, hx3, hx4, hx5, hx6
|
||||
|
||||
def decoder(self, hx1, hx2, hx3, hx4, hx5, hx6):
|
||||
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))
|
||||
|
||||
return hx1d, hx2d, hx3d, hx4d, hx5d
|
||||
|
||||
def side_output(self, hx1d, hx2d, hx3d, hx4d, hx5d, hx6):
|
||||
# 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)
|
||||
|
||||
def forward(self, x, memory_sequence=None):
|
||||
if self.is_train:
|
||||
hx1, hx2, hx3, hx4, hx5, hx6 = self.extract_feature(x)
|
||||
|
||||
if memory_sequence is None:
|
||||
memory_hx6 = [hx6 for _ in range(self.sequence_num)]
|
||||
else:
|
||||
memory_hx6 = []
|
||||
for single_memory in memory_sequence:
|
||||
_, _, _, _, _, hx6 = self.extract_feature(single_memory)
|
||||
memory_hx6.append(hx6)
|
||||
|
||||
memory_hx6 = [mhx6.unsqueeze(0) for mhx6 in memory_hx6] # T, BxCxHxW
|
||||
memory_hx6 = torch.cat(memory_hx6, dim=0)
|
||||
memory_keys = self.memory_key_conv(memory_hx6)
|
||||
memory_values = self.memory_value_conv(memory_hx6)
|
||||
query_key = self.query_key_conv(hx6)
|
||||
query_value = self.query_value_conv(hx6)
|
||||
merge_hx6 = self.memory_module(memory_keys, memory_values, query_key, query_value)
|
||||
merge_hx6 = self.bottleneck(merge_hx6)
|
||||
|
||||
hx1d, hx2d, hx3d, hx4d, hx5d = self.decoder(hx1, hx2, hx3, hx4, hx5, merge_hx6)
|
||||
|
||||
return self.side_output(hx1d, hx2d, hx3d, hx4d, hx5d, hx6)
|
||||
else:
|
||||
return self.test(x)
|
||||
|
||||
def test(self, x):
|
||||
with torch.no_grad():
|
||||
hx1, hx2, hx3, hx4, hx5, hx6 = self.extract_feature(x)
|
||||
|
||||
memory_hx6 = [hx6 for _ in range(self.sequence_num)]
|
||||
|
||||
memory_hx6 = [mhx6.unsqueeze(0) for mhx6 in memory_hx6] # T, BxCxHxW
|
||||
memory_hx6 = torch.cat(memory_hx6, dim=0)
|
||||
memory_keys = self.memory_key_conv(memory_hx6)
|
||||
memory_values = self.memory_value_conv(memory_hx6)
|
||||
query_key = self.query_key_conv(hx6)
|
||||
query_value = self.query_value_conv(hx6)
|
||||
merge_hx6 = self.memory_module(memory_keys, memory_values, query_key, query_value)
|
||||
merge_hx6 = self.bottleneck(merge_hx6)
|
||||
|
||||
hx1d, hx2d, hx3d, hx4d, hx5d = self.decoder(hx1, hx2, hx3, hx4, hx5, merge_hx6)
|
||||
|
||||
mask, _, _, _, _, _, _ = self.side_output(hx1d, hx2d, hx3d, hx4d, hx5d, hx6)
|
||||
|
||||
return mask
|
||||
#
|
||||
# ### 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,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)
|
||||
Reference in New Issue
Block a user