初始化换发型项目：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,86 @@
|
||||
import os
|
||||
import torch
|
||||
|
||||
from seg.networks.deeplabv3_plus import get_deeplabv3_plus
|
||||
import numpy as np
|
||||
import cv2
|
||||
import time
|
||||
from utils import landmark_processor
|
||||
|
||||
def label_to_mask(label_np):
|
||||
label_np = label_np.astype(np.int32)[:, :, np.newaxis]
|
||||
mask = np.zeros((label_np.shape[0], label_np.shape[1], 3), dtype=np.uint8)
|
||||
for id, color in enumerate(label_map):
|
||||
index = (label_np == id).all(axis=2)
|
||||
mask[index] = color
|
||||
return mask
|
||||
|
||||
label_map = [
|
||||
[0, 0, 0], #
|
||||
[128, 128, 128],
|
||||
[255, 255, 255],
|
||||
]
|
||||
|
||||
class Evaluator(object):
|
||||
def __init__(self, gpu_id, output_img_size, nclass, seg_model_path=None):
|
||||
self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
|
||||
# print("gpu_id: ", gpu_id)
|
||||
|
||||
# create network
|
||||
self.model = get_deeplabv3_plus(backbone='xception', nclass=nclass)
|
||||
model_path = os.path.join(seg_model_path)
|
||||
self.model.load_state_dict(torch.load(model_path, map_location=lambda storage, loc: storage))
|
||||
# print("seg device: ", self.model.device)
|
||||
self.model.to(self.device)
|
||||
self.model.eval()
|
||||
|
||||
# images = torch.randn((1, 3, 512, 512)).to(self.device)
|
||||
# torch.onnx.export(self.model, images,
|
||||
# "deeplabv3_hair512_360_0520_wl.onnx",
|
||||
# verbose=True,
|
||||
# opset_version=11,
|
||||
# input_names=['data'],
|
||||
# do_constant_folding=True,
|
||||
# output_names=['output'])
|
||||
|
||||
# exit()
|
||||
|
||||
self.output_img_size = output_img_size
|
||||
self.nclass = nclass
|
||||
|
||||
|
||||
def process_data(self, img):
|
||||
img = (img.astype(np.float32) / 255).transpose((2, 0, 1))
|
||||
img = torch.from_numpy(img).unsqueeze(0)
|
||||
|
||||
return img
|
||||
|
||||
def eval(self, img, pts1k):
|
||||
orig_h, orig_w, _ = img.shape
|
||||
|
||||
M1 = landmark_processor.get_transform_mat_hair(pts1k, self.output_img_size, ratio=0.28, h_ratio=0.3)
|
||||
crop_img = cv2.warpAffine(img, M1, (self.output_img_size, self.output_img_size), flags=cv2.INTER_LANCZOS4)
|
||||
crop_img = self.process_data(crop_img)
|
||||
crop_img = crop_img.to(self.device)
|
||||
|
||||
with torch.no_grad():
|
||||
# torch.cuda.synchronize()
|
||||
outputs = self.model(crop_img)
|
||||
pred = torch.argmax(outputs[0], 1)
|
||||
|
||||
pred = pred[0].detach().cpu().numpy()
|
||||
predict = pred.astype(np.float32)
|
||||
|
||||
pred_mask = label_to_mask(predict)
|
||||
|
||||
M1_invert = cv2.invertAffineTransform(M1)
|
||||
img_pred = cv2.warpAffine(pred_mask, M1_invert, (orig_w, orig_h), flags=cv2.INTER_CUBIC) #flags=cv2.INTER_NEAREST
|
||||
orig_mask = img_pred.copy()
|
||||
|
||||
# show_concat = np.concatenate((img, orig_mask), axis=1)
|
||||
# cv2.imshow("show_concat", show_concat)
|
||||
# cv2.waitKey()
|
||||
return orig_mask
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,462 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
__all__ = ['_ConvBNReLU', '_DWConvBNReLU', 'InvertedResidual', '_ASPP', '_FCNHead',
|
||||
'_Hswish', '_ConvBNHswish', 'SEModule', 'Bottleneck', 'ShuffleNetUnit',
|
||||
'ShuffleNetV2Unit', 'InvertedIGCV3', 'MBConvBlock']
|
||||
|
||||
|
||||
class _ConvBNReLU(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0,
|
||||
dilation=1, groups=1, relu6=False, norm_layer=nn.BatchNorm2d, **kwargs):
|
||||
super(_ConvBNReLU, self).__init__()
|
||||
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias=False)
|
||||
self.bn = norm_layer(out_channels)
|
||||
self.relu = nn.ReLU6(True) if relu6 else nn.ReLU(True)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv(x)
|
||||
x = self.bn(x)
|
||||
x = self.relu(x)
|
||||
return x
|
||||
|
||||
|
||||
class _FCNHead(nn.Module):
|
||||
def __init__(self, in_channels, channels, norm_layer=nn.BatchNorm2d, norm_kwargs=None, **kwargs):
|
||||
super(_FCNHead, self).__init__()
|
||||
inter_channels = in_channels // 4
|
||||
self.block = nn.Sequential(
|
||||
nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False),
|
||||
norm_layer(inter_channels, **({} if norm_kwargs is None else norm_kwargs)),
|
||||
nn.ReLU(True),
|
||||
nn.Dropout(0.1),
|
||||
nn.Conv2d(inter_channels, channels, 1)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return self.block(x)
|
||||
|
||||
|
||||
# -----------------------------------------------------------------
|
||||
# For MobileNet
|
||||
# -----------------------------------------------------------------
|
||||
class _DWConvBNReLU(nn.Module):
|
||||
"""Depthwise Separable Convolution in MobileNet.
|
||||
depthwise convolution + pointwise convolution
|
||||
"""
|
||||
|
||||
def __init__(self, in_channels, dw_channels, out_channels, stride, dilation=1, norm_layer=nn.BatchNorm2d, **kwargs):
|
||||
super(_DWConvBNReLU, self).__init__()
|
||||
self.conv = nn.Sequential(
|
||||
_ConvBNReLU(in_channels, dw_channels, 3, stride, dilation, dilation, in_channels, norm_layer=norm_layer),
|
||||
_ConvBNReLU(dw_channels, out_channels, 1, norm_layer=norm_layer))
|
||||
|
||||
def forward(self, x):
|
||||
return self.conv(x)
|
||||
|
||||
|
||||
# -----------------------------------------------------------------
|
||||
# For MobileNetV2
|
||||
# -----------------------------------------------------------------
|
||||
class InvertedResidual(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, stride, expand_ratio,
|
||||
dilation=1, norm_layer=nn.BatchNorm2d, **kwargs):
|
||||
super(InvertedResidual, self).__init__()
|
||||
assert stride in [1, 2]
|
||||
self.use_res_connect = stride == 1 and in_channels == out_channels
|
||||
|
||||
layers = list()
|
||||
inter_channels = int(round(in_channels * expand_ratio))
|
||||
if expand_ratio != 1:
|
||||
# pw
|
||||
layers.append(_ConvBNReLU(in_channels, inter_channels, 1, relu6=True, norm_layer=norm_layer))
|
||||
layers.extend([
|
||||
# dw
|
||||
_ConvBNReLU(inter_channels, inter_channels, 3, stride, dilation, dilation,
|
||||
groups=inter_channels, relu6=True, norm_layer=norm_layer),
|
||||
# pw-linear
|
||||
nn.Conv2d(inter_channels, out_channels, 1, bias=False),
|
||||
norm_layer(out_channels)])
|
||||
self.conv = nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
if self.use_res_connect:
|
||||
return x + self.conv(x)
|
||||
else:
|
||||
return self.conv(x)
|
||||
|
||||
|
||||
# -----------------------------------------------------------------
|
||||
# ASPP: For MobileNetV2
|
||||
# -----------------------------------------------------------------
|
||||
class _AsppPooling(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, norm_layer, **kwargs):
|
||||
super(_AsppPooling, self).__init__()
|
||||
self.gap = nn.Sequential(
|
||||
nn.AdaptiveAvgPool2d(1),
|
||||
nn.Conv2d(in_channels, out_channels, 1, bias=False),
|
||||
norm_layer(out_channels),
|
||||
nn.ReLU(True)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
# size = x.size()[2:]
|
||||
size = (48, 48)
|
||||
# print("size: ", size)
|
||||
pool = self.gap(x)
|
||||
# out = F.interpolate(pool, size, mode='bilinear', align_corners=True)
|
||||
out = F.interpolate(pool, size, mode='nearest')
|
||||
return out
|
||||
|
||||
|
||||
class _ASPP(nn.Module):
|
||||
def __init__(self, in_channels, atrous_rates, norm_layer=nn.BatchNorm2d, **kwargs):
|
||||
super(_ASPP, self).__init__()
|
||||
out_channels = 256
|
||||
self.b0 = nn.Sequential(
|
||||
nn.Conv2d(in_channels, out_channels, 1, bias=False),
|
||||
norm_layer(out_channels),
|
||||
nn.ReLU(True)
|
||||
)
|
||||
|
||||
rate1, rate2, rate3 = tuple(atrous_rates)
|
||||
self.b1 = _ConvBNReLU(in_channels, out_channels, 3, padding=rate1, dilation=rate1, norm_layer=norm_layer)
|
||||
self.b2 = _ConvBNReLU(in_channels, out_channels, 3, padding=rate2, dilation=rate2, norm_layer=norm_layer)
|
||||
self.b3 = _ConvBNReLU(in_channels, out_channels, 3, padding=rate3, dilation=rate3, norm_layer=norm_layer)
|
||||
self.b4 = _AsppPooling(in_channels, out_channels, norm_layer=norm_layer)
|
||||
|
||||
self.project = nn.Sequential(
|
||||
nn.Conv2d(5 * out_channels, out_channels, 1, bias=False),
|
||||
norm_layer(out_channels),
|
||||
nn.ReLU(True),
|
||||
nn.Dropout2d(0.5)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
feat1 = self.b0(x)
|
||||
feat2 = self.b1(x)
|
||||
feat3 = self.b2(x)
|
||||
feat4 = self.b3(x)
|
||||
feat5 = self.b4(x)
|
||||
x = torch.cat((feat1, feat2, feat3, feat4, feat5), dim=1)
|
||||
x = self.project(x)
|
||||
return x
|
||||
|
||||
|
||||
# -----------------------------------------------------------------
|
||||
# For MobileNetV3
|
||||
# -----------------------------------------------------------------
|
||||
class _Hswish(nn.Module):
|
||||
def __init__(self, inplace=True):
|
||||
super(_Hswish, self).__init__()
|
||||
self.relu6 = nn.ReLU6(inplace)
|
||||
|
||||
def forward(self, x):
|
||||
return x * self.relu6(x + 3.) / 6.
|
||||
|
||||
|
||||
class _Hsigmoid(nn.Module):
|
||||
def __init__(self, inplace=True):
|
||||
super(_Hsigmoid, self).__init__()
|
||||
self.relu6 = nn.ReLU6(inplace)
|
||||
|
||||
def forward(self, x):
|
||||
return self.relu6(x + 3.) / 6.
|
||||
|
||||
|
||||
class _ConvBNHswish(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0,
|
||||
dilation=1, groups=1, norm_layer=nn.BatchNorm2d, **kwargs):
|
||||
super(_ConvBNHswish, self).__init__()
|
||||
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias=False)
|
||||
self.bn = norm_layer(out_channels)
|
||||
self.act = _Hswish(True)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv(x)
|
||||
x = self.bn(x)
|
||||
x = self.act(x)
|
||||
return x
|
||||
|
||||
|
||||
class SEModule(nn.Module):
|
||||
def __init__(self, in_channels, reduction=4):
|
||||
super(SEModule, self).__init__()
|
||||
self.avg_pool = nn.AdaptiveAvgPool2d(1)
|
||||
self.fc = nn.Sequential(
|
||||
nn.Linear(in_channels, in_channels // reduction, bias=False),
|
||||
nn.ReLU(True),
|
||||
nn.Linear(in_channels // reduction, in_channels, bias=False),
|
||||
_Hsigmoid(True)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
n, c, _, _ = x.size()
|
||||
out = self.avg_pool(x).view(n, c)
|
||||
out = self.fc(out).view(n, c, 1, 1)
|
||||
return x * out.expand_as(x)
|
||||
|
||||
|
||||
class Identity(nn.Module):
|
||||
def __init__(self, in_channels):
|
||||
super(Identity, self).__init__()
|
||||
|
||||
def forward(self, x):
|
||||
return x
|
||||
|
||||
|
||||
class Bottleneck(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, exp_size, kernel_size, stride, dilation=1, se=False, nl='RE',
|
||||
norm_layer=nn.BatchNorm2d, **kwargs):
|
||||
super(Bottleneck, self).__init__()
|
||||
assert stride in [1, 2]
|
||||
self.use_res_connect = stride == 1 and in_channels == out_channels
|
||||
if nl == 'HS':
|
||||
act = _Hswish
|
||||
else:
|
||||
act = nn.ReLU
|
||||
if se:
|
||||
SELayer = SEModule
|
||||
else:
|
||||
SELayer = Identity
|
||||
|
||||
self.conv = nn.Sequential(
|
||||
# pw
|
||||
nn.Conv2d(in_channels, exp_size, 1, bias=False),
|
||||
norm_layer(exp_size),
|
||||
act(True),
|
||||
# dw
|
||||
nn.Conv2d(exp_size, exp_size, kernel_size, stride, (kernel_size - 1) // 2 * dilation,
|
||||
dilation, groups=exp_size, bias=False),
|
||||
norm_layer(exp_size),
|
||||
SELayer(exp_size),
|
||||
act(True),
|
||||
# pw-linear
|
||||
nn.Conv2d(exp_size, out_channels, 1, bias=False),
|
||||
norm_layer(out_channels)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
if self.use_res_connect:
|
||||
return x + self.conv(x)
|
||||
else:
|
||||
return self.conv(x)
|
||||
|
||||
|
||||
# -----------------------------------------------------------------
|
||||
# For ShuffleNet
|
||||
# -----------------------------------------------------------------
|
||||
def channel_shuffle(x, groups):
|
||||
n, c, h, w = x.size()
|
||||
|
||||
channels_per_group = c // groups
|
||||
x = x.view(n, groups, channels_per_group, h, w)
|
||||
x = torch.transpose(x, 1, 2).contiguous()
|
||||
x = x.view(n, -1, h, w)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class ShuffleNetUnit(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, stride, groups, dilation=1, norm_layer=nn.BatchNorm2d, **kwargs):
|
||||
super(ShuffleNetUnit, self).__init__()
|
||||
self.stride = stride
|
||||
self.groups = groups
|
||||
self.dilation = dilation
|
||||
assert stride in [1, 2, 3]
|
||||
|
||||
inter_channels = out_channels // 4
|
||||
|
||||
if stride > 1:
|
||||
self.shortcut = nn.AvgPool2d(3, stride, 1)
|
||||
out_channels -= in_channels
|
||||
elif dilation > 1:
|
||||
out_channels -= in_channels
|
||||
|
||||
g = 1 if in_channels == 24 else groups
|
||||
self.conv1 = _ConvBNReLU(in_channels, inter_channels, 1, groups=g, norm_layer=norm_layer)
|
||||
self.conv2 = _ConvBNReLU(inter_channels, inter_channels, 3, stride, dilation,
|
||||
dilation, groups, norm_layer=norm_layer)
|
||||
self.conv3 = nn.Sequential(
|
||||
nn.Conv2d(inter_channels, out_channels, 1, groups=groups, bias=False),
|
||||
norm_layer(out_channels))
|
||||
|
||||
def forward(self, x):
|
||||
out = self.conv1(x)
|
||||
out = channel_shuffle(out, self.groups)
|
||||
out = self.conv2(out)
|
||||
out = self.conv3(out)
|
||||
if self.stride > 1:
|
||||
x = self.shortcut(x)
|
||||
out = torch.cat([out, x], dim=1)
|
||||
elif self.dilation > 1:
|
||||
out = torch.cat([out, x], dim=1)
|
||||
else:
|
||||
out = out + x
|
||||
out = F.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
# -----------------------------------------------------------------
|
||||
# For ShuffleNetV2
|
||||
# -----------------------------------------------------------------
|
||||
class _DWConv(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, bias=False):
|
||||
super(_DWConv, self).__init__()
|
||||
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride,
|
||||
padding, dilation, groups=in_channels, bias=bias)
|
||||
|
||||
def forward(self, x):
|
||||
return self.conv(x)
|
||||
|
||||
|
||||
class ShuffleNetV2Unit(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, stride, dilation=1, norm_layer=nn.BatchNorm2d, **kwargs):
|
||||
super(ShuffleNetV2Unit, self).__init__()
|
||||
assert stride in [1, 2, 3]
|
||||
self.stride = stride
|
||||
self.dilation = dilation
|
||||
|
||||
inter_channels = out_channels // 2
|
||||
|
||||
if (stride > 1) or (dilation > 1):
|
||||
self.branch1 = nn.Sequential(
|
||||
_DWConv(in_channels, in_channels, 3, stride, dilation, dilation),
|
||||
norm_layer(in_channels),
|
||||
_ConvBNReLU(in_channels, inter_channels, 1, norm_layer=norm_layer))
|
||||
self.branch2 = nn.Sequential(
|
||||
_ConvBNReLU(in_channels if (stride > 1) else inter_channels, inter_channels, 1, norm_layer=norm_layer),
|
||||
_DWConv(inter_channels, inter_channels, 3, stride, dilation, dilation),
|
||||
norm_layer(inter_channels),
|
||||
_ConvBNReLU(inter_channels, inter_channels, 1, norm_layer=norm_layer))
|
||||
|
||||
def forward(self, x):
|
||||
if (self.stride == 1) and (self.dilation == 1):
|
||||
x1, x2 = x.chunk(2, dim=1)
|
||||
out = torch.cat((x1, self.branch2(x2)), dim=1)
|
||||
else:
|
||||
out = torch.cat((self.branch1(x), self.branch2(x)), dim=1)
|
||||
out = channel_shuffle(out, 2)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
# -----------------------------------------------------------------
|
||||
# For IGCV3
|
||||
# -----------------------------------------------------------------
|
||||
class PermutationBlock(nn.Module):
|
||||
def __init__(self, groups):
|
||||
super(PermutationBlock, self).__init__()
|
||||
self.groups = groups
|
||||
|
||||
def forward(self, x):
|
||||
n, c, h, w = x.size()
|
||||
x = x.view(n, self.groups, c // self.groups, h, w).permute(0, 2, 1, 3, 4).contiguous().view(n, c, h, w)
|
||||
return x
|
||||
|
||||
|
||||
class InvertedIGCV3(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, stride, expand_ratio,
|
||||
dilation=1, norm_layer=nn.BatchNorm2d, **kwargs):
|
||||
super(InvertedIGCV3, self).__init__()
|
||||
assert stride in [1, 2]
|
||||
self.use_res_connect = stride == 1 and in_channels == out_channels
|
||||
|
||||
layers = list()
|
||||
inter_channels = int(round(in_channels * expand_ratio))
|
||||
if expand_ratio != 1:
|
||||
# pw
|
||||
layers.append(_ConvBNReLU(in_channels, inter_channels, 1,
|
||||
groups=2, relu6=True, norm_layer=norm_layer))
|
||||
# permutation
|
||||
layers.append(PermutationBlock(groups=2))
|
||||
layers.extend([
|
||||
# dw
|
||||
_ConvBNReLU(inter_channels, inter_channels, 3, stride, dilation, dilation,
|
||||
groups=inter_channels, relu6=True, norm_layer=norm_layer),
|
||||
# pw-linear
|
||||
nn.Conv2d(inter_channels, out_channels, 1, groups=2, bias=False),
|
||||
norm_layer(out_channels),
|
||||
# permutation
|
||||
PermutationBlock(groups=int(round(out_channels / 2)))
|
||||
])
|
||||
self.conv = nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
if self.use_res_connect:
|
||||
return x + self.conv(x)
|
||||
else:
|
||||
return self.conv(x)
|
||||
|
||||
|
||||
# -----------------------------------------------------------------
|
||||
# For EfficientNet
|
||||
# -----------------------------------------------------------------
|
||||
class _Swish(nn.Module):
|
||||
def __init__(self):
|
||||
super(_Swish, self).__init__()
|
||||
self.sigmoid = nn.Sigmoid()
|
||||
|
||||
def forward(self, x):
|
||||
return x * self.sigmoid(x)
|
||||
|
||||
|
||||
class SEModuleV2(nn.Module):
|
||||
def __init__(self, in_channels, se_ratio=0.25):
|
||||
super(SEModuleV2, self).__init__()
|
||||
self.avg_pool = nn.AdaptiveAvgPool2d(1)
|
||||
se_channels = max(1, int(in_channels * se_ratio))
|
||||
self.fc = nn.Sequential(
|
||||
nn.Conv2d(in_channels, se_channels, 1, bias=False),
|
||||
_Swish(),
|
||||
nn.Conv2d(se_channels, in_channels, 1, bias=False),
|
||||
nn.Sigmoid()
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
n, c, _, _ = x.size()
|
||||
out = self.avg_pool(x)
|
||||
out = self.fc(out)
|
||||
return x * out.expand_as(x)
|
||||
|
||||
|
||||
class MBConvBlock(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, kernel_size, stride, expand_ratio,
|
||||
dilation=1, se_ratio=0.25, drop_connect_rate=0.2, norm_layer=nn.BatchNorm2d, **kwargs):
|
||||
super(MBConvBlock, self).__init__()
|
||||
assert stride in [1, 2]
|
||||
self.use_res_connect = stride == 1 and in_channels == out_channels
|
||||
self.drop_connect_rate = drop_connect_rate
|
||||
use_se = (se_ratio is not None) and (0 < se_ratio <= 1.)
|
||||
if use_se:
|
||||
SELayer = SEModuleV2
|
||||
else:
|
||||
SELayer = Identity
|
||||
|
||||
layers = list()
|
||||
inter_channels = int(round(in_channels * expand_ratio))
|
||||
if expand_ratio != 1:
|
||||
layers.append(_ConvBNHswish(in_channels, inter_channels, 1, norm_layer=norm_layer))
|
||||
layers.extend([
|
||||
# dw
|
||||
_ConvBNHswish(inter_channels, inter_channels, kernel_size, stride, kernel_size // 2 * dilation, dilation,
|
||||
groups=inter_channels, norm_layer=norm_layer), # check act function
|
||||
SELayer(inter_channels, se_ratio),
|
||||
# pw-linear
|
||||
nn.Conv2d(inter_channels, out_channels, 1, bias=False),
|
||||
norm_layer(out_channels)
|
||||
])
|
||||
self.conv = nn.Sequential(*layers)
|
||||
|
||||
if drop_connect_rate:
|
||||
self.dropout = nn.Dropout2d(drop_connect_rate)
|
||||
|
||||
def forward(self, x):
|
||||
out = self.conv(x)
|
||||
if self.use_res_connect:
|
||||
if self.drop_connect_rate:
|
||||
out = self.dropout(out)
|
||||
out = x + out
|
||||
return out
|
||||
@@ -0,0 +1,187 @@
|
||||
"""Pyramid Scene Parsing Network"""
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from seg.networks.segbase import SegBaseModel
|
||||
from seg.networks.fcn import _FCNHead
|
||||
|
||||
__all__ = ['DeepLabV3', 'get_deeplabv3', 'get_deeplabv3_resnet50_voc', 'get_deeplabv3_resnet101_voc',
|
||||
'get_deeplabv3_resnet152_voc', 'get_deeplabv3_resnet50_ade', 'get_deeplabv3_resnet101_ade',
|
||||
'get_deeplabv3_resnet152_ade']
|
||||
|
||||
|
||||
class DeepLabV3(SegBaseModel):
|
||||
r"""DeepLabV3
|
||||
|
||||
Parameters
|
||||
----------
|
||||
nclass : int
|
||||
Number of categories for the training dataset.
|
||||
backbone : string
|
||||
Pre-trained dilated backbone network type (default:'resnet50'; 'resnet50',
|
||||
'resnet101' or 'resnet152').
|
||||
norm_layer : object
|
||||
Normalization layer used in backbone network (default: :class:`nn.BatchNorm`;
|
||||
for Synchronized Cross-GPU BachNormalization).
|
||||
aux : bool
|
||||
Auxiliary loss.
|
||||
|
||||
Reference:
|
||||
Chen, Liang-Chieh, et al. "Rethinking atrous convolution for semantic image segmentation."
|
||||
arXiv preprint arXiv:1706.05587 (2017).
|
||||
"""
|
||||
|
||||
def __init__(self, nclass, backbone='resnet50', aux=False, pretrained_base=True, **kwargs):
|
||||
super(DeepLabV3, self).__init__(nclass, aux, backbone, pretrained_base=pretrained_base, **kwargs)
|
||||
self.head = _DeepLabHead(nclass, **kwargs)
|
||||
if self.aux:
|
||||
self.auxlayer = _FCNHead(1024, nclass, **kwargs)
|
||||
|
||||
self.__setattr__('exclusive', ['head', 'auxlayer'] if aux else ['head'])
|
||||
|
||||
def forward(self, x):
|
||||
size = x.size()[2:]
|
||||
_, _, c3, c4 = self.base_forward(x)
|
||||
outputs = []
|
||||
x = self.head(c4)
|
||||
x = F.interpolate(x, size, mode='bilinear', align_corners=True)
|
||||
outputs.append(x)
|
||||
|
||||
if self.aux:
|
||||
auxout = self.auxlayer(c3)
|
||||
auxout = F.interpolate(auxout, size, mode='bilinear', align_corners=True)
|
||||
outputs.append(auxout)
|
||||
return tuple(outputs)
|
||||
|
||||
|
||||
class _DeepLabHead(nn.Module):
|
||||
def __init__(self, nclass, norm_layer=nn.BatchNorm2d, norm_kwargs=None, **kwargs):
|
||||
super(_DeepLabHead, self).__init__()
|
||||
self.aspp = _ASPP(2048, [12, 24, 36], norm_layer=norm_layer, norm_kwargs=norm_kwargs, **kwargs)
|
||||
self.block = nn.Sequential(
|
||||
nn.Conv2d(256, 256, 3, padding=1, bias=False),
|
||||
norm_layer(256, **({} if norm_kwargs is None else norm_kwargs)),
|
||||
nn.ReLU(True),
|
||||
nn.Dropout(0.1),
|
||||
nn.Conv2d(256, nclass, 1)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.aspp(x)
|
||||
return self.block(x)
|
||||
|
||||
|
||||
class _ASPPConv(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, atrous_rate, norm_layer, norm_kwargs):
|
||||
super(_ASPPConv, self).__init__()
|
||||
self.block = nn.Sequential(
|
||||
nn.Conv2d(in_channels, out_channels, 3, padding=atrous_rate, dilation=atrous_rate, bias=False),
|
||||
norm_layer(out_channels, **({} if norm_kwargs is None else norm_kwargs)),
|
||||
nn.ReLU(True)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return self.block(x)
|
||||
|
||||
|
||||
class _AsppPooling(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, norm_layer, norm_kwargs, **kwargs):
|
||||
super(_AsppPooling, self).__init__()
|
||||
self.gap = nn.Sequential(
|
||||
nn.AdaptiveAvgPool2d(1),
|
||||
nn.Conv2d(in_channels, out_channels, 1, bias=False),
|
||||
norm_layer(out_channels, **({} if norm_kwargs is None else norm_kwargs)),
|
||||
nn.ReLU(True)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
size = x.size()[2:]
|
||||
# print("before gap: ", x.size())
|
||||
pool = self.gap(x)
|
||||
out = F.interpolate(pool, size, mode='bilinear', align_corners=True)
|
||||
return out
|
||||
|
||||
|
||||
class _ASPP(nn.Module):
|
||||
def __init__(self, in_channels, atrous_rates, norm_layer, norm_kwargs=None, **kwargs):
|
||||
super(_ASPP, self).__init__()
|
||||
out_channels = 256
|
||||
self.b0 = nn.Sequential(
|
||||
nn.Conv2d(in_channels, out_channels, 1, bias=False),
|
||||
norm_layer(out_channels, **({} if norm_kwargs is None else norm_kwargs)),
|
||||
nn.ReLU(True)
|
||||
)
|
||||
|
||||
rate1, rate2, rate3 = tuple(atrous_rates)
|
||||
self.b1 = _ASPPConv(in_channels, out_channels, rate1, norm_layer, norm_kwargs)
|
||||
self.b2 = _ASPPConv(in_channels, out_channels, rate2, norm_layer, norm_kwargs)
|
||||
self.b3 = _ASPPConv(in_channels, out_channels, rate3, norm_layer, norm_kwargs)
|
||||
self.b4 = _AsppPooling(in_channels, out_channels, norm_layer=norm_layer, norm_kwargs=norm_kwargs)
|
||||
|
||||
self.project = nn.Sequential(
|
||||
nn.Conv2d(5 * out_channels, out_channels, 1, bias=False),
|
||||
norm_layer(out_channels, **({} if norm_kwargs is None else norm_kwargs)),
|
||||
nn.ReLU(True),
|
||||
nn.Dropout(0.5)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
feat1 = self.b0(x)
|
||||
feat2 = self.b1(x)
|
||||
feat3 = self.b2(x)
|
||||
feat4 = self.b3(x)
|
||||
# print("before b4: ", x.size())
|
||||
feat5 = self.b4(x)
|
||||
x = torch.cat((feat1, feat2, feat3, feat4, feat5), dim=1)
|
||||
x = self.project(x)
|
||||
return x
|
||||
|
||||
|
||||
def get_deeplabv3(dataset='pascal_voc', backbone='resnet50', pretrained=False, root='~/.torch/models',
|
||||
pretrained_base=True, **kwargs):
|
||||
acronyms = {
|
||||
'pascal_voc': 'pascal_voc',
|
||||
'pascal_aug': 'pascal_aug',
|
||||
'ade20k': 'ade',
|
||||
'coco': 'coco',
|
||||
'citys': 'citys',
|
||||
}
|
||||
from ..data.dataloader import datasets
|
||||
model = DeepLabV3(datasets[dataset].NUM_CLASS, backbone=backbone, pretrained_base=pretrained_base, **kwargs)
|
||||
if pretrained:
|
||||
from .model_store import get_model_file
|
||||
device = torch.device(kwargs['local_rank'])
|
||||
model.load_state_dict(torch.load(get_model_file('deeplabv3_%s_%s' % (backbone, acronyms[dataset]), root=root),
|
||||
map_location=device))
|
||||
return model
|
||||
|
||||
|
||||
def get_deeplabv3_resnet50_voc(**kwargs):
|
||||
return get_deeplabv3('pascal_voc', 'resnet50', **kwargs)
|
||||
|
||||
|
||||
def get_deeplabv3_resnet101_voc(**kwargs):
|
||||
return get_deeplabv3('pascal_voc', 'resnet101', **kwargs)
|
||||
|
||||
|
||||
def get_deeplabv3_resnet152_voc(**kwargs):
|
||||
return get_deeplabv3('pascal_voc', 'resnet152', **kwargs)
|
||||
|
||||
|
||||
def get_deeplabv3_resnet50_ade(**kwargs):
|
||||
return get_deeplabv3('ade20k', 'resnet50', **kwargs)
|
||||
|
||||
|
||||
def get_deeplabv3_resnet101_ade(**kwargs):
|
||||
return get_deeplabv3('ade20k', 'resnet101', **kwargs)
|
||||
|
||||
|
||||
def get_deeplabv3_resnet152_ade(**kwargs):
|
||||
return get_deeplabv3('ade20k', 'resnet152', **kwargs)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
model = get_deeplabv3_resnet50_voc()
|
||||
img = torch.randn(2, 3, 480, 480)
|
||||
output = model(img)
|
||||
@@ -0,0 +1,160 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from seg.networks.xception import get_xception
|
||||
from seg.networks.deeplabv3 import _ASPP
|
||||
from seg.networks.fcn import _FCNHead
|
||||
from seg.networks.basic import _ConvBNReLU
|
||||
|
||||
__all__ = ['DeepLabV3Plus', 'get_deeplabv3_plus', 'get_deeplabv3_plus_xception_voc']
|
||||
|
||||
|
||||
class DeepLabV3Plus(nn.Module):
|
||||
r"""DeepLabV3Plus
|
||||
Parameters
|
||||
----------
|
||||
nclass : int
|
||||
Number of categories for the training dataset.
|
||||
backbone : string
|
||||
Pre-trained dilated backbone network type (default:'xception').
|
||||
norm_layer : object
|
||||
Normalization layer used in backbone network (default: :class:`nn.BatchNorm`;
|
||||
for Synchronized Cross-GPU BachNormalization).
|
||||
aux : bool
|
||||
Auxiliary loss.
|
||||
|
||||
Reference:
|
||||
Chen, Liang-Chieh, et al. "Encoder-Decoder with Atrous Separable Convolution for Semantic
|
||||
Image Segmentation."
|
||||
"""
|
||||
|
||||
def __init__(self, nclass, backbone='xception', aux=True, pretrained_base=True, dilated=True, **kwargs):
|
||||
super(DeepLabV3Plus, self).__init__()
|
||||
self.aux = aux
|
||||
self.nclass = nclass
|
||||
output_stride = 8 if dilated else 32
|
||||
|
||||
self.pretrained = get_xception(pretrained=pretrained_base, output_stride=output_stride, **kwargs)
|
||||
|
||||
# deeplabv3 plus
|
||||
self.head = _DeepLabHead(nclass, **kwargs)
|
||||
if aux:
|
||||
self.auxlayer = _FCNHead(728, nclass, **kwargs)
|
||||
|
||||
def base_forward(self, x):
|
||||
# Entry flow
|
||||
x = self.pretrained.conv1(x)
|
||||
x = self.pretrained.bn1(x)
|
||||
x = self.pretrained.relu(x)
|
||||
|
||||
x = self.pretrained.conv2(x)
|
||||
x = self.pretrained.bn2(x)
|
||||
x = self.pretrained.relu(x)
|
||||
|
||||
x = self.pretrained.block1(x)
|
||||
# add relu here
|
||||
x = self.pretrained.relu(x)
|
||||
low_level_feat = x
|
||||
|
||||
x = self.pretrained.block2(x)
|
||||
x = self.pretrained.block3(x)
|
||||
|
||||
# Middle flow
|
||||
x = self.pretrained.midflow(x)
|
||||
mid_level_feat = x
|
||||
|
||||
# Exit flow
|
||||
x = self.pretrained.block20(x)
|
||||
x = self.pretrained.relu(x)
|
||||
x = self.pretrained.conv3(x)
|
||||
x = self.pretrained.bn3(x)
|
||||
x = self.pretrained.relu(x)
|
||||
|
||||
x = self.pretrained.conv4(x)
|
||||
x = self.pretrained.bn4(x)
|
||||
x = self.pretrained.relu(x)
|
||||
|
||||
x = self.pretrained.conv5(x)
|
||||
x = self.pretrained.bn5(x)
|
||||
x = self.pretrained.relu(x)
|
||||
return low_level_feat, mid_level_feat, x
|
||||
|
||||
def forward(self, x):
|
||||
# print("x size: ", x.size())
|
||||
size = x.size()[2:]
|
||||
c1, c3, c4 = self.base_forward(x)
|
||||
# print("c1 size: ", c1.size())
|
||||
# print("c4 size: ", c4.size())
|
||||
outputs = list()
|
||||
x = self.head(c4, c1)
|
||||
x = F.interpolate(x, size, mode='bilinear', align_corners=True)
|
||||
outputs.append(x)
|
||||
if self.aux:
|
||||
auxout = self.auxlayer(c3)
|
||||
auxout = F.interpolate(auxout, size, mode='bilinear', align_corners=True)
|
||||
outputs.append(auxout)
|
||||
|
||||
# for save onnx
|
||||
# y = torch.max(x, 1)[1].to(torch.float32)
|
||||
# return y
|
||||
|
||||
return tuple(outputs)
|
||||
|
||||
|
||||
class _DeepLabHead(nn.Module):
|
||||
def __init__(self, nclass, c1_channels=128, norm_layer=nn.BatchNorm2d, **kwargs):
|
||||
super(_DeepLabHead, self).__init__()
|
||||
self.aspp = _ASPP(2048, [12, 24, 36], norm_layer=norm_layer, **kwargs)
|
||||
self.c1_block = _ConvBNReLU(c1_channels, 48, 3, padding=1, norm_layer=norm_layer)
|
||||
self.block = nn.Sequential(
|
||||
_ConvBNReLU(304, 256, 3, padding=1, norm_layer=norm_layer),
|
||||
nn.Dropout(0.5),
|
||||
_ConvBNReLU(256, 256, 3, padding=1, norm_layer=norm_layer),
|
||||
nn.Dropout(0.1),
|
||||
nn.Conv2d(256, nclass, 1))
|
||||
|
||||
def forward(self, x, c1):
|
||||
size = c1.size()[2:]
|
||||
c1 = self.c1_block(c1)
|
||||
# print("c1", c1.size())
|
||||
# print("before aspp: ", x.size())
|
||||
x = self.aspp(x)
|
||||
# print("after aspp: ", x.size())
|
||||
x = F.interpolate(x, size, mode='bilinear', align_corners=True)
|
||||
return self.block(torch.cat([x, c1], dim=1))
|
||||
|
||||
|
||||
def get_deeplabv3_plus(dataset='pascal_voc', backbone='xception', pretrained=False, root='../ckpt',
|
||||
pretrained_base=False, nclass=3, **kwargs):
|
||||
acronyms = {
|
||||
'pascal_voc': 'pascal_voc',
|
||||
'pascal_aug': 'pascal_aug',
|
||||
'ade20k': 'ade',
|
||||
'coco': 'coco',
|
||||
'citys': 'citys',
|
||||
}
|
||||
#from light.data import datasets
|
||||
|
||||
model = DeepLabV3Plus(nclass, backbone=backbone, pretrained_base=pretrained_base, **kwargs)
|
||||
if pretrained:
|
||||
pass
|
||||
# if dataset not in acronyms.keys():
|
||||
# print("root:", root)
|
||||
# model_path = os.path.join(root, "deeplabv3_plus_28.pth")
|
||||
# model.load_state_dict(torch.load(model_path), strict=False)
|
||||
# else:
|
||||
# from .model_store import get_model_file
|
||||
# device = torch.device(kwargs['local_rank'])
|
||||
# model.load_state_dict(
|
||||
# torch.load(get_model_file('deeplabv3_plus_%s_%s' % (backbone, acronyms[dataset]), root=root),
|
||||
# map_location=device))
|
||||
return model
|
||||
|
||||
|
||||
def get_deeplabv3_plus_xception_voc(**kwargs):
|
||||
return get_deeplabv3_plus('pascal_voc', 'xception', **kwargs)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
model = get_deeplabv3_plus_xception_voc()
|
||||
@@ -0,0 +1,222 @@
|
||||
import os
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from seg.networks.vgg import vgg16
|
||||
|
||||
__all__ = ['get_fcn32s', 'get_fcn16s', 'get_fcn8s',
|
||||
'get_fcn32s_vgg16_voc', 'get_fcn16s_vgg16_voc', 'get_fcn8s_vgg16_voc']
|
||||
|
||||
|
||||
class FCN32s(nn.Module):
|
||||
"""There are some difference from original fcn"""
|
||||
|
||||
def __init__(self, nclass, backbone='vgg16', aux=False, pretrained_base=True,
|
||||
norm_layer=nn.BatchNorm2d, **kwargs):
|
||||
super(FCN32s, self).__init__()
|
||||
self.aux = aux
|
||||
if backbone == 'vgg16':
|
||||
self.pretrained = vgg16(pretrained=pretrained_base).features
|
||||
else:
|
||||
raise RuntimeError('unknown backbone: {}'.format(backbone))
|
||||
self.head = _FCNHead(512, nclass, norm_layer)
|
||||
if aux:
|
||||
self.auxlayer = _FCNHead(512, nclass, norm_layer)
|
||||
|
||||
self.__setattr__('exclusive', ['head', 'auxlayer'] if aux else ['head'])
|
||||
|
||||
def forward(self, x):
|
||||
size = x.size()[2:]
|
||||
pool5 = self.pretrained(x)
|
||||
|
||||
outputs = []
|
||||
out = self.head(pool5)
|
||||
out = F.interpolate(out, size, mode='bilinear', align_corners=True)
|
||||
outputs.append(out)
|
||||
|
||||
if self.aux:
|
||||
auxout = self.auxlayer(pool5)
|
||||
auxout = F.interpolate(auxout, size, mode='bilinear', align_corners=True)
|
||||
outputs.append(auxout)
|
||||
|
||||
return tuple(outputs)
|
||||
|
||||
|
||||
class FCN16s(nn.Module):
|
||||
def __init__(self, nclass, backbone='vgg16', aux=False, pretrained_base=True, norm_layer=nn.BatchNorm2d, **kwargs):
|
||||
super(FCN16s, self).__init__()
|
||||
self.aux = aux
|
||||
if backbone == 'vgg16':
|
||||
self.pretrained = vgg16(pretrained=pretrained_base).features
|
||||
else:
|
||||
raise RuntimeError('unknown backbone: {}'.format(backbone))
|
||||
self.pool4 = nn.Sequential(*self.pretrained[:24])
|
||||
self.pool5 = nn.Sequential(*self.pretrained[24:])
|
||||
self.head = _FCNHead(512, nclass, norm_layer)
|
||||
self.score_pool4 = nn.Conv2d(512, nclass, 1)
|
||||
if aux:
|
||||
self.auxlayer = _FCNHead(512, nclass, norm_layer)
|
||||
|
||||
self.__setattr__('exclusive', ['head', 'score_pool4', 'auxlayer'] if aux else ['head', 'score_pool4'])
|
||||
|
||||
def forward(self, x):
|
||||
pool4 = self.pool4(x)
|
||||
pool5 = self.pool5(pool4)
|
||||
|
||||
outputs = []
|
||||
score_fr = self.head(pool5)
|
||||
|
||||
score_pool4 = self.score_pool4(pool4)
|
||||
|
||||
upscore2 = F.interpolate(score_fr, score_pool4.size()[2:], mode='bilinear', align_corners=True)
|
||||
fuse_pool4 = upscore2 + score_pool4
|
||||
|
||||
out = F.interpolate(fuse_pool4, x.size()[2:], mode='bilinear', align_corners=True)
|
||||
outputs.append(out)
|
||||
|
||||
if self.aux:
|
||||
auxout = self.auxlayer(pool5)
|
||||
auxout = F.interpolate(auxout, x.size()[2:], mode='bilinear', align_corners=True)
|
||||
outputs.append(auxout)
|
||||
|
||||
return tuple(outputs)
|
||||
|
||||
|
||||
class FCN8s(nn.Module):
|
||||
def __init__(self, nclass, backbone='vgg16', aux=False, pretrained_base=True, norm_layer=nn.BatchNorm2d, **kwargs):
|
||||
super(FCN8s, self).__init__()
|
||||
self.aux = aux
|
||||
if backbone == 'vgg16':
|
||||
self.pretrained = vgg16(pretrained=pretrained_base).features
|
||||
else:
|
||||
raise RuntimeError('unknown backbone: {}'.format(backbone))
|
||||
self.pool3 = nn.Sequential(*self.pretrained[:17])
|
||||
self.pool4 = nn.Sequential(*self.pretrained[17:24])
|
||||
self.pool5 = nn.Sequential(*self.pretrained[24:])
|
||||
self.head = _FCNHead(512, nclass, norm_layer)
|
||||
self.score_pool3 = nn.Conv2d(256, nclass, 1)
|
||||
self.score_pool4 = nn.Conv2d(512, nclass, 1)
|
||||
if aux:
|
||||
self.auxlayer = _FCNHead(512, nclass, norm_layer)
|
||||
|
||||
self.__setattr__('exclusive',
|
||||
['head', 'score_pool3', 'score_pool4', 'auxlayer'] if aux else ['head', 'score_pool3',
|
||||
'score_pool4'])
|
||||
|
||||
def forward(self, x):
|
||||
pool3 = self.pool3(x)
|
||||
pool4 = self.pool4(pool3)
|
||||
pool5 = self.pool5(pool4)
|
||||
|
||||
outputs = []
|
||||
score_fr = self.head(pool5)
|
||||
|
||||
score_pool4 = self.score_pool4(pool4)
|
||||
score_pool3 = self.score_pool3(pool3)
|
||||
|
||||
upscore2 = F.interpolate(score_fr, score_pool4.size()[2:], mode='bilinear', align_corners=True)
|
||||
fuse_pool4 = upscore2 + score_pool4
|
||||
|
||||
upscore_pool4 = F.interpolate(fuse_pool4, score_pool3.size()[2:], mode='bilinear', align_corners=True)
|
||||
fuse_pool3 = upscore_pool4 + score_pool3
|
||||
|
||||
out = F.interpolate(fuse_pool3, x.size()[2:], mode='bilinear', align_corners=True)
|
||||
outputs.append(out)
|
||||
|
||||
if self.aux:
|
||||
auxout = self.auxlayer(pool5)
|
||||
auxout = F.interpolate(auxout, x.size()[2:], mode='bilinear', align_corners=True)
|
||||
outputs.append(auxout)
|
||||
|
||||
return tuple(outputs)
|
||||
|
||||
|
||||
class _FCNHead(nn.Module):
|
||||
def __init__(self, in_channels, channels, norm_layer=nn.BatchNorm2d, **kwargs):
|
||||
super(_FCNHead, self).__init__()
|
||||
inter_channels = in_channels // 4
|
||||
self.block = nn.Sequential(
|
||||
nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False),
|
||||
norm_layer(inter_channels),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Dropout(0.1),
|
||||
nn.Conv2d(inter_channels, channels, 1)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return self.block(x)
|
||||
|
||||
|
||||
def get_fcn32s(dataset='pascal_voc', backbone='vgg16', pretrained=False, root='~/.torch/models',
|
||||
pretrained_base=True, **kwargs):
|
||||
acronyms = {
|
||||
'pascal_voc': 'pascal_voc',
|
||||
'pascal_aug': 'pascal_aug',
|
||||
'ade20k': 'ade',
|
||||
'coco': 'coco',
|
||||
'citys': 'citys',
|
||||
}
|
||||
from ..data.dataloader import datasets
|
||||
model = FCN32s(datasets[dataset].NUM_CLASS, backbone=backbone, pretrained_base=pretrained_base, **kwargs)
|
||||
if pretrained:
|
||||
from .model_store import get_model_file
|
||||
device = torch.device(kwargs['local_rank'])
|
||||
model.load_state_dict(torch.load(get_model_file('fcn32s_%s_%s' % (backbone, acronyms[dataset]), root=root),
|
||||
map_location=device))
|
||||
return model
|
||||
|
||||
|
||||
def get_fcn16s(dataset='pascal_voc', backbone='vgg16', pretrained=False, root='~/.torch/models',
|
||||
pretrained_base=True, **kwargs):
|
||||
acronyms = {
|
||||
'pascal_voc': 'pascal_voc',
|
||||
'pascal_aug': 'pascal_aug',
|
||||
'ade20k': 'ade',
|
||||
'coco': 'coco',
|
||||
'citys': 'citys',
|
||||
}
|
||||
from ..data.dataloader import datasets
|
||||
model = FCN16s(datasets[dataset].NUM_CLASS, backbone=backbone, pretrained_base=pretrained_base, **kwargs)
|
||||
if pretrained:
|
||||
from .model_store import get_model_file
|
||||
device = torch.device(kwargs['local_rank'])
|
||||
model.load_state_dict(torch.load(get_model_file('fcn16s_%s_%s' % (backbone, acronyms[dataset]), root=root),
|
||||
map_location=device))
|
||||
return model
|
||||
|
||||
|
||||
def get_fcn8s(dataset='pascal_voc', backbone='vgg16', pretrained=False, root='~/.torch/models',
|
||||
pretrained_base=True, **kwargs):
|
||||
acronyms = {
|
||||
'pascal_voc': 'pascal_voc',
|
||||
'pascal_aug': 'pascal_aug',
|
||||
'ade20k': 'ade',
|
||||
'coco': 'coco',
|
||||
'citys': 'citys',
|
||||
}
|
||||
from ..data.dataloader import datasets
|
||||
model = FCN8s(datasets[dataset].NUM_CLASS, backbone=backbone, pretrained_base=pretrained_base, **kwargs)
|
||||
if pretrained:
|
||||
from .model_store import get_model_file
|
||||
device = torch.device(kwargs['local_rank'])
|
||||
model.load_state_dict(torch.load(get_model_file('fcn8s_%s_%s' % (backbone, acronyms[dataset]), root=root),
|
||||
map_location=device))
|
||||
return model
|
||||
|
||||
|
||||
def get_fcn32s_vgg16_voc(**kwargs):
|
||||
return get_fcn32s('pascal_voc', 'vgg16', **kwargs)
|
||||
|
||||
|
||||
def get_fcn16s_vgg16_voc(**kwargs):
|
||||
return get_fcn16s('pascal_voc', 'vgg16', **kwargs)
|
||||
|
||||
|
||||
def get_fcn8s_vgg16_voc(**kwargs):
|
||||
return get_fcn8s('pascal_voc', 'vgg16', **kwargs)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
model = FCN16s(21)
|
||||
print(model)
|
||||
@@ -0,0 +1,68 @@
|
||||
"""Joint Pyramid Upsampling"""
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
__all__ = ['JPU']
|
||||
|
||||
|
||||
class SeparableConv2d(nn.Module):
|
||||
def __init__(self, inplanes, planes, kernel_size=3, stride=1, padding=1,
|
||||
dilation=1, bias=False, norm_layer=nn.BatchNorm2d):
|
||||
super(SeparableConv2d, self).__init__()
|
||||
self.conv = nn.Conv2d(inplanes, inplanes, kernel_size, stride, padding, dilation, groups=inplanes, bias=bias)
|
||||
self.bn = norm_layer(inplanes)
|
||||
self.pointwise = nn.Conv2d(inplanes, planes, 1, bias=bias)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv(x)
|
||||
x = self.bn(x)
|
||||
x = self.pointwise(x)
|
||||
return x
|
||||
|
||||
|
||||
# copy from: https://github.com/wuhuikai/FastFCN/blob/master/encoding/nn/customize.py
|
||||
class JPU(nn.Module):
|
||||
def __init__(self, in_channels, width=512, norm_layer=nn.BatchNorm2d, **kwargs):
|
||||
super(JPU, self).__init__()
|
||||
|
||||
self.conv5 = nn.Sequential(
|
||||
nn.Conv2d(in_channels[-1], width, 3, padding=1, bias=False),
|
||||
norm_layer(width),
|
||||
nn.ReLU(True))
|
||||
self.conv4 = nn.Sequential(
|
||||
nn.Conv2d(in_channels[-2], width, 3, padding=1, bias=False),
|
||||
norm_layer(width),
|
||||
nn.ReLU(True))
|
||||
self.conv3 = nn.Sequential(
|
||||
nn.Conv2d(in_channels[-3], width, 3, padding=1, bias=False),
|
||||
norm_layer(width),
|
||||
nn.ReLU(True))
|
||||
|
||||
self.dilation1 = nn.Sequential(
|
||||
SeparableConv2d(3 * width, width, 3, padding=1, dilation=1, bias=False),
|
||||
norm_layer(width),
|
||||
nn.ReLU(True))
|
||||
self.dilation2 = nn.Sequential(
|
||||
SeparableConv2d(3 * width, width, 3, padding=2, dilation=2, bias=False),
|
||||
norm_layer(width),
|
||||
nn.ReLU(True))
|
||||
self.dilation3 = nn.Sequential(
|
||||
SeparableConv2d(3 * width, width, 3, padding=4, dilation=4, bias=False),
|
||||
norm_layer(width),
|
||||
nn.ReLU(True))
|
||||
self.dilation4 = nn.Sequential(
|
||||
SeparableConv2d(3 * width, width, 3, padding=8, dilation=8, bias=False),
|
||||
norm_layer(width),
|
||||
nn.ReLU(True))
|
||||
|
||||
def forward(self, *inputs):
|
||||
feats = [self.conv5(inputs[-1]), self.conv4(inputs[-2]), self.conv3(inputs[-3])]
|
||||
size = feats[-1].size()[2:]
|
||||
feats[-2] = F.interpolate(feats[-2], size, mode='bilinear', align_corners=True)
|
||||
feats[-3] = F.interpolate(feats[-3], size, mode='bilinear', align_corners=True)
|
||||
feat = torch.cat(feats, dim=1)
|
||||
feat = torch.cat([self.dilation1(feat), self.dilation2(feat), self.dilation3(feat), self.dilation4(feat)],
|
||||
dim=1)
|
||||
|
||||
return inputs[0], inputs[1], inputs[2], feat
|
||||
@@ -0,0 +1,264 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.utils.model_zoo as model_zoo
|
||||
|
||||
__all__ = ['ResNetV1b', 'resnet18_v1b', 'resnet34_v1b', 'resnet50_v1b',
|
||||
'resnet101_v1b', 'resnet152_v1b', 'resnet152_v1s', 'resnet101_v1s', 'resnet50_v1s']
|
||||
|
||||
model_urls = {
|
||||
'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
|
||||
'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
|
||||
'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
|
||||
'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
|
||||
'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
|
||||
}
|
||||
|
||||
|
||||
class BasicBlockV1b(nn.Module):
|
||||
expansion = 1
|
||||
|
||||
def __init__(self, inplanes, planes, stride=1, dilation=1, downsample=None,
|
||||
previous_dilation=1, norm_layer=nn.BatchNorm2d):
|
||||
super(BasicBlockV1b, self).__init__()
|
||||
self.conv1 = nn.Conv2d(inplanes, planes, 3, stride,
|
||||
dilation, dilation, bias=False)
|
||||
self.bn1 = norm_layer(planes)
|
||||
self.relu = nn.ReLU(True)
|
||||
self.conv2 = nn.Conv2d(planes, planes, 3, 1, previous_dilation,
|
||||
dilation=previous_dilation, bias=False)
|
||||
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 BottleneckV1b(nn.Module):
|
||||
expansion = 4
|
||||
|
||||
def __init__(self, inplanes, planes, stride=1, dilation=1, downsample=None,
|
||||
previous_dilation=1, norm_layer=nn.BatchNorm2d):
|
||||
super(BottleneckV1b, self).__init__()
|
||||
self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False)
|
||||
self.bn1 = norm_layer(planes)
|
||||
self.conv2 = nn.Conv2d(planes, planes, 3, stride,
|
||||
dilation, dilation, bias=False)
|
||||
self.bn2 = norm_layer(planes)
|
||||
self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False)
|
||||
self.bn3 = norm_layer(planes * self.expansion)
|
||||
self.relu = nn.ReLU(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.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 ResNetV1b(nn.Module):
|
||||
|
||||
def __init__(self, block, layers, num_classes=1000, dilated=True, deep_stem=False,
|
||||
zero_init_residual=False, norm_layer=nn.BatchNorm2d):
|
||||
self.inplanes = 128 if deep_stem else 64
|
||||
super(ResNetV1b, self).__init__()
|
||||
if deep_stem:
|
||||
self.conv1 = nn.Sequential(
|
||||
nn.Conv2d(3, 64, 3, 2, 1, bias=False),
|
||||
norm_layer(64),
|
||||
nn.ReLU(True),
|
||||
nn.Conv2d(64, 64, 3, 1, 1, bias=False),
|
||||
norm_layer(64),
|
||||
nn.ReLU(True),
|
||||
nn.Conv2d(64, 128, 3, 1, 1, bias=False)
|
||||
)
|
||||
else:
|
||||
self.conv1 = nn.Conv2d(3, 64, 7, 2, 3, bias=False)
|
||||
self.bn1 = norm_layer(self.inplanes)
|
||||
self.relu = nn.ReLU(True)
|
||||
self.maxpool = nn.MaxPool2d(3, 2, 1)
|
||||
self.layer1 = self._make_layer(block, 64, layers[0], norm_layer=norm_layer)
|
||||
self.layer2 = self._make_layer(block, 128, layers[1], stride=2, norm_layer=norm_layer)
|
||||
if dilated:
|
||||
self.layer3 = self._make_layer(block, 256, layers[2], stride=1, dilation=2, norm_layer=norm_layer)
|
||||
self.layer4 = self._make_layer(block, 512, layers[3], stride=1, dilation=4, norm_layer=norm_layer)
|
||||
else:
|
||||
self.layer3 = self._make_layer(block, 256, layers[2], stride=2, norm_layer=norm_layer)
|
||||
self.layer4 = self._make_layer(block, 512, layers[3], stride=2, norm_layer=norm_layer)
|
||||
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
|
||||
self.fc = nn.Linear(512 * block.expansion, num_classes)
|
||||
|
||||
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.init.constant_(m.weight, 1)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
if zero_init_residual:
|
||||
for m in self.modules():
|
||||
if isinstance(m, BottleneckV1b):
|
||||
nn.init.constant_(m.bn3.weight, 0)
|
||||
elif isinstance(m, BasicBlockV1b):
|
||||
nn.init.constant_(m.bn2.weight, 0)
|
||||
|
||||
def _make_layer(self, block, planes, blocks, stride=1, dilation=1, norm_layer=nn.BatchNorm2d):
|
||||
downsample = None
|
||||
if stride != 1 or self.inplanes != planes * block.expansion:
|
||||
downsample = nn.Sequential(
|
||||
nn.Conv2d(self.inplanes, planes * block.expansion, 1, stride, bias=False),
|
||||
norm_layer(planes * block.expansion),
|
||||
)
|
||||
|
||||
layers = []
|
||||
if dilation in (1, 2):
|
||||
layers.append(block(self.inplanes, planes, stride, dilation=1, downsample=downsample,
|
||||
previous_dilation=dilation, norm_layer=norm_layer))
|
||||
elif dilation == 4:
|
||||
layers.append(block(self.inplanes, planes, stride, dilation=2, downsample=downsample,
|
||||
previous_dilation=dilation, norm_layer=norm_layer))
|
||||
else:
|
||||
raise RuntimeError("=> unknown dilation size: {}".format(dilation))
|
||||
self.inplanes = planes * block.expansion
|
||||
for _ in range(1, blocks):
|
||||
layers.append(block(self.inplanes, planes, dilation=dilation,
|
||||
previous_dilation=dilation, norm_layer=norm_layer))
|
||||
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.relu(x)
|
||||
x = self.maxpool(x)
|
||||
|
||||
x = self.layer1(x)
|
||||
x = self.layer2(x)
|
||||
x = self.layer3(x)
|
||||
x = self.layer4(x)
|
||||
|
||||
x = self.avgpool(x)
|
||||
x = x.view(x.size(0), -1)
|
||||
x = self.fc(x)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
def resnet18_v1b(pretrained=False, **kwargs):
|
||||
model = ResNetV1b(BasicBlockV1b, [2, 2, 2, 2], **kwargs)
|
||||
if pretrained:
|
||||
old_dict = model_zoo.load_url(model_urls['resnet18'])
|
||||
model_dict = model.state_dict()
|
||||
old_dict = {k: v for k, v in old_dict.items() if (k in model_dict)}
|
||||
model_dict.update(old_dict)
|
||||
model.load_state_dict(model_dict)
|
||||
return model
|
||||
|
||||
|
||||
def resnet34_v1b(pretrained=False, **kwargs):
|
||||
model = ResNetV1b(BasicBlockV1b, [3, 4, 6, 3], **kwargs)
|
||||
if pretrained:
|
||||
old_dict = model_zoo.load_url(model_urls['resnet34'])
|
||||
model_dict = model.state_dict()
|
||||
old_dict = {k: v for k, v in old_dict.items() if (k in model_dict)}
|
||||
model_dict.update(old_dict)
|
||||
model.load_state_dict(model_dict)
|
||||
return model
|
||||
|
||||
|
||||
def resnet50_v1b(pretrained=False, **kwargs):
|
||||
model = ResNetV1b(BottleneckV1b, [3, 4, 6, 3], **kwargs)
|
||||
if pretrained:
|
||||
old_dict = model_zoo.load_url(model_urls['resnet50'])
|
||||
model_dict = model.state_dict()
|
||||
old_dict = {k: v for k, v in old_dict.items() if (k in model_dict)}
|
||||
model_dict.update(old_dict)
|
||||
model.load_state_dict(model_dict)
|
||||
return model
|
||||
|
||||
|
||||
def resnet101_v1b(pretrained=False, **kwargs):
|
||||
model = ResNetV1b(BottleneckV1b, [3, 4, 23, 3], **kwargs)
|
||||
if pretrained:
|
||||
old_dict = model_zoo.load_url(model_urls['resnet101'])
|
||||
model_dict = model.state_dict()
|
||||
old_dict = {k: v for k, v in old_dict.items() if (k in model_dict)}
|
||||
model_dict.update(old_dict)
|
||||
model.load_state_dict(model_dict)
|
||||
return model
|
||||
|
||||
|
||||
def resnet152_v1b(pretrained=False, **kwargs):
|
||||
model = ResNetV1b(BottleneckV1b, [3, 8, 36, 3], **kwargs)
|
||||
if pretrained:
|
||||
old_dict = model_zoo.load_url(model_urls['resnet152'])
|
||||
model_dict = model.state_dict()
|
||||
old_dict = {k: v for k, v in old_dict.items() if (k in model_dict)}
|
||||
model_dict.update(old_dict)
|
||||
model.load_state_dict(model_dict)
|
||||
return model
|
||||
|
||||
|
||||
def resnet50_v1s(pretrained=False, root='~/.torch/models', **kwargs):
|
||||
model = ResNetV1b(BottleneckV1b, [3, 4, 6, 3], deep_stem=True, **kwargs)
|
||||
if pretrained:
|
||||
from ..model_store import get_resnet_file
|
||||
model.load_state_dict(torch.load(get_resnet_file('resnet50', root=root)), strict=False)
|
||||
return model
|
||||
|
||||
|
||||
def resnet101_v1s(pretrained=False, root='~/.torch/models', **kwargs):
|
||||
model = ResNetV1b(BottleneckV1b, [3, 4, 23, 3], deep_stem=True, **kwargs)
|
||||
if pretrained:
|
||||
from ..model_store import get_resnet_file
|
||||
model.load_state_dict(torch.load(get_resnet_file('resnet101', root=root)), strict=False)
|
||||
return model
|
||||
|
||||
|
||||
def resnet152_v1s(pretrained=False, root='~/.torch/models', **kwargs):
|
||||
model = ResNetV1b(BottleneckV1b, [3, 8, 36, 3], deep_stem=True, **kwargs)
|
||||
if pretrained:
|
||||
from ..model_store import get_resnet_file
|
||||
model.load_state_dict(torch.load(get_resnet_file('resnet152', root=root)), strict=False)
|
||||
return model
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
import torch
|
||||
|
||||
img = torch.randn(4, 3, 224, 224)
|
||||
model = resnet50_v1b(True)
|
||||
output = model(img)
|
||||
@@ -0,0 +1,60 @@
|
||||
"""Base Model for Semantic Segmentation"""
|
||||
import torch.nn as nn
|
||||
|
||||
from seg.networks.jpu import JPU
|
||||
from seg.networks.resnetv1b import resnet50_v1s, resnet101_v1s, resnet152_v1s
|
||||
|
||||
__all__ = ['SegBaseModel']
|
||||
|
||||
|
||||
class SegBaseModel(nn.Module):
|
||||
r"""Base Model for Semantic Segmentation
|
||||
|
||||
Parameters
|
||||
----------
|
||||
backbone : string
|
||||
Pre-trained dilated backbone network type (default:'resnet50'; 'resnet50',
|
||||
'resnet101' or 'resnet152').
|
||||
"""
|
||||
|
||||
def __init__(self, nclass, aux, backbone='resnet50', jpu=False, pretrained_base=True, **kwargs):
|
||||
super(SegBaseModel, self).__init__()
|
||||
dilated = False if jpu else True
|
||||
self.aux = aux
|
||||
self.nclass = nclass
|
||||
if backbone == 'resnet50':
|
||||
self.pretrained = resnet50_v1s(pretrained=pretrained_base, dilated=dilated, **kwargs)
|
||||
elif backbone == 'resnet101':
|
||||
self.pretrained = resnet101_v1s(pretrained=pretrained_base, dilated=dilated, **kwargs)
|
||||
elif backbone == 'resnet152':
|
||||
self.pretrained = resnet152_v1s(pretrained=pretrained_base, dilated=dilated, **kwargs)
|
||||
else:
|
||||
raise RuntimeError('unknown backbone: {}'.format(backbone))
|
||||
|
||||
self.jpu = JPU([512, 1024, 2048], width=512, **kwargs) if jpu else None
|
||||
|
||||
def base_forward(self, x):
|
||||
"""forwarding pre-trained network"""
|
||||
x = self.pretrained.conv1(x)
|
||||
x = self.pretrained.bn1(x)
|
||||
x = self.pretrained.relu(x)
|
||||
x = self.pretrained.maxpool(x)
|
||||
c1 = self.pretrained.layer1(x)
|
||||
c2 = self.pretrained.layer2(c1)
|
||||
c3 = self.pretrained.layer3(c2)
|
||||
c4 = self.pretrained.layer4(c3)
|
||||
|
||||
if self.jpu:
|
||||
return self.jpu(c1, c2, c3, c4)
|
||||
else:
|
||||
return c1, c2, c3, c4
|
||||
|
||||
def evaluate(self, x):
|
||||
"""evaluating network with inputs and targets"""
|
||||
return self.forward(x)[0]
|
||||
|
||||
def demo(self, x):
|
||||
pred = self.forward(x)
|
||||
if self.aux:
|
||||
pred = pred[0]
|
||||
return pred
|
||||
@@ -0,0 +1,191 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.utils.model_zoo as model_zoo
|
||||
|
||||
__all__ = [
|
||||
'VGG', 'vgg11', 'vgg11_bn', 'vgg13', 'vgg13_bn', 'vgg16', 'vgg16_bn',
|
||||
'vgg19_bn', 'vgg19',
|
||||
]
|
||||
|
||||
model_urls = {
|
||||
'vgg11': 'https://download.pytorch.org/models/vgg11-bbd30ac9.pth',
|
||||
'vgg13': 'https://download.pytorch.org/models/vgg13-c768596a.pth',
|
||||
'vgg16': 'https://download.pytorch.org/models/vgg16-397923af.pth',
|
||||
'vgg19': 'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth',
|
||||
'vgg11_bn': 'https://download.pytorch.org/models/vgg11_bn-6002323d.pth',
|
||||
'vgg13_bn': 'https://download.pytorch.org/models/vgg13_bn-abd245e5.pth',
|
||||
'vgg16_bn': 'https://download.pytorch.org/models/vgg16_bn-6c64b313.pth',
|
||||
'vgg19_bn': 'https://download.pytorch.org/models/vgg19_bn-c79401a0.pth',
|
||||
}
|
||||
|
||||
|
||||
class VGG(nn.Module):
|
||||
def __init__(self, features, num_classes=1000, init_weights=True):
|
||||
super(VGG, self).__init__()
|
||||
self.features = features
|
||||
self.avgpool = nn.AdaptiveAvgPool2d((7, 7))
|
||||
self.classifier = nn.Sequential(
|
||||
nn.Linear(512 * 7 * 7, 4096),
|
||||
nn.ReLU(True),
|
||||
nn.Dropout(),
|
||||
nn.Linear(4096, 4096),
|
||||
nn.ReLU(True),
|
||||
nn.Dropout(),
|
||||
nn.Linear(4096, num_classes)
|
||||
)
|
||||
if init_weights:
|
||||
self._initialize_weights()
|
||||
|
||||
def forward(self, x):
|
||||
x = self.features(x)
|
||||
x = self.avgpool(x)
|
||||
x = x.view(x.size(0), -1)
|
||||
x = self.classifier(x)
|
||||
return x
|
||||
|
||||
def _initialize_weights(self):
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
nn.init.constant_(m.weight, 1)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
elif isinstance(m, nn.Linear):
|
||||
nn.init.normal_(m.weight, 0, 0.01)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
|
||||
def make_layers(cfg, batch_norm=False):
|
||||
layers = []
|
||||
in_channels = 3
|
||||
for v in cfg:
|
||||
if v == 'M':
|
||||
layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
|
||||
else:
|
||||
conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
|
||||
if batch_norm:
|
||||
layers += (conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True))
|
||||
else:
|
||||
layers += [conv2d, nn.ReLU(inplace=True)]
|
||||
in_channels = v
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
|
||||
cfg = {
|
||||
'A': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
|
||||
'B': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
|
||||
'D': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],
|
||||
'E': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'],
|
||||
}
|
||||
|
||||
|
||||
def vgg11(pretrained=False, **kwargs):
|
||||
"""VGG 11-layer model (configuration "A")
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
"""
|
||||
if pretrained:
|
||||
kwargs['init_weights'] = False
|
||||
model = VGG(make_layers(cfg['A']), **kwargs)
|
||||
if pretrained:
|
||||
model.load_state_dict(model_zoo.load_url(model_urls['vgg11']))
|
||||
return model
|
||||
|
||||
|
||||
def vgg11_bn(pretrained=False, **kwargs):
|
||||
"""VGG 11-layer model (configuration "A") with batch normalization
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
"""
|
||||
if pretrained:
|
||||
kwargs['init_weights'] = False
|
||||
model = VGG(make_layers(cfg['A'], batch_norm=True), **kwargs)
|
||||
if pretrained:
|
||||
model.load_state_dict(model_zoo.load_url(model_urls['vgg11_bn']))
|
||||
return model
|
||||
|
||||
|
||||
def vgg13(pretrained=False, **kwargs):
|
||||
"""VGG 13-layer model (configuration "B")
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
"""
|
||||
if pretrained:
|
||||
kwargs['init_weights'] = False
|
||||
model = VGG(make_layers(cfg['B']), **kwargs)
|
||||
if pretrained:
|
||||
model.load_state_dict(model_zoo.load_url(model_urls['vgg13']))
|
||||
return model
|
||||
|
||||
|
||||
def vgg13_bn(pretrained=False, **kwargs):
|
||||
"""VGG 13-layer model (configuration "B") with batch normalization
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
"""
|
||||
if pretrained:
|
||||
kwargs['init_weights'] = False
|
||||
model = VGG(make_layers(cfg['B'], batch_norm=True), **kwargs)
|
||||
if pretrained:
|
||||
model.load_state_dict(model_zoo.load_url(model_urls['vgg13_bn']))
|
||||
return model
|
||||
|
||||
|
||||
def vgg16(pretrained=False, **kwargs):
|
||||
"""VGG 16-layer model (configuration "D")
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
"""
|
||||
if pretrained:
|
||||
kwargs['init_weights'] = False
|
||||
model = VGG(make_layers(cfg['D']), **kwargs)
|
||||
if pretrained:
|
||||
model.load_state_dict(model_zoo.load_url(model_urls['vgg16']))
|
||||
return model
|
||||
|
||||
|
||||
def vgg16_bn(pretrained=False, **kwargs):
|
||||
"""VGG 16-layer model (configuration "D") with batch normalization
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
"""
|
||||
if pretrained:
|
||||
kwargs['init_weights'] = False
|
||||
model = VGG(make_layers(cfg['D'], batch_norm=True), **kwargs)
|
||||
if pretrained:
|
||||
model.load_state_dict(model_zoo.load_url(model_urls['vgg16_bn']))
|
||||
return model
|
||||
|
||||
|
||||
def vgg19(pretrained=False, **kwargs):
|
||||
"""VGG 19-layer model (configuration "E")
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
"""
|
||||
if pretrained:
|
||||
kwargs['init_weights'] = False
|
||||
model = VGG(make_layers(cfg['E']), **kwargs)
|
||||
if pretrained:
|
||||
model.load_state_dict(model_zoo.load_url(model_urls['vgg19']))
|
||||
return model
|
||||
|
||||
|
||||
def vgg19_bn(pretrained=False, **kwargs):
|
||||
"""VGG 19-layer model (configuration 'E') with batch normalization
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
"""
|
||||
if pretrained:
|
||||
kwargs['init_weights'] = False
|
||||
model = VGG(make_layers(cfg['E'], batch_norm=True), **kwargs)
|
||||
if pretrained:
|
||||
model.load_state_dict(model_zoo.load_url(model_urls['vgg19_bn']))
|
||||
return model
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
img = torch.randn((4, 3, 480, 480))
|
||||
model = vgg16(pretrained=False)
|
||||
out = model(img)
|
||||
@@ -0,0 +1,411 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
__all__ = ['Enc', 'FCAttention', 'Xception65', 'Xception71', 'get_xception', 'get_xception_71', 'get_xception_a']
|
||||
|
||||
|
||||
class SeparableConv2d(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, dilation=1, bias=False, norm_layer=None):
|
||||
super(SeparableConv2d, self).__init__()
|
||||
self.kernel_size = kernel_size
|
||||
self.dilation = dilation
|
||||
|
||||
self.conv1 = nn.Conv2d(in_channels, in_channels, kernel_size, stride, 0, dilation, groups=in_channels,
|
||||
bias=bias)
|
||||
self.bn = norm_layer(in_channels)
|
||||
self.pointwise = nn.Conv2d(in_channels, out_channels, 1, bias=bias)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.fix_padding(x, self.kernel_size, self.dilation)
|
||||
x = self.conv1(x)
|
||||
x = self.bn(x)
|
||||
x = self.pointwise(x)
|
||||
|
||||
return x
|
||||
|
||||
def fix_padding(self, x, kernel_size, dilation):
|
||||
kernel_size_effective = kernel_size + (kernel_size - 1) * (dilation - 1)
|
||||
pad_total = kernel_size_effective - 1
|
||||
pad_beg = pad_total // 2
|
||||
pad_end = pad_total - pad_beg
|
||||
padded_inputs = F.pad(x, (pad_beg, pad_end, pad_beg, pad_end))
|
||||
return padded_inputs
|
||||
|
||||
|
||||
class Block(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, reps, stride=1, dilation=1, norm_layer=None,
|
||||
start_with_relu=True, grow_first=True, is_last=False):
|
||||
super(Block, self).__init__()
|
||||
if out_channels != in_channels or stride != 1:
|
||||
self.skip = nn.Conv2d(in_channels, out_channels, 1, stride, bias=False)
|
||||
self.skipbn = norm_layer(out_channels)
|
||||
else:
|
||||
self.skip = None
|
||||
self.relu = nn.ReLU(True)
|
||||
rep = list()
|
||||
filters = in_channels
|
||||
if grow_first:
|
||||
if start_with_relu:
|
||||
rep.append(self.relu)
|
||||
rep.append(SeparableConv2d(in_channels, out_channels, 3, 1, dilation, norm_layer=norm_layer))
|
||||
rep.append(norm_layer(out_channels))
|
||||
filters = out_channels
|
||||
for i in range(reps - 1):
|
||||
if grow_first or start_with_relu:
|
||||
rep.append(self.relu)
|
||||
rep.append(SeparableConv2d(filters, filters, 3, 1, dilation, norm_layer=norm_layer))
|
||||
rep.append(norm_layer(filters))
|
||||
if not grow_first:
|
||||
rep.append(self.relu)
|
||||
rep.append(SeparableConv2d(in_channels, out_channels, 3, 1, dilation, norm_layer=norm_layer))
|
||||
if stride != 1:
|
||||
rep.append(self.relu)
|
||||
rep.append(SeparableConv2d(out_channels, out_channels, 3, stride, norm_layer=norm_layer))
|
||||
rep.append(norm_layer(out_channels))
|
||||
elif is_last:
|
||||
rep.append(self.relu)
|
||||
rep.append(SeparableConv2d(out_channels, out_channels, 3, 1, dilation, norm_layer=norm_layer))
|
||||
rep.append(norm_layer(out_channels))
|
||||
self.rep = nn.Sequential(*rep)
|
||||
|
||||
def forward(self, x):
|
||||
out = self.rep(x)
|
||||
if self.skip is not None:
|
||||
skip = self.skipbn(self.skip(x))
|
||||
else:
|
||||
skip = x
|
||||
out = out + skip
|
||||
return out
|
||||
|
||||
|
||||
class Xception65(nn.Module):
|
||||
"""Modified Aligned Xception
|
||||
"""
|
||||
|
||||
def __init__(self, num_classes=1000, output_stride=32, norm_layer=nn.BatchNorm2d):
|
||||
super(Xception65, self).__init__()
|
||||
if output_stride == 32:
|
||||
entry_block3_stride = 2
|
||||
exit_block20_stride = 2
|
||||
middle_block_dilation = 1
|
||||
exit_block_dilations = (1, 1)
|
||||
elif output_stride == 16:
|
||||
entry_block3_stride = 2
|
||||
exit_block20_stride = 1
|
||||
middle_block_dilation = 1
|
||||
exit_block_dilations = (1, 2)
|
||||
elif output_stride == 8:
|
||||
entry_block3_stride = 1
|
||||
exit_block20_stride = 1
|
||||
middle_block_dilation = 2
|
||||
exit_block_dilations = (2, 4)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
# Entry flow
|
||||
self.conv1 = nn.Conv2d(3, 32, 3, 2, 1, bias=False)
|
||||
self.bn1 = norm_layer(32)
|
||||
self.relu = nn.ReLU(True)
|
||||
|
||||
self.conv2 = nn.Conv2d(32, 64, 3, 1, 1, bias=False)
|
||||
self.bn2 = norm_layer(64)
|
||||
|
||||
self.block1 = Block(64, 128, reps=2, stride=2, norm_layer=norm_layer, start_with_relu=False)
|
||||
self.block2 = Block(128, 256, reps=2, stride=2, norm_layer=norm_layer, start_with_relu=False, grow_first=True)
|
||||
self.block3 = Block(256, 728, reps=2, stride=entry_block3_stride, norm_layer=norm_layer,
|
||||
start_with_relu=True, grow_first=True, is_last=True)
|
||||
|
||||
# Middle flow
|
||||
midflow = list()
|
||||
for i in range(4, 20):
|
||||
midflow.append(Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation, norm_layer=norm_layer,
|
||||
start_with_relu=True, grow_first=True))
|
||||
self.midflow = nn.Sequential(*midflow)
|
||||
|
||||
# Exit flow
|
||||
self.block20 = Block(728, 1024, reps=2, stride=exit_block20_stride, dilation=exit_block_dilations[0],
|
||||
norm_layer=norm_layer, start_with_relu=True, grow_first=False, is_last=True)
|
||||
self.conv3 = SeparableConv2d(1024, 1536, 3, 1, dilation=exit_block_dilations[1], norm_layer=norm_layer)
|
||||
self.bn3 = norm_layer(1536)
|
||||
self.conv4 = SeparableConv2d(1536, 1536, 3, stride=1, dilation=exit_block_dilations[1], norm_layer=norm_layer)
|
||||
self.bn4 = norm_layer(1536)
|
||||
self.conv5 = SeparableConv2d(1536, 2048, 3, 1, dilation=exit_block_dilations[1], norm_layer=norm_layer)
|
||||
self.bn5 = norm_layer(2048)
|
||||
self.avgpool = nn.AdaptiveAvgPool2d(1)
|
||||
self.fc = nn.Linear(2048, num_classes)
|
||||
|
||||
def forward(self, x):
|
||||
# Entry flow
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.relu(x)
|
||||
|
||||
x = self.conv2(x)
|
||||
x = self.bn2(x)
|
||||
x = self.relu(x)
|
||||
|
||||
x = self.block1(x)
|
||||
x = self.relu(x)
|
||||
# c1 = x
|
||||
x = self.block2(x)
|
||||
# c2 = x
|
||||
x = self.block3(x)
|
||||
|
||||
# Middle flow
|
||||
x = self.midflow(x)
|
||||
# c3 = x
|
||||
|
||||
# Exit flow
|
||||
x = self.block20(x)
|
||||
x = self.relu(x)
|
||||
x = self.conv3(x)
|
||||
x = self.bn3(x)
|
||||
x = self.relu(x)
|
||||
|
||||
x = self.conv4(x)
|
||||
x = self.bn4(x)
|
||||
x = self.relu(x)
|
||||
|
||||
x = self.conv5(x)
|
||||
x = self.bn5(x)
|
||||
x = self.relu(x)
|
||||
|
||||
x = self.avgpool(x)
|
||||
x = x.view(x.size(0), -1)
|
||||
x = self.fc(x)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class Xception71(nn.Module):
|
||||
"""Modified Aligned Xception
|
||||
"""
|
||||
|
||||
def __init__(self, num_classes=1000, output_stride=32, norm_layer=nn.BatchNorm2d):
|
||||
super(Xception71, self).__init__()
|
||||
if output_stride == 32:
|
||||
entry_block3_stride = 2
|
||||
exit_block20_stride = 2
|
||||
middle_block_dilation = 1
|
||||
exit_block_dilations = (1, 1)
|
||||
elif output_stride == 16:
|
||||
entry_block3_stride = 2
|
||||
exit_block20_stride = 1
|
||||
middle_block_dilation = 1
|
||||
exit_block_dilations = (1, 2)
|
||||
elif output_stride == 8:
|
||||
entry_block3_stride = 1
|
||||
exit_block20_stride = 1
|
||||
middle_block_dilation = 2
|
||||
exit_block_dilations = (2, 4)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
# Entry flow
|
||||
self.conv1 = nn.Conv2d(3, 32, 3, 2, 1, bias=False)
|
||||
self.bn1 = norm_layer(32)
|
||||
self.relu = nn.ReLU(True)
|
||||
|
||||
self.conv2 = nn.Conv2d(32, 64, 3, 1, 1, bias=False)
|
||||
self.bn2 = norm_layer(64)
|
||||
|
||||
self.block1 = Block(64, 128, reps=2, stride=2, norm_layer=norm_layer, start_with_relu=False)
|
||||
self.block2 = nn.Sequential(
|
||||
Block(128, 256, reps=2, stride=2, norm_layer=norm_layer, start_with_relu=False, grow_first=True),
|
||||
Block(256, 728, reps=2, stride=2, norm_layer=norm_layer, start_with_relu=False, grow_first=True))
|
||||
self.block3 = Block(728, 728, reps=2, stride=entry_block3_stride, norm_layer=norm_layer,
|
||||
start_with_relu=True, grow_first=True, is_last=True)
|
||||
|
||||
# Middle flow
|
||||
midflow = list()
|
||||
for i in range(4, 20):
|
||||
midflow.append(Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation, norm_layer=norm_layer,
|
||||
start_with_relu=True, grow_first=True))
|
||||
self.midflow = nn.Sequential(*midflow)
|
||||
|
||||
# Exit flow
|
||||
self.block20 = Block(728, 1024, reps=2, stride=exit_block20_stride, dilation=exit_block_dilations[0],
|
||||
norm_layer=norm_layer, start_with_relu=True, grow_first=False, is_last=True)
|
||||
self.conv3 = SeparableConv2d(1024, 1536, 3, 1, dilation=exit_block_dilations[1], norm_layer=norm_layer)
|
||||
self.bn3 = norm_layer(1536)
|
||||
self.conv4 = SeparableConv2d(1536, 1536, 3, stride=1, dilation=exit_block_dilations[1], norm_layer=norm_layer)
|
||||
self.bn4 = norm_layer(1536)
|
||||
self.conv5 = SeparableConv2d(1536, 2048, 3, 1, dilation=exit_block_dilations[1], norm_layer=norm_layer)
|
||||
self.bn5 = norm_layer(2048)
|
||||
self.avgpool = nn.AdaptiveAvgPool2d(1)
|
||||
self.fc = nn.Linear(2048, num_classes)
|
||||
|
||||
def forward(self, x):
|
||||
# Entry flow
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.relu(x)
|
||||
|
||||
x = self.conv2(x)
|
||||
x = self.bn2(x)
|
||||
x = self.relu(x)
|
||||
|
||||
x = self.block1(x)
|
||||
x = self.relu(x)
|
||||
# c1 = x
|
||||
x = self.block2(x)
|
||||
# c2 = x
|
||||
x = self.block3(x)
|
||||
|
||||
# Middle flow
|
||||
x = self.midflow(x)
|
||||
# c3 = x
|
||||
|
||||
# Exit flow
|
||||
x = self.block20(x)
|
||||
x = self.relu(x)
|
||||
x = self.conv3(x)
|
||||
x = self.bn3(x)
|
||||
x = self.relu(x)
|
||||
|
||||
x = self.conv4(x)
|
||||
x = self.bn4(x)
|
||||
x = self.relu(x)
|
||||
|
||||
x = self.conv5(x)
|
||||
x = self.bn5(x)
|
||||
x = self.relu(x)
|
||||
|
||||
x = self.avgpool(x)
|
||||
x = x.view(x.size(0), -1)
|
||||
x = self.fc(x)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
# -------------------------------------------------
|
||||
# For DFANet
|
||||
# -------------------------------------------------
|
||||
class BlockA(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, stride=1, dilation=1, norm_layer=None, start_with_relu=True):
|
||||
super(BlockA, self).__init__()
|
||||
if out_channels != in_channels or stride != 1:
|
||||
self.skip = nn.Conv2d(in_channels, out_channels, 1, stride, bias=False)
|
||||
self.skipbn = norm_layer(out_channels)
|
||||
else:
|
||||
self.skip = None
|
||||
self.relu = nn.ReLU(True)
|
||||
rep = list()
|
||||
inter_channels = out_channels // 4
|
||||
|
||||
if start_with_relu:
|
||||
rep.append(self.relu)
|
||||
rep.append(SeparableConv2d(in_channels, inter_channels, 3, 1, dilation, norm_layer=norm_layer))
|
||||
rep.append(norm_layer(inter_channels))
|
||||
|
||||
rep.append(self.relu)
|
||||
rep.append(SeparableConv2d(inter_channels, inter_channels, 3, 1, dilation, norm_layer=norm_layer))
|
||||
rep.append(norm_layer(inter_channels))
|
||||
|
||||
if stride != 1:
|
||||
rep.append(self.relu)
|
||||
rep.append(SeparableConv2d(inter_channels, out_channels, 3, stride, norm_layer=norm_layer))
|
||||
rep.append(norm_layer(out_channels))
|
||||
else:
|
||||
rep.append(self.relu)
|
||||
rep.append(SeparableConv2d(inter_channels, out_channels, 3, 1, norm_layer=norm_layer))
|
||||
rep.append(norm_layer(out_channels))
|
||||
self.rep = nn.Sequential(*rep)
|
||||
|
||||
def forward(self, x):
|
||||
out = self.rep(x)
|
||||
if self.skip is not None:
|
||||
skip = self.skipbn(self.skip(x))
|
||||
else:
|
||||
skip = x
|
||||
out = out + skip
|
||||
return out
|
||||
|
||||
|
||||
class Enc(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, blocks, norm_layer=None):
|
||||
super(Enc, self).__init__()
|
||||
block = list()
|
||||
block.append(BlockA(in_channels, out_channels, 2, norm_layer=norm_layer))
|
||||
for i in range(blocks - 1):
|
||||
block.append(BlockA(out_channels, out_channels, 1, norm_layer=norm_layer))
|
||||
self.block = nn.Sequential(*block)
|
||||
|
||||
def forward(self, x):
|
||||
return self.block(x)
|
||||
|
||||
|
||||
class FCAttention(nn.Module):
|
||||
def __init__(self, in_channels, norm_layer=None):
|
||||
super(FCAttention, self).__init__()
|
||||
self.avgpool = nn.AdaptiveAvgPool2d(1)
|
||||
self.fc = nn.Linear(in_channels, 1000)
|
||||
self.conv = nn.Sequential(
|
||||
nn.Conv2d(1000, in_channels, 1, bias=False),
|
||||
norm_layer(in_channels),
|
||||
nn.ReLU(True))
|
||||
|
||||
def forward(self, x):
|
||||
n, c, _, _ = x.size()
|
||||
att = self.avgpool(x).view(n, c)
|
||||
att = self.fc(att).view(n, 1000, 1, 1)
|
||||
att = self.conv(att)
|
||||
return x * att.expand_as(x)
|
||||
|
||||
|
||||
class XceptionA(nn.Module):
|
||||
def __init__(self, num_classes=1000, norm_layer=nn.BatchNorm2d):
|
||||
super(XceptionA, self).__init__()
|
||||
self.conv1 = nn.Sequential(nn.Conv2d(3, 8, 3, 2, 1, bias=False),
|
||||
norm_layer(8),
|
||||
nn.ReLU(True))
|
||||
|
||||
self.enc2 = Enc(8, 48, 4, norm_layer=norm_layer)
|
||||
self.enc3 = Enc(48, 96, 6, norm_layer=norm_layer)
|
||||
self.enc4 = Enc(96, 192, 4, norm_layer=norm_layer)
|
||||
|
||||
self.fca = FCAttention(192, norm_layer=norm_layer)
|
||||
self.avgpool = nn.AdaptiveAvgPool2d(1)
|
||||
self.fc = nn.Linear(192, num_classes)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x)
|
||||
|
||||
x = self.enc2(x)
|
||||
x = self.enc3(x)
|
||||
x = self.enc4(x)
|
||||
x = self.fca(x)
|
||||
|
||||
x = self.avgpool(x)
|
||||
x = x.view(x.size(0), -1)
|
||||
x = self.fc(x)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
# Constructor
|
||||
def get_xception(pretrained=False, root='~/.torch/models', **kwargs):
|
||||
model = Xception65(**kwargs)
|
||||
if pretrained:
|
||||
from ..model_store import get_model_file
|
||||
model.load_state_dict(torch.load(get_model_file('xception', root=root)))
|
||||
return model
|
||||
|
||||
|
||||
def get_xception_71(pretrained=False, root='~/.torch/models', **kwargs):
|
||||
model = Xception71(**kwargs)
|
||||
if pretrained:
|
||||
from ..model_store import get_model_file
|
||||
model.load_state_dict(torch.load(get_model_file('xception71', root=root)))
|
||||
return model
|
||||
|
||||
|
||||
def get_xception_a(pretrained=False, root='~/.torch/models', **kwargs):
|
||||
model = XceptionA(**kwargs)
|
||||
if pretrained:
|
||||
from ..model_store import get_model_file
|
||||
model.load_state_dict(torch.load(get_model_file('xception_a', root=root)))
|
||||
return model
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
model = get_xception_a()
|
||||
@@ -0,0 +1,102 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding: utf-8 -*-
|
||||
|
||||
# @File : setup.py
|
||||
# @Time : 2020/1/15
|
||||
# @Author : yangchaojie (yangchaojie@immomo.com)
|
||||
|
||||
import os
|
||||
import sys
|
||||
import shutil
|
||||
import numpy
|
||||
import tempfile
|
||||
|
||||
from setuptools import setup
|
||||
from setuptools.extension import Extension
|
||||
|
||||
from Cython.Build import cythonize
|
||||
from Cython.Distutils import build_ext
|
||||
|
||||
import platform
|
||||
|
||||
|
||||
def get_root_path(root):
|
||||
if os.path.dirname(root) in ['', '.']:
|
||||
return os.path.basename(root)
|
||||
else:
|
||||
return get_root_path(os.path.dirname(root))
|
||||
|
||||
|
||||
def copy_file(src, dest):
|
||||
if os.path.exists(dest):
|
||||
return
|
||||
|
||||
if not os.path.exists(os.path.dirname(dest)):
|
||||
os.makedirs(os.path.dirname(dest))
|
||||
if os.path.isdir(src):
|
||||
shutil.copytree(src, dest)
|
||||
else:
|
||||
shutil.copyfile(src, dest)
|
||||
|
||||
|
||||
def touch_init_file():
|
||||
init_file_name = os.path.join(tempfile.mkdtemp(), '__init__.py')
|
||||
with open(init_file_name, 'w'):
|
||||
pass
|
||||
return init_file_name
|
||||
|
||||
|
||||
|
||||
|
||||
def compose_extensions(root='.'):
|
||||
for file_ in os.listdir(root):
|
||||
abs_file = os.path.join(root, file_)
|
||||
|
||||
if os.path.isfile(abs_file):
|
||||
if abs_file.endswith('.py'):
|
||||
extensions.append(Extension(get_root_path(abs_file) + '.*', [abs_file]))
|
||||
elif abs_file.endswith('.c') or abs_file.endswith('.pyc'):
|
||||
continue
|
||||
else:
|
||||
copy_file(abs_file, os.path.join(build_root_dir, abs_file))
|
||||
if abs_file.endswith('__init__.py'):
|
||||
copy_file(init_file, os.path.join(build_root_dir, abs_file))
|
||||
|
||||
else:
|
||||
if os.path.basename(abs_file) in ignore_folders :
|
||||
continue
|
||||
if os.path.basename(abs_file) in conf_folders:
|
||||
copy_file(abs_file, os.path.join(build_root_dir, abs_file))
|
||||
compose_extensions(abs_file)
|
||||
|
||||
|
||||
|
||||
# if __name__ == '__main__':
|
||||
build_root_dir = 'build/lib.' + platform.system().lower() + '-' + platform.machine() + '-' + str(
|
||||
sys.version_info.major) + '.' + str(sys.version_info.minor)
|
||||
|
||||
print(build_root_dir)
|
||||
|
||||
extensions = []
|
||||
ignore_folders = ['build', 'new_ref_zao_color_0818', 'ref_hair_online_0703_local', 'ref_分好类别', '.git']
|
||||
conf_folders = ['conf']
|
||||
|
||||
|
||||
init_file = touch_init_file()
|
||||
print(init_file)
|
||||
|
||||
|
||||
compose_extensions()
|
||||
os.remove(init_file)
|
||||
|
||||
setup(
|
||||
name='moxie_hairstyle',
|
||||
version='1.0',
|
||||
ext_modules=cythonize(
|
||||
extensions,
|
||||
nthreads=16,
|
||||
compiler_directives=dict(always_allow_keywords=True),
|
||||
include_path=[numpy.get_include()]),
|
||||
cmdclass=dict(build_ext=build_ext))
|
||||
|
||||
# python setup.py build_ext
|
||||
@@ -0,0 +1,24 @@
|
||||
from seg.hairseg_single_model import Evaluator
|
||||
import os
|
||||
import cv2
|
||||
import numpy as np
|
||||
if __name__ == "__main__":
|
||||
|
||||
data_path = "/home/liyang/project/matting/合格/origin"
|
||||
dst_path = "/home/liyang/project/matting/合格/origin_seg_res1102"
|
||||
if not os.path.exists(dst_path):
|
||||
os.mkdir(dst_path)
|
||||
seg_model = Evaluator(gpu_id=0, output_img_size=512, nclass=3)
|
||||
for imgs in os.listdir(data_path):
|
||||
if imgs.endswith(".txt"):
|
||||
continue
|
||||
|
||||
img_path = os.path.join(data_path, imgs)
|
||||
img = cv2.imread(img_path)
|
||||
if img_path.endswith('.png'):
|
||||
kpts_1k = np.loadtxt(img_path.replace('.png', '_landmark1k.txt'))
|
||||
elif img_path.endswith('.jpg'):
|
||||
kpts_1k = np.loadtxt(img_path.replace('.jpg', '_landmark1k.txt'))
|
||||
output_img_size = 512
|
||||
mask = seg_model.eval(img, kpts_1k)
|
||||
cv2.imwrite(os.path.join(dst_path, imgs), mask)
|
||||
Reference in New Issue
Block a user