初始化换发型项目：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,330 @@
|
||||
import sys
|
||||
# sys.path.append("/Users/momo/human_seg_train")
|
||||
# print(sys.path)
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from bodyseg.backbone.backbone import get_norm
|
||||
|
||||
class ConvBNReLU(nn.Sequential):
|
||||
def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, groups=1, norm_layer=None):
|
||||
padding = (kernel_size - 1) // 2
|
||||
if norm_layer is None:
|
||||
norm_layer = nn.BatchNorm2d
|
||||
super(ConvBNReLU, self).__init__(
|
||||
nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=False),
|
||||
norm_layer(out_planes),
|
||||
nn.ReLU6(inplace=True)
|
||||
)
|
||||
|
||||
class InvertedResidual(nn.Module):
|
||||
def __init__(self, inp, oup, stride, expand_ratio, norm_layer=None):
|
||||
super(InvertedResidual, self).__init__()
|
||||
self.stride = stride
|
||||
assert stride in [1, 2]
|
||||
|
||||
if norm_layer is None:
|
||||
norm_layer = nn.BatchNorm2d
|
||||
|
||||
hidden_dim = int(round(inp * expand_ratio))
|
||||
self.use_res_connect = self.stride == 1 and inp == oup
|
||||
|
||||
layers = []
|
||||
if expand_ratio != 1:
|
||||
# pw
|
||||
layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1, norm_layer=norm_layer))
|
||||
layers.extend([
|
||||
# dw
|
||||
ConvBNReLU(hidden_dim, hidden_dim, stride=stride, groups=hidden_dim, norm_layer=norm_layer),
|
||||
# pw-linear
|
||||
nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
|
||||
norm_layer(oup),
|
||||
])
|
||||
self.conv = nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
if self.use_res_connect:
|
||||
return x + self.conv(x)
|
||||
else:
|
||||
return self.conv(x)
|
||||
|
||||
class UpSampleBlock(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, expand_ratio=6):
|
||||
super(UpSampleBlock, self).__init__()
|
||||
self.refine = InvertedResidual(in_channels, out_channels, 1, expand_ratio)
|
||||
|
||||
def forward(self, x0, x1):
|
||||
x = torch.cat([x0, x1], dim=1)
|
||||
x = self.refine(x)
|
||||
return x
|
||||
|
||||
class BodySegNet_32_thin_sigmod(nn.Module):
|
||||
def __init__(self, input_channels=3, class_nums=1, output_onnx=False):
|
||||
super(BodySegNet_32_thin_sigmod, self).__init__()
|
||||
self.class_nums = class_nums
|
||||
self.output_onnx = output_onnx
|
||||
|
||||
self.stage_1 = nn.Sequential(
|
||||
ConvBNReLU(input_channels, 16, kernel_size=3, stride=2)
|
||||
)
|
||||
self.stage_2 = nn.Sequential(
|
||||
ConvBNReLU(16, 16, kernel_size=3, stride=2, groups=16),
|
||||
ConvBNReLU(16, 16, kernel_size=1, stride=1),
|
||||
)
|
||||
self.stage_3 = nn.Sequential(
|
||||
InvertedResidual(16, 24, stride=2, expand_ratio=6),
|
||||
InvertedResidual(24, 24, stride=1, expand_ratio=6),
|
||||
InvertedResidual(24, 24, stride=1, expand_ratio=6),
|
||||
)
|
||||
self.stage_4 = nn.Sequential(
|
||||
InvertedResidual(24, 32, stride=2, expand_ratio=6),
|
||||
InvertedResidual(32, 32, stride=1, expand_ratio=6),
|
||||
InvertedResidual(32, 32, stride=1, expand_ratio=6),
|
||||
InvertedResidual(32, 32, stride=1, expand_ratio=6),
|
||||
)
|
||||
self.stage_5 = nn.Sequential(
|
||||
InvertedResidual(32, 48, stride=2, expand_ratio=6),
|
||||
InvertedResidual(48, 48, stride=1, expand_ratio=6),
|
||||
InvertedResidual(48, 48, stride=1, expand_ratio=6),
|
||||
InvertedResidual(48, 48, stride=1, expand_ratio=6)
|
||||
)
|
||||
self.up_to_4 = UpSampleBlock(48 + 32, 16)
|
||||
self.up_to_3 = UpSampleBlock(16 + 24, 16)
|
||||
self.up_to_2 = UpSampleBlock(16 + 16, 16)
|
||||
self.up_to_1 = UpSampleBlock(16 + 16, 16)
|
||||
self.last_layer = nn.Sequential(
|
||||
ConvBNReLU(16, 16, kernel_size=1, stride=1),
|
||||
nn.Conv2d(16, self.class_nums, kernel_size=1, stride=1)
|
||||
)
|
||||
self._initialize_weights()
|
||||
|
||||
def forward(self, x):
|
||||
feature_S = []
|
||||
x1 = self.stage_1(x)
|
||||
x2 = self.stage_2(x1)
|
||||
x3 = self.stage_3(x2)
|
||||
x4 = self.stage_4(x3)
|
||||
feature = self.stage_5(x4)
|
||||
|
||||
feature = F.interpolate(feature, size=x4.size()[2:], mode='bilinear', align_corners=True)
|
||||
feature = self.up_to_4(x4, feature)
|
||||
feature = F.interpolate(feature, size=x3.size()[2:], mode='bilinear', align_corners=True)
|
||||
feature = self.up_to_3(x3, feature)
|
||||
feature = F.interpolate(feature, size=x2.size()[2:], mode='bilinear', align_corners=True)
|
||||
feature = self.up_to_2(x2, feature)
|
||||
feature = F.interpolate(feature, size=x1.size()[2:], mode='bilinear', align_corners=True)
|
||||
feature = self.up_to_1(x1, feature)
|
||||
feature_S.append(feature)
|
||||
feature = self.last_layer(feature)
|
||||
feature_S.append(feature)
|
||||
output = F.interpolate(feature, size=x.size()[2:], mode='bilinear', align_corners=True)
|
||||
# output = torch.sigmoid(output)
|
||||
if self.output_onnx:
|
||||
output = torch.argmax(output, dim=1).to(torch.float32)
|
||||
|
||||
return output, feature_S
|
||||
|
||||
def _initialize_weights(self):
|
||||
for name, m in self.named_modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
if 'first' in name:
|
||||
nn.init.normal_(m.weight, 0, 0.01)
|
||||
else:
|
||||
nn.init.normal_(m.weight, 0, 1.0 / m.weight.shape[1])
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
nn.init.constant_(m.weight, 1)
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0.0001)
|
||||
nn.init.constant_(m.running_mean, 0)
|
||||
elif isinstance(m, nn.BatchNorm1d):
|
||||
nn.init.constant_(m.weight, 1)
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0.0001)
|
||||
nn.init.constant_(m.running_mean, 0)
|
||||
elif isinstance(m, nn.Linear):
|
||||
nn.init.normal_(m.weight, 0, 0.01)
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
class _ASPPModule(nn.Module):
|
||||
def __init__(self, inplanes, planes, kernel_size, padding, dilation, BatchNorm):
|
||||
super(_ASPPModule, self).__init__()
|
||||
self.atrous_conv = nn.Conv2d(inplanes, planes, kernel_size=kernel_size,
|
||||
stride=1, padding=padding, dilation=dilation, bias=False)
|
||||
self.bn = BatchNorm(planes)
|
||||
self.relu = nn.ReLU()
|
||||
|
||||
self._init_weight()
|
||||
|
||||
def forward(self, x):
|
||||
x = self.atrous_conv(x)
|
||||
x = self.bn(x)
|
||||
|
||||
return self.relu(x)
|
||||
|
||||
def _init_weight(self):
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
torch.nn.init.kaiming_normal_(m.weight)
|
||||
|
||||
class ASPP(nn.Module):
|
||||
def __init__(self, backbone, output_stride, BatchNorm):
|
||||
super(ASPP, self).__init__()
|
||||
if backbone == 'drn':
|
||||
inplanes = 512
|
||||
elif backbone == 'mobilenet':
|
||||
inplanes = 320
|
||||
elif backbone == 'resnet_fpn':
|
||||
inplanes = 256
|
||||
else:
|
||||
inplanes = 2048
|
||||
if output_stride == 16:
|
||||
dilations = [1, 6, 12, 18]
|
||||
elif output_stride == 8:
|
||||
dilations = [1, 12, 24, 36]
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
self.aspp1 = _ASPPModule(inplanes, 256, 1, padding=0, dilation=dilations[0], BatchNorm=BatchNorm)
|
||||
self.aspp2 = _ASPPModule(inplanes, 256, 3, padding=dilations[1], dilation=dilations[1], BatchNorm=BatchNorm)
|
||||
self.aspp3 = _ASPPModule(inplanes, 256, 3, padding=dilations[2], dilation=dilations[2], BatchNorm=BatchNorm)
|
||||
self.aspp4 = _ASPPModule(inplanes, 256, 3, padding=dilations[3], dilation=dilations[3], BatchNorm=BatchNorm)
|
||||
|
||||
self.global_avg_pool = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)),
|
||||
nn.Conv2d(inplanes, 256, 1, stride=1, bias=False),
|
||||
BatchNorm(256),
|
||||
nn.ReLU())
|
||||
self.conv1 = nn.Conv2d(1280, 256, 1, bias=False)
|
||||
self.bn1 = BatchNorm(256)
|
||||
self.relu = nn.ReLU()
|
||||
self.dropout = nn.Dropout(0.5)
|
||||
self._init_weight()
|
||||
|
||||
def forward(self, x):
|
||||
x1 = self.aspp1(x)
|
||||
x2 = self.aspp2(x)
|
||||
x3 = self.aspp3(x)
|
||||
x4 = self.aspp4(x)
|
||||
x5 = self.global_avg_pool(x)
|
||||
x5 = F.interpolate(x5, size=x4.size()[2:], mode='bilinear', align_corners=True)
|
||||
# x5 = F.interpolate(x5, size=x4.size()[2:], mode='nearest')
|
||||
x = torch.cat((x1, x2, x3, x4, x5), dim=1)
|
||||
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.relu(x)
|
||||
|
||||
return self.dropout(x)
|
||||
|
||||
def _init_weight(self):
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
torch.nn.init.kaiming_normal_(m.weight)
|
||||
|
||||
class Decoder(nn.Module):
|
||||
def __init__(self,num_classes, backbone, BatchNorm):
|
||||
super(Decoder, self).__init__()
|
||||
if backbone == 'resnet_fpn' or backbone == 'drn':
|
||||
low_level_inplanes = 256
|
||||
elif backbone == 'xception':
|
||||
low_level_inplanes = 128
|
||||
elif backbone == 'mobilenet':
|
||||
low_level_inplanes = 24
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
self.conv1 = nn.Conv2d(low_level_inplanes, 16, 1, bias=False)
|
||||
self.bn1 = BatchNorm(16)
|
||||
self.relu = nn.ReLU()
|
||||
self.last_conv = nn.Sequential(nn.Conv2d(304, 256, kernel_size=3, stride=1, padding=1, bias=False),
|
||||
BatchNorm(256),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.5),
|
||||
nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1, bias=False),
|
||||
BatchNorm(256),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.1),
|
||||
nn.Conv2d(256, num_classes, kernel_size=1, stride=1))
|
||||
|
||||
self.up_to_1 = UpSampleBlock(16 + 16, 16)
|
||||
self.last_layer = nn.Sequential(
|
||||
ConvBNReLU(16, 16, kernel_size=1, stride=1),
|
||||
nn.Conv2d(16, 1, kernel_size=1, stride=1)
|
||||
)
|
||||
self._init_weight()
|
||||
|
||||
# x(1,256,8,6) low(1,256,32,24) -》 x(1,1,32,24)
|
||||
def forward(self, x, low_level_feat):
|
||||
feature_T = []
|
||||
# deeplab part
|
||||
#(1,256,32,24) -> (1,16,64,48)
|
||||
low_level_feat = self.conv1(low_level_feat)
|
||||
low_level_feat = self.bn1(low_level_feat)
|
||||
low_level_feat = self.relu(low_level_feat)
|
||||
|
||||
# x(1, 256, 8, 6)-> (1,16,32,24)
|
||||
x = F.interpolate(x, size=low_level_feat.size()[2:], mode='bilinear', align_corners=True)
|
||||
low_level_feat = F.interpolate(low_level_feat,
|
||||
size=[low_level_feat.size()[2] * 2, low_level_feat.size()[3] * 2],
|
||||
mode='bilinear', align_corners=True)
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.relu(x)
|
||||
|
||||
# bodyseg part
|
||||
feature = F.interpolate(x, size=[x.size()[2] * 2,x.size()[3] * 2], mode='bilinear', align_corners=True)
|
||||
# 输入up_to_1 x(low)(1,16,64,48),上采样2倍后的feature(1,16,64,48)
|
||||
feature = self.up_to_1(low_level_feat, feature)
|
||||
feature_T.append(feature)
|
||||
# (1,1,64,48)
|
||||
feature = self.last_layer(feature)
|
||||
feature_T.append(feature)
|
||||
# (1,1,128,96)
|
||||
output = F.interpolate(feature, size=[128,96], mode='bilinear', align_corners=True)
|
||||
# output = torch.sigmoid(output)
|
||||
|
||||
return output, feature_T
|
||||
|
||||
def _init_weight(self):
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
torch.nn.init.kaiming_normal_(m.weight)
|
||||
|
||||
def build_backbone(backbone, output_stride, BatchNorm, input_channel=3):
|
||||
from bodyseg.backbone.fpn import build_resnet_fpn_backbone
|
||||
from bodyseg.backbone.xception import AlignedXception
|
||||
|
||||
return build_resnet_fpn_backbone(input_channel)
|
||||
|
||||
def build_aspp(backbone, output_stride, BatchNorm):
|
||||
return ASPP(backbone, output_stride, BatchNorm)
|
||||
|
||||
def build_decoder(num_classes, backbone, BatchNorm):
|
||||
return Decoder(num_classes, backbone, BatchNorm)
|
||||
|
||||
class DeepLab(nn.Module):
|
||||
def __init__(self, input_channel=3, class_num=1):
|
||||
super(DeepLab, self).__init__()
|
||||
|
||||
BatchNorm = get_norm("BN")
|
||||
|
||||
self.backbone = build_backbone("resnet_fpn", 16, BatchNorm, input_channel=input_channel)
|
||||
self.aspp = build_aspp("resnet_fpn", 16, BatchNorm)
|
||||
self.decoder = build_decoder(class_num, "resnet_fpn", BatchNorm)
|
||||
|
||||
def forward(self, input):
|
||||
# input(1,3,128,96)
|
||||
#output: "p2"(1,256,32,24), "p3"(1,256,16,12), "p4"(1,256,8,6)
|
||||
output = self.backbone(input)
|
||||
# "p4"(1,256,8,6) "p2"(1,256,32,24)
|
||||
x, low_level_feat = output['p4'], output['p2']
|
||||
# x(1,256,8,6)-》(1,256,8,6)
|
||||
x = self.aspp(x)
|
||||
# x(1,256,8,6) low(1,256,32,24) -》 x(1,1,32,24)
|
||||
x, feature_T = self.decoder(x, low_level_feat)
|
||||
# x(1,1,128,96)
|
||||
x = F.interpolate(x, size=input.size()[2:], mode='bilinear', align_corners=True)
|
||||
# x = F.interpolate(x, size=input.size()[2:], mode='nearest')
|
||||
return x, feature_T
|
||||
Reference in New Issue
Block a user