初始化:换发型/换发色/训练发型服务
包含: - hair_service_sd: 主服务(换发型/换发色/生发,端口8801) - photo_service: LoRA调度+训练(端口32678) - hair_grow_service: 调试测试页(端口8888,含4个测试页) - 批量训练脚本(batch_train_hairstyles.py) - 发际线mask自动识别(hairline_mask.py,4种方案) - 手绘mask换发型(hair_swap_manual.py) - 文档:README.md + LARGE_FILES.md + docs/ 大文件(模型权重200G、训练数据123G)已排除,见 LARGE_FILES.md OSS/COS密钥已脱敏为环境变量,原文件备份在本地
This commit is contained in:
@@ -0,0 +1,127 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torchvision.models.detection.backbone_utils as backbone_utils
|
||||
import torchvision.models._utils as _utils
|
||||
import torch.nn.functional as F
|
||||
from collections import OrderedDict
|
||||
|
||||
from models.net import MobileNetV1 as MobileNetV1
|
||||
from models.net import FPN as FPN
|
||||
from models.net import SSH as SSH
|
||||
|
||||
|
||||
|
||||
class ClassHead(nn.Module):
|
||||
def __init__(self,inchannels=512,num_anchors=3):
|
||||
super(ClassHead,self).__init__()
|
||||
self.num_anchors = num_anchors
|
||||
self.conv1x1 = nn.Conv2d(inchannels,self.num_anchors*2,kernel_size=(1,1),stride=1,padding=0)
|
||||
|
||||
def forward(self,x):
|
||||
out = self.conv1x1(x)
|
||||
out = out.permute(0,2,3,1).contiguous()
|
||||
|
||||
return out.view(out.shape[0], -1, 2)
|
||||
|
||||
class BboxHead(nn.Module):
|
||||
def __init__(self,inchannels=512,num_anchors=3):
|
||||
super(BboxHead,self).__init__()
|
||||
self.conv1x1 = nn.Conv2d(inchannels,num_anchors*4,kernel_size=(1,1),stride=1,padding=0)
|
||||
|
||||
def forward(self,x):
|
||||
out = self.conv1x1(x)
|
||||
out = out.permute(0,2,3,1).contiguous()
|
||||
|
||||
return out.view(out.shape[0], -1, 4)
|
||||
|
||||
class LandmarkHead(nn.Module):
|
||||
def __init__(self,inchannels=512,num_anchors=3):
|
||||
super(LandmarkHead,self).__init__()
|
||||
self.conv1x1 = nn.Conv2d(inchannels,num_anchors*10,kernel_size=(1,1),stride=1,padding=0)
|
||||
|
||||
def forward(self,x):
|
||||
out = self.conv1x1(x)
|
||||
out = out.permute(0,2,3,1).contiguous()
|
||||
|
||||
return out.view(out.shape[0], -1, 10)
|
||||
|
||||
class RetinaFace(nn.Module):
|
||||
def __init__(self, cfg = None, phase = 'train'):
|
||||
"""
|
||||
:param cfg: Network related settings.
|
||||
:param phase: train or test.
|
||||
"""
|
||||
super(RetinaFace,self).__init__()
|
||||
self.phase = phase
|
||||
backbone = None
|
||||
if cfg['name'] == 'mobilenet0.25':
|
||||
backbone = MobileNetV1()
|
||||
if cfg['pretrain']:
|
||||
checkpoint = torch.load("./weights/mobilenetV1X0.25_pretrain.tar", map_location=torch.device('cpu'))
|
||||
from collections import OrderedDict
|
||||
new_state_dict = OrderedDict()
|
||||
for k, v in checkpoint['state_dict'].items():
|
||||
name = k[7:] # remove module.
|
||||
new_state_dict[name] = v
|
||||
# load params
|
||||
backbone.load_state_dict(new_state_dict)
|
||||
elif cfg['name'] == 'Resnet50':
|
||||
import torchvision.models as models
|
||||
backbone = models.resnet50(pretrained=cfg['pretrain'])
|
||||
|
||||
self.body = _utils.IntermediateLayerGetter(backbone, cfg['return_layers'])
|
||||
in_channels_stage2 = cfg['in_channel']
|
||||
in_channels_list = [
|
||||
in_channels_stage2 * 2,
|
||||
in_channels_stage2 * 4,
|
||||
in_channels_stage2 * 8,
|
||||
]
|
||||
out_channels = cfg['out_channel']
|
||||
self.fpn = FPN(in_channels_list,out_channels)
|
||||
self.ssh1 = SSH(out_channels, out_channels)
|
||||
self.ssh2 = SSH(out_channels, out_channels)
|
||||
self.ssh3 = SSH(out_channels, out_channels)
|
||||
|
||||
self.ClassHead = self._make_class_head(fpn_num=3, inchannels=cfg['out_channel'])
|
||||
self.BboxHead = self._make_bbox_head(fpn_num=3, inchannels=cfg['out_channel'])
|
||||
self.LandmarkHead = self._make_landmark_head(fpn_num=3, inchannels=cfg['out_channel'])
|
||||
|
||||
def _make_class_head(self,fpn_num=3,inchannels=64,anchor_num=2):
|
||||
classhead = nn.ModuleList()
|
||||
for i in range(fpn_num):
|
||||
classhead.append(ClassHead(inchannels,anchor_num))
|
||||
return classhead
|
||||
|
||||
def _make_bbox_head(self,fpn_num=3,inchannels=64,anchor_num=2):
|
||||
bboxhead = nn.ModuleList()
|
||||
for i in range(fpn_num):
|
||||
bboxhead.append(BboxHead(inchannels,anchor_num))
|
||||
return bboxhead
|
||||
|
||||
def _make_landmark_head(self,fpn_num=3,inchannels=64,anchor_num=2):
|
||||
landmarkhead = nn.ModuleList()
|
||||
for i in range(fpn_num):
|
||||
landmarkhead.append(LandmarkHead(inchannels,anchor_num))
|
||||
return landmarkhead
|
||||
|
||||
def forward(self,inputs):
|
||||
out = self.body(inputs)
|
||||
|
||||
# FPN
|
||||
fpn = self.fpn(out)
|
||||
|
||||
# SSH
|
||||
feature1 = self.ssh1(fpn[0])
|
||||
feature2 = self.ssh2(fpn[1])
|
||||
feature3 = self.ssh3(fpn[2])
|
||||
features = [feature1, feature2, feature3]
|
||||
|
||||
bbox_regressions = torch.cat([self.BboxHead[i](feature) for i, feature in enumerate(features)], dim=1)
|
||||
classifications = torch.cat([self.ClassHead[i](feature) for i, feature in enumerate(features)],dim=1)
|
||||
ldm_regressions = torch.cat([self.LandmarkHead[i](feature) for i, feature in enumerate(features)], dim=1)
|
||||
|
||||
if self.phase == 'train':
|
||||
output = (bbox_regressions, classifications, ldm_regressions)
|
||||
else:
|
||||
output = (bbox_regressions, F.softmax(classifications, dim=-1), ldm_regressions)
|
||||
return output
|
||||
Reference in New Issue
Block a user