初始化换发型项目：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,209 @@
|
||||
import os
|
||||
import cv2
|
||||
import argparse
|
||||
import numpy as np
|
||||
|
||||
import torch
|
||||
from torch.nn import functional as F
|
||||
|
||||
import utils
|
||||
from matting import networks
|
||||
from utils.data_preprocess import *
|
||||
from utils import landmark_processor
|
||||
from time import time
|
||||
|
||||
def single_inference(model, image_dict, device, return_offset=False):
|
||||
|
||||
with torch.no_grad():
|
||||
image, trimap = image_dict['image'], image_dict['trimap']
|
||||
alpha_shape = image_dict['alpha_shape']
|
||||
image = image.to(device)
|
||||
trimap = trimap.to(device)
|
||||
|
||||
matte_start = time()
|
||||
alpha_pred, info_dict = model(image, trimap)
|
||||
matte_end = time()
|
||||
# print("matte time cost : {:.4f}".format(matte_end-matte_start))
|
||||
|
||||
trimap_argmax = trimap.argmax(dim=1, keepdim=True)
|
||||
|
||||
alpha_pred[trimap_argmax == 2] = 1
|
||||
alpha_pred[trimap_argmax == 0] = 0
|
||||
|
||||
h, w = alpha_shape
|
||||
test_pred = alpha_pred[0, 0, ...].detach().cpu().numpy() * 255
|
||||
test_pred = test_pred.astype(np.uint8)
|
||||
test_pred = test_pred[32:h+32, 32:w+32]
|
||||
|
||||
if return_offset:
|
||||
short_side = h if h < w else w
|
||||
ratio = 512 / short_side
|
||||
offset_1 = utils.flow_to_image(info_dict['offset_1'][0][0, ...].data.cpu().numpy()).astype(np.uint8)
|
||||
# write softmax_scale to offset image
|
||||
scale = info_dict['offset_1'][1].cpu()
|
||||
offset_1 = cv2.resize(offset_1, (int(w * ratio), int(h * ratio)), interpolation=cv2.INTER_NEAREST)
|
||||
text = 'unknown: {:.2f}, known: {:.2f}'.format(scale[-1, 0].item(), scale[-1,1].item())
|
||||
offset_1 = cv2.putText(offset_1, text, (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.8, 0, thickness=2)
|
||||
|
||||
offset_2 = utils.flow_to_image(info_dict['offset_2'][0][0, ...].data.cpu().numpy()).astype(np.uint8)
|
||||
# write softmax_scale to offset image
|
||||
scale = info_dict['offset_2'][1].cpu()
|
||||
offset_2 = cv2.resize(offset_2, (int(w * ratio), int(h * ratio)), interpolation=cv2.INTER_NEAREST)
|
||||
text = 'unknown: {:.2f}, known: {:.2f}'.format(scale[-1, 0].item(), scale[-1,1].item())
|
||||
offset_2 = cv2.putText(offset_2, text, (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.8, 0, thickness=2)
|
||||
|
||||
return test_pred, (offset_1, offset_2)
|
||||
else:
|
||||
return test_pred, None
|
||||
|
||||
def generator_tensor_dict(image, trimap):
|
||||
|
||||
sample = {'image': image, 'trimap': trimap, 'alpha_shape': trimap.shape}
|
||||
|
||||
# reshape
|
||||
h, w = sample["alpha_shape"]
|
||||
|
||||
if h % 32 == 0 and w % 32 == 0:
|
||||
padded_image = np.pad(sample['image'], ((32, 32), (32, 32), (0, 0)), mode="reflect")
|
||||
padded_trimap = np.pad(sample['trimap'], ((32, 32), (32, 32)), mode="reflect")
|
||||
sample['image'] = padded_image
|
||||
sample['trimap'] = padded_trimap
|
||||
else:
|
||||
target_h = 32 * ((h - 1) // 32 + 1)
|
||||
target_w = 32 * ((w - 1) // 32 + 1)
|
||||
pad_h = target_h - h
|
||||
pad_w = target_w - w
|
||||
padded_image = np.pad(sample['image'], ((32, pad_h+32), (32, pad_w+32), (0, 0)), mode="reflect")
|
||||
padded_trimap = np.pad(sample['trimap'], ((32, pad_h+32), (32, pad_w+32)), mode="reflect")
|
||||
sample['image'] = padded_image
|
||||
sample['trimap'] = padded_trimap
|
||||
|
||||
# ImageNet mean & std
|
||||
mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
|
||||
std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
|
||||
# convert GBR images to RGB
|
||||
image, trimap = sample['image'][:, :, ::-1], sample['trimap']
|
||||
# swap color axis
|
||||
image = image.transpose((2, 0, 1)).astype(np.float32)
|
||||
trimap[trimap < 85] = 0
|
||||
trimap[trimap >= 170] = 2
|
||||
trimap[trimap >= 85] = 1
|
||||
# normalize image
|
||||
image /= 255.
|
||||
|
||||
# to tensor
|
||||
sample['image'], sample['trimap'] = torch.from_numpy(image), torch.from_numpy(trimap).to(torch.long)
|
||||
sample['image'] = sample['image'].sub_(mean).div_(std)
|
||||
|
||||
# if CONFIG.model.trimap_channel == 3:
|
||||
sample['trimap'] = F.one_hot(sample['trimap'], num_classes=3).permute(2, 0, 1).float()
|
||||
# elif CONFIG.model.trimap_channel == 1:
|
||||
# sample['trimap'] = sample['trimap'][None, ...].float()
|
||||
# else:
|
||||
# raise NotImplementedError("CONFIG.model.trimap_channel can only be 3 or 1")
|
||||
|
||||
# add first channel
|
||||
sample['image'], sample['trimap'] = sample['image'][None, ...], sample['trimap'][None, ...]
|
||||
|
||||
return sample
|
||||
|
||||
if __name__ == '__main__':
|
||||
print('Torch Version: ', torch.__version__)
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--checkpoint', type=str, default='/home/liyang/project/matting/GCA-Matting/checkpoints/gca-dist-align-truth-1101/best_model.pth',
|
||||
help="path of checkpoint")
|
||||
parser.add_argument('--image-dir', type=str, default='/home/liyang/project/matting/合格/origin', help="input image dir")
|
||||
parser.add_argument('--trimap-dir', type=str, default='/home/liyang/project/matting/合格/trimap', help="input trimap dir")
|
||||
parser.add_argument('--output', type=str, default='/home/liyang/project/matting/合格/pred6', help="output dir")
|
||||
|
||||
# Parse configuration
|
||||
args = parser.parse_args()
|
||||
|
||||
# # Check if toml config file is loaded
|
||||
# if CONFIG.is_default:
|
||||
# raise ValueError("No .toml config loaded.")
|
||||
|
||||
args.output = os.path.join(args.output, args.checkpoint.split('/')[-1])
|
||||
utils.make_dir(args.output)
|
||||
|
||||
# build model
|
||||
model = networks.get_generator(encoder="resnet_gca_encoder_29", decoder="res_gca_decoder_22", num_class=1)
|
||||
model.cuda()
|
||||
print("model: ", model)
|
||||
|
||||
# load checkpoint
|
||||
checkpoint = torch.load(args.checkpoint)
|
||||
model.load_state_dict(utils.remove_prefix_state_dict(checkpoint['state_dict']), strict=True)
|
||||
|
||||
# inference
|
||||
model = model.eval()
|
||||
|
||||
for image_name in os.listdir(args.image_dir):
|
||||
|
||||
if not is_image_file(image_name):
|
||||
continue
|
||||
# assume image and trimap have the same file name
|
||||
img_basename, ext = os.path.splitext(image_name)
|
||||
|
||||
image_path = os.path.join(args.image_dir, image_name)
|
||||
trimap_path = os.path.join(args.trimap_dir, image_name.replace(ext, "_alpha.png"))
|
||||
# trimap_path = os.path.join(args.trimap_dir, image_name)
|
||||
print('Image: ', image_path, ' Tirmap: ', trimap_path)
|
||||
|
||||
# read images
|
||||
img_basename, img_ext = os.path.splitext(image_name)
|
||||
# img_pt_path = image_path.replace(img_ext, "_landmark1k.txt")
|
||||
# img_landmark1k = np.loadtxt(img_pt_path)
|
||||
|
||||
image = cv2.imread(image_path)
|
||||
trimap = cv2.imread(trimap_path, 0)
|
||||
|
||||
ori_h, ori_w, _ = image.shape
|
||||
tri_h, tri_w = trimap.shape
|
||||
if image.shape[1] != trimap.shape[1] or image.shape[0] != trimap.shape[0]:
|
||||
image = cv2.resize(image, (trimap.shape[1], trimap.shape[0]), interpolation=cv2.INTER_CUBIC)
|
||||
# img_landmark1k = landmark_processor.resize_points(img_landmark1k, ori_w, ori_h,
|
||||
# tri_w, tri_h)
|
||||
|
||||
# hair_mat = landmark_processor.get_transform_mat_hair(img_landmark1k, 640, ratio=0.3, w_ratio=0.5,
|
||||
# h_ratio=0.4)
|
||||
# hair_img_landmark = landmark_processor.transform_points(img_landmark1k, hair_mat)
|
||||
#
|
||||
# image = cv2.warpAffine(image, hair_mat, (640, 640), flags=cv2.INTER_CUBIC)
|
||||
# trimap = cv2.warpAffine(trimap, hair_mat, (640, 640), flags=cv2.INTER_CUBIC)
|
||||
|
||||
if tri_h > 1920 or tri_w > 1920:
|
||||
if tri_h > 1920:
|
||||
new_tri_h = 1920
|
||||
new_tri_w = int(tri_w * 1920 / tri_h)
|
||||
else:
|
||||
new_tri_w = 1920
|
||||
new_tri_h = int(tri_h * 1920 / tri_w)
|
||||
|
||||
image_resize = cv2.resize(image, (new_tri_w, new_tri_h), interpolation=cv2.INTER_CUBIC)
|
||||
trimap_resize = cv2.resize(trimap, (new_tri_w, new_tri_h), interpolation=cv2.INTER_NEAREST)
|
||||
else:
|
||||
image_resize = image
|
||||
trimap_resize = trimap
|
||||
|
||||
# image_dict = generator_tensor_dict(image, trimap)
|
||||
image_dict = generator_tensor_dict(image_resize, trimap_resize)
|
||||
pred, offset = single_inference(model, image_dict)
|
||||
|
||||
# torch.cuda.empty_cache()
|
||||
|
||||
pred = cv2.resize(pred, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_CUBIC)
|
||||
# pred = cv2.warpAffine(pred, cv2.invertAffineTransform(hair_mat), (tri_w, tri_h), flags=cv2.INTER_CUBIC)
|
||||
|
||||
# offset[0] = cv2.resize(offset[0], (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST)
|
||||
# offset[1] = cv2.resize(offset[1], (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST)
|
||||
|
||||
# cv2.imshow("image_resize", image_resize)
|
||||
# cv2.imshow("trimap_resize", trimap_resize)
|
||||
# cv2.imshow("pred", pred)
|
||||
# cv2.waitKey()
|
||||
cv2.imwrite(os.path.join(args.output, image_name.replace(img_ext, "_noalign.png")), pred)
|
||||
# if offset is not None:
|
||||
# cv2.imwrite(os.path.join(args.output, os.path.splitext(image_name)[0]+'_offset1.png'), offset[0])
|
||||
# cv2.imwrite(os.path.join(args.output, os.path.splitext(image_name)[0]+'_offset2.png'), offset[1])
|
||||
@@ -0,0 +1,197 @@
|
||||
import os
|
||||
import cv2
|
||||
import argparse
|
||||
import numpy as np
|
||||
|
||||
import torch
|
||||
from torch.nn import functional as F
|
||||
|
||||
import utils
|
||||
from hair_matting.matting import networks
|
||||
from utils.data_preprocess import *
|
||||
from utils import landmark_processor
|
||||
from time import time
|
||||
|
||||
def single_inference(model, image_dict, device, return_offset=False):
|
||||
|
||||
with torch.no_grad():
|
||||
image, trimap = image_dict['image'], image_dict['trimap']
|
||||
alpha_shape = image_dict['alpha_shape']
|
||||
image = image.to(device)
|
||||
trimap = trimap.to(device)
|
||||
|
||||
matte_start = time()
|
||||
|
||||
alpha_pred, info_dict = model(image, trimap)
|
||||
|
||||
fg_pred = alpha_pred[:, :-1, :, :]
|
||||
alpha_pred = alpha_pred[:, -1, :, :].unsqueeze(1)
|
||||
matte_end = time()
|
||||
print("matte time cost : {:.4f}".format(matte_end-matte_start))
|
||||
|
||||
trimap_argmax = trimap.argmax(dim=1, keepdim=True)
|
||||
|
||||
alpha_pred[trimap_argmax == 2] = 1
|
||||
alpha_pred[trimap_argmax == 0] = 0
|
||||
|
||||
h, w = alpha_shape
|
||||
test_fg_pred = fg_pred[0].detach().cpu().numpy().transpose(1, 2, 0)[:, :, ::-1] * 255
|
||||
test_fg_pred = test_fg_pred.astype(np.uint8)
|
||||
test_fg_pred = test_fg_pred[32:h+32, 32:w+32]
|
||||
|
||||
test_pred = alpha_pred[0, 0, ...].detach().cpu().numpy() * 255
|
||||
test_pred = test_pred.astype(np.uint8)
|
||||
test_pred = test_pred[32:h+32, 32:w+32]
|
||||
|
||||
if return_offset:
|
||||
short_side = h if h < w else w
|
||||
ratio = 512 / short_side
|
||||
offset_1 = utils.flow_to_image(info_dict['offset_1'][0][0, ...].data.cpu().numpy()).astype(np.uint8)
|
||||
# write softmax_scale to offset image
|
||||
scale = info_dict['offset_1'][1].cpu()
|
||||
offset_1 = cv2.resize(offset_1, (int(w * ratio), int(h * ratio)), interpolation=cv2.INTER_NEAREST)
|
||||
text = 'unknown: {:.2f}, known: {:.2f}'.format(scale[-1, 0].item(), scale[-1,1].item())
|
||||
offset_1 = cv2.putText(offset_1, text, (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.8, 0, thickness=2)
|
||||
|
||||
offset_2 = utils.flow_to_image(info_dict['offset_2'][0][0, ...].data.cpu().numpy()).astype(np.uint8)
|
||||
# write softmax_scale to offset image
|
||||
scale = info_dict['offset_2'][1].cpu()
|
||||
offset_2 = cv2.resize(offset_2, (int(w * ratio), int(h * ratio)), interpolation=cv2.INTER_NEAREST)
|
||||
text = 'unknown: {:.2f}, known: {:.2f}'.format(scale[-1, 0].item(), scale[-1,1].item())
|
||||
offset_2 = cv2.putText(offset_2, text, (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.8, 0, thickness=2)
|
||||
|
||||
return test_fg_pred, test_pred, (offset_1, offset_2)
|
||||
else:
|
||||
return test_fg_pred, test_pred, None
|
||||
|
||||
def generator_tensor_dict(image, trimap):
|
||||
|
||||
sample = {'image': image, 'trimap': trimap, 'alpha_shape': trimap.shape}
|
||||
|
||||
# reshape
|
||||
h, w = sample["alpha_shape"]
|
||||
|
||||
if h % 32 == 0 and w % 32 == 0:
|
||||
padded_image = np.pad(sample['image'], ((32, 32), (32, 32), (0, 0)), mode="reflect")
|
||||
padded_trimap = np.pad(sample['trimap'], ((32, 32), (32, 32)), mode="reflect")
|
||||
sample['image'] = padded_image
|
||||
sample['trimap'] = padded_trimap
|
||||
else:
|
||||
target_h = 32 * ((h - 1) // 32 + 1)
|
||||
target_w = 32 * ((w - 1) // 32 + 1)
|
||||
pad_h = target_h - h
|
||||
pad_w = target_w - w
|
||||
padded_image = np.pad(sample['image'], ((32, pad_h+32), (32, pad_w+32), (0, 0)), mode="reflect")
|
||||
padded_trimap = np.pad(sample['trimap'], ((32, pad_h+32), (32, pad_w+32)), mode="reflect")
|
||||
sample['image'] = padded_image
|
||||
sample['trimap'] = padded_trimap
|
||||
|
||||
# ImageNet mean & std
|
||||
mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
|
||||
std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
|
||||
# convert GBR images to RGB
|
||||
image, trimap = sample['image'][:, :, ::-1], sample['trimap']
|
||||
# swap color axis
|
||||
image = image.transpose((2, 0, 1)).astype(np.float32)
|
||||
trimap[trimap < 85] = 0
|
||||
trimap[trimap >= 170] = 2
|
||||
trimap[trimap >= 85] = 1
|
||||
# normalize image
|
||||
image /= 255.
|
||||
|
||||
# to tensor
|
||||
sample['image'], sample['trimap'] = torch.from_numpy(image), torch.from_numpy(trimap).to(torch.long)
|
||||
sample['image'] = sample['image'].sub_(mean).div_(std)
|
||||
|
||||
sample['trimap'] = F.one_hot(sample['trimap'], num_classes=3).permute(2, 0, 1).float()
|
||||
|
||||
# add first channel
|
||||
sample['image'], sample['trimap'] = sample['image'][None, ...], sample['trimap'][None, ...]
|
||||
|
||||
return sample
|
||||
|
||||
if __name__ == '__main__':
|
||||
print('Torch Version: ', torch.__version__)
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--checkpoint', type=str, default='/home/liyang/project/matting/GCA-Matting/checkpoints/gca-dist-align-truth-1101/best_model.pth',
|
||||
help="path of checkpoint")
|
||||
parser.add_argument('--image-dir', type=str, default='/home/liyang/project/matting/合格/origin', help="input image dir")
|
||||
parser.add_argument('--trimap-dir', type=str, default='/home/liyang/project/matting/合格/trimap', help="input trimap dir")
|
||||
parser.add_argument('--output', type=str, default='/home/liyang/project/matting/合格/pred6', help="output dir")
|
||||
|
||||
# Parse configuration
|
||||
args = parser.parse_args()
|
||||
|
||||
# # Check if toml config file is loaded
|
||||
# if CONFIG.is_default:
|
||||
# raise ValueError("No .toml config loaded.")
|
||||
|
||||
args.output = os.path.join(args.output, args.checkpoint.split('/')[-1])
|
||||
utils.make_dir(args.output)
|
||||
|
||||
# build model
|
||||
model = networks.get_generator(encoder=CONFIG.model.arch.encoder, decoder=CONFIG.model.arch.decoder)
|
||||
model.cuda()
|
||||
print("model: ", model)
|
||||
|
||||
# load checkpoint
|
||||
checkpoint = torch.load(args.checkpoint)
|
||||
model.load_state_dict(utils.remove_prefix_state_dict(checkpoint['state_dict']), strict=True)
|
||||
|
||||
# inference
|
||||
model = model.eval()
|
||||
export_onnx_file = "test.onnx"
|
||||
torch.onnx.export(model, x, export_onnx_file, opset_version=10, do_constant_folding=True, input_names=["image", "trimap"], # 输入名
|
||||
output_names=["fg_pred", "alpha_pred", "None"])
|
||||
|
||||
for image_name in os.listdir(args.image_dir):
|
||||
|
||||
if not is_image_file(image_name):
|
||||
continue
|
||||
# assume image and trimap have the same file name
|
||||
img_basename, ext = os.path.splitext(image_name)
|
||||
|
||||
image_path = os.path.join(args.image_dir, image_name)
|
||||
trimap_path = os.path.join(args.trimap_dir, image_name.replace(ext, "_alpha.png"))
|
||||
# trimap_path = os.path.join(args.trimap_dir, image_name)
|
||||
print('Image: ', image_path, ' Tirmap: ', trimap_path)
|
||||
|
||||
# read images
|
||||
img_basename, img_ext = os.path.splitext(image_name)
|
||||
|
||||
image = cv2.imread(image_path)
|
||||
trimap = cv2.imread(trimap_path, 0)
|
||||
|
||||
ori_h, ori_w, _ = image.shape
|
||||
tri_h, tri_w = trimap.shape
|
||||
if image.shape[1] != trimap.shape[1] or image.shape[0] != trimap.shape[0]:
|
||||
image = cv2.resize(image, (trimap.shape[1], trimap.shape[0]), interpolation=cv2.INTER_CUBIC)
|
||||
|
||||
|
||||
if tri_h > 1920 or tri_w > 1920:
|
||||
if tri_h > 1920:
|
||||
new_tri_h = 1920
|
||||
new_tri_w = int(tri_w * 1920 / tri_h)
|
||||
else:
|
||||
new_tri_w = 1920
|
||||
new_tri_h = int(tri_h * 1920 / tri_w)
|
||||
|
||||
image_resize = cv2.resize(image, (new_tri_w, new_tri_h), interpolation=cv2.INTER_CUBIC)
|
||||
trimap_resize = cv2.resize(trimap, (new_tri_w, new_tri_h), interpolation=cv2.INTER_NEAREST)
|
||||
else:
|
||||
image_resize = image
|
||||
trimap_resize = trimap
|
||||
|
||||
# image_dict = generator_tensor_dict(image, trimap)
|
||||
image_dict = generator_tensor_dict(image_resize, trimap_resize)
|
||||
fg_pred, pred, offset = single_inference(model, image_dict)
|
||||
|
||||
# torch.cuda.empty_cache()
|
||||
|
||||
fg_pred = cv2.resize(fg_pred, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_CUBIC)
|
||||
pred = cv2.resize(pred, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_CUBIC)
|
||||
|
||||
cv2.imwrite(os.path.join(args.output, image_name.replace(img_ext, "_noalign.png")), pred)
|
||||
cv2.imwrite(os.path.join(args.output, image_name.replace(img_ext, "_fg_p.png")), fg_pred)
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
from .generators import *
|
||||
@@ -0,0 +1,28 @@
|
||||
from .resnet_dec import ResNet_D_Dec, BasicBlock
|
||||
from .res_shortcut_dec import ResShortCut_D_Dec
|
||||
from .res_gca_dec import ResGuidedCxtAtten_Dec
|
||||
|
||||
|
||||
__all__ = ['res_shortcut_decoder_22', 'res_gca_decoder_22']
|
||||
|
||||
|
||||
def _res_shortcut_D_dec(block, layers, **kwargs):
|
||||
model = ResShortCut_D_Dec(block, layers, **kwargs)
|
||||
return model
|
||||
|
||||
|
||||
def _res_gca_D_dec(block, layers, num_class, **kwargs):
|
||||
model = ResGuidedCxtAtten_Dec(block, layers, num_class, **kwargs)
|
||||
return model
|
||||
|
||||
|
||||
def res_shortcut_decoder_22(**kwargs):
|
||||
"""Constructs a resnet_encoder_14 model.
|
||||
"""
|
||||
return _res_shortcut_D_dec(BasicBlock, [2, 3, 3, 2], **kwargs)
|
||||
|
||||
|
||||
def res_gca_decoder_22(num_class=1, **kwargs):
|
||||
"""Constructs a resnet_encoder_14 model.
|
||||
"""
|
||||
return _res_gca_D_dec(BasicBlock, [2, 3, 3, 2], num_class, **kwargs)
|
||||
@@ -0,0 +1,28 @@
|
||||
from hair_matting.matting.networks.ops import GuidedCxtAtten, SpectralNorm
|
||||
from hair_matting.matting.networks.decoders.res_shortcut_dec import ResShortCut_D_Dec
|
||||
|
||||
|
||||
class ResGuidedCxtAtten_Dec(ResShortCut_D_Dec):
|
||||
|
||||
def __init__(self, block, layers, num_class=1, norm_layer=None, large_kernel=False):
|
||||
super(ResGuidedCxtAtten_Dec, self).__init__(block, layers, num_class, norm_layer, large_kernel)
|
||||
self.gca = GuidedCxtAtten(128, 128)
|
||||
self.num_class = num_class
|
||||
|
||||
def forward(self, x, mid_fea):
|
||||
fea1, fea2, fea3, fea4, fea5 = mid_fea['shortcut']
|
||||
im = mid_fea['image_fea']
|
||||
x = self.layer1(x) + fea5 # N x 256 x 32 x 32
|
||||
x = self.layer2(x) + fea4 # N x 128 x 64 x 64
|
||||
x, offset = self.gca(im, x, mid_fea['unknown']) # contextual attention
|
||||
x = self.layer3(x) + fea3 # N x 64 x 128 x 128
|
||||
x = self.layer4(x) + fea2 # N x 32 x 256 x 256
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.leaky_relu(x) + fea1
|
||||
x = self.conv2(x)
|
||||
|
||||
alpha = (self.tanh(x) + 1.0) / 2.0
|
||||
|
||||
return alpha, {'offset_1': mid_fea['offset_1'], 'offset_2': offset}
|
||||
|
||||
@@ -0,0 +1,24 @@
|
||||
from hair_matting.matting.networks.decoders.resnet_dec import ResNet_D_Dec
|
||||
|
||||
|
||||
class ResShortCut_D_Dec(ResNet_D_Dec):
|
||||
|
||||
def __init__(self, block, layers, num_class=1, norm_layer=None, large_kernel=False, late_downsample=False):
|
||||
super(ResShortCut_D_Dec, self).__init__(block, layers, num_class, norm_layer, large_kernel,
|
||||
late_downsample=late_downsample)
|
||||
|
||||
def forward(self, x, mid_fea):
|
||||
fea1, fea2, fea3, fea4, fea5 = mid_fea['shortcut']
|
||||
x = self.layer1(x) + fea5
|
||||
x = self.layer2(x) + fea4
|
||||
x = self.layer3(x) + fea3
|
||||
x = self.layer4(x) + fea2
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.leaky_relu(x) + fea1
|
||||
x = self.conv2(x)
|
||||
|
||||
alpha = (self.tanh(x) + 1.0) / 2.0
|
||||
|
||||
return alpha, None
|
||||
|
||||
@@ -0,0 +1,142 @@
|
||||
import logging
|
||||
import torch.nn as nn
|
||||
from hair_matting.matting.networks.ops import SpectralNorm
|
||||
|
||||
def conv5x5(in_planes, out_planes, stride=1, groups=1, dilation=1):
|
||||
"""5x5 convolution with padding"""
|
||||
return nn.Conv2d(in_planes, out_planes, kernel_size=5, stride=stride,
|
||||
padding=2, groups=groups, bias=False, dilation=dilation)
|
||||
|
||||
|
||||
def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
|
||||
"""3x3 convolution with padding"""
|
||||
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
|
||||
padding=dilation, groups=groups, bias=False, dilation=dilation)
|
||||
|
||||
|
||||
def conv1x1(in_planes, out_planes, stride=1):
|
||||
"""1x1 convolution"""
|
||||
return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
|
||||
|
||||
|
||||
class BasicBlock(nn.Module):
|
||||
expansion = 1
|
||||
|
||||
def __init__(self, inplanes, planes, stride=1, upsample=None, norm_layer=None, large_kernel=False):
|
||||
super(BasicBlock, self).__init__()
|
||||
if norm_layer is None:
|
||||
norm_layer = nn.BatchNorm2d
|
||||
self.stride = stride
|
||||
conv = conv5x5 if large_kernel else conv3x3
|
||||
# Both self.conv1 and self.downsample layers downsample the input when stride != 1
|
||||
if self.stride > 1:
|
||||
self.conv1 = SpectralNorm(nn.ConvTranspose2d(inplanes, inplanes, kernel_size=4, stride=2, padding=1, bias=False))
|
||||
else:
|
||||
self.conv1 = SpectralNorm(conv(inplanes, inplanes))
|
||||
self.bn1 = norm_layer(inplanes)
|
||||
self.activation = nn.LeakyReLU(0.2, inplace=True)
|
||||
self.conv2 = SpectralNorm(conv(inplanes, planes))
|
||||
self.bn2 = norm_layer(planes)
|
||||
self.upsample = upsample
|
||||
|
||||
def forward(self, x):
|
||||
identity = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.activation(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
|
||||
if self.upsample is not None:
|
||||
identity = self.upsample(x)
|
||||
|
||||
out += identity
|
||||
out = self.activation(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class ResNet_D_Dec(nn.Module):
|
||||
|
||||
def __init__(self, block, layers, num_class=1, norm_layer=None, large_kernel=False, late_downsample=False):
|
||||
super(ResNet_D_Dec, self).__init__()
|
||||
self.logger = logging.getLogger("Logger")
|
||||
if norm_layer is None:
|
||||
norm_layer = nn.BatchNorm2d
|
||||
self._norm_layer = norm_layer
|
||||
self.large_kernel = large_kernel
|
||||
self.kernel_size = 5 if self.large_kernel else 3
|
||||
|
||||
self.inplanes = 512 if layers[0] > 0 else 256
|
||||
self.late_downsample = late_downsample
|
||||
self.midplanes = 64 if late_downsample else 32
|
||||
|
||||
self.conv1 = SpectralNorm(nn.ConvTranspose2d(self.midplanes, 32, kernel_size=4, stride=2, padding=1, bias=False))
|
||||
self.bn1 = norm_layer(32)
|
||||
self.leaky_relu = nn.LeakyReLU(0.2, inplace=True)
|
||||
self.conv2 = nn.Conv2d(32, num_class, kernel_size=self.kernel_size, stride=1, padding=self.kernel_size//2)
|
||||
self.upsample = nn.UpsamplingNearest2d(scale_factor=2)
|
||||
self.tanh = nn.Tanh()
|
||||
self.layer1 = self._make_layer(block, 256, layers[0], stride=2)
|
||||
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
|
||||
self.layer3 = self._make_layer(block, 64, layers[2], stride=2)
|
||||
self.layer4 = self._make_layer(block, self.midplanes, layers[3], stride=2)
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
if hasattr(m, "weight_bar"):
|
||||
nn.init.xavier_uniform_(m.weight_bar)
|
||||
else:
|
||||
nn.init.xavier_uniform_(m.weight)
|
||||
elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
|
||||
nn.init.constant_(m.weight, 1)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
# Zero-initialize the last BN in each residual branch,
|
||||
# so that the residual branch starts with zeros, and each residual block behaves like an identity.
|
||||
# This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
|
||||
for m in self.modules():
|
||||
if isinstance(m, BasicBlock):
|
||||
nn.init.constant_(m.bn2.weight, 0)
|
||||
|
||||
self.logger.debug(self)
|
||||
|
||||
def _make_layer(self, block, planes, blocks, stride=1):
|
||||
if blocks == 0:
|
||||
return nn.Sequential(nn.Identity())
|
||||
norm_layer = self._norm_layer
|
||||
upsample = None
|
||||
if stride != 1:
|
||||
upsample = nn.Sequential(
|
||||
nn.UpsamplingNearest2d(scale_factor=2),
|
||||
SpectralNorm(conv1x1(self.inplanes, planes * block.expansion)),
|
||||
norm_layer(planes * block.expansion),
|
||||
)
|
||||
elif self.inplanes != planes * block.expansion:
|
||||
upsample = nn.Sequential(
|
||||
SpectralNorm(conv1x1(self.inplanes, planes * block.expansion)),
|
||||
norm_layer(planes * block.expansion),
|
||||
)
|
||||
|
||||
layers = [block(self.inplanes, planes, stride, upsample, norm_layer, self.large_kernel)]
|
||||
self.inplanes = planes * block.expansion
|
||||
for _ in range(1, blocks):
|
||||
layers.append(block(self.inplanes, planes, norm_layer=norm_layer, large_kernel=self.large_kernel))
|
||||
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x, mid_fea):
|
||||
x = self.layer1(x) # N x 256 x 32 x 32
|
||||
x = self.layer2(x) # N x 128 x 64 x 64
|
||||
x = self.layer3(x) # N x 64 x 128 x 128
|
||||
x = self.layer4(x) # N x 32 x 256 x 256
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.leaky_relu(x)
|
||||
x = self.conv2(x)
|
||||
|
||||
alpha = (self.tanh(x) + 1.0) / 2.0
|
||||
|
||||
return alpha, None
|
||||
@@ -0,0 +1,39 @@
|
||||
import logging
|
||||
from .resnet_enc import ResNet_D, BasicBlock
|
||||
from .res_shortcut_enc import ResShortCut_D
|
||||
from .res_gca_enc import ResGuidedCxtAtten
|
||||
|
||||
|
||||
__all__ = ['res_shortcut_encoder_29', 'resnet_gca_encoder_29']
|
||||
|
||||
|
||||
def _res_shortcut_D(block, layers, **kwargs):
|
||||
model = ResShortCut_D(block, layers, **kwargs)
|
||||
return model
|
||||
|
||||
|
||||
def _res_gca_D(block, layers, **kwargs):
|
||||
model = ResGuidedCxtAtten(block, layers, **kwargs)
|
||||
return model
|
||||
|
||||
|
||||
def resnet_gca_encoder_29(**kwargs):
|
||||
"""Constructs a resnet_encoder_29 model.
|
||||
"""
|
||||
return _res_gca_D(BasicBlock, [3, 4, 4, 2], **kwargs)
|
||||
|
||||
|
||||
def res_shortcut_encoder_29(**kwargs):
|
||||
"""Constructs a resnet_encoder_25 model.
|
||||
"""
|
||||
return _res_shortcut_D(BasicBlock, [3, 4, 4, 2], **kwargs)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import torch
|
||||
logging.basicConfig(level=logging.DEBUG, format='[%(asctime)s] %(levelname)s: %(message)s',
|
||||
datefmt='%m-%d %H:%M:%S')
|
||||
resnet_encoder = res_shortcut_encoder_29()
|
||||
x = torch.randn(4,6,512,512)
|
||||
z = resnet_encoder(x)
|
||||
print(z[0].shape)
|
||||
@@ -0,0 +1,97 @@
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
# from utils import CONFIG
|
||||
from hair_matting.matting.networks.encoders.resnet_enc import ResNet_D
|
||||
from hair_matting.matting.networks.ops import GuidedCxtAtten, SpectralNorm
|
||||
|
||||
class ResGuidedCxtAtten(ResNet_D):
|
||||
|
||||
def __init__(self, block, layers, norm_layer=None, late_downsample=False):
|
||||
super(ResGuidedCxtAtten, self).__init__(block, layers, norm_layer, late_downsample=late_downsample)
|
||||
first_inplane = 3 + 3
|
||||
self.shortcut_inplane = [first_inplane, self.midplanes, 64, 128, 256]
|
||||
self.shortcut_plane = [32, self.midplanes, 64, 128, 256]
|
||||
|
||||
self.shortcut = nn.ModuleList()
|
||||
for stage, inplane in enumerate(self.shortcut_inplane):
|
||||
self.shortcut.append(self._make_shortcut(inplane, self.shortcut_plane[stage]))
|
||||
|
||||
self.guidance_head = nn.Sequential(
|
||||
nn.ReflectionPad2d(1),
|
||||
SpectralNorm(nn.Conv2d(3, 16, kernel_size=3, padding=0, stride=2, bias=False)),
|
||||
nn.ReLU(inplace=True),
|
||||
self._norm_layer(16),
|
||||
nn.ReflectionPad2d(1),
|
||||
SpectralNorm(nn.Conv2d(16, 32, kernel_size=3, padding=0, stride=2, bias=False)),
|
||||
nn.ReLU(inplace=True),
|
||||
self._norm_layer(32),
|
||||
nn.ReflectionPad2d(1),
|
||||
SpectralNorm(nn.Conv2d(32, 128, kernel_size=3, padding=0, stride=2, bias=False)),
|
||||
nn.ReLU(inplace=True),
|
||||
self._norm_layer(128)
|
||||
)
|
||||
|
||||
self.gca = GuidedCxtAtten(128, 128)
|
||||
|
||||
# initialize guidance head
|
||||
for layers in range(len(self.guidance_head)):
|
||||
m = self.guidance_head[layers]
|
||||
if isinstance(m, nn.Conv2d):
|
||||
if hasattr(m, "weight_bar"):
|
||||
nn.init.xavier_uniform_(m.weight_bar)
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
nn.init.constant_(m.weight, 1)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
def _make_shortcut(self, inplane, planes):
|
||||
return nn.Sequential(
|
||||
SpectralNorm(nn.Conv2d(inplane, planes, kernel_size=3, padding=1, bias=False)),
|
||||
nn.ReLU(inplace=True),
|
||||
self._norm_layer(planes),
|
||||
SpectralNorm(nn.Conv2d(planes, planes, kernel_size=3, padding=1, bias=False)),
|
||||
nn.ReLU(inplace=True),
|
||||
self._norm_layer(planes)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.activation(out)
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
x1 = self.activation(out) # N x 32 x 256 x 256
|
||||
out = self.conv3(x1)
|
||||
out = self.bn3(out)
|
||||
out = self.activation(out)
|
||||
|
||||
im_fea = self.guidance_head(x[:, :3, ...]) # downsample origin image and extract features
|
||||
# if CONFIG.model.trimap_channel == 3:
|
||||
unknown = F.interpolate(x[:, 4:5, ...], scale_factor=1/8, mode='nearest')
|
||||
# else:
|
||||
# unknown = F.interpolate(x[:,3:,...].eq(1.).float(), scale_factor=1/8, mode='nearest')
|
||||
|
||||
x2 = self.layer1(out) # N x 64 x 128 x 128
|
||||
x3= self.layer2(x2) # N x 128 x 64 x 64
|
||||
x3, offset = self.gca(im_fea, x3, unknown) # contextual attention
|
||||
x4 = self.layer3(x3) # N x 256 x 32 x 32
|
||||
out = self.layer_bottleneck(x4) # N x 512 x 16 x 16
|
||||
|
||||
fea1 = self.shortcut[0](x) # input image and trimap
|
||||
fea2 = self.shortcut[1](x1)
|
||||
fea3 = self.shortcut[2](x2)
|
||||
fea4 = self.shortcut[3](x3)
|
||||
fea5 = self.shortcut[4](x4)
|
||||
|
||||
return out, {'shortcut': (fea1, fea2, fea3, fea4, fea5),
|
||||
'image_fea': im_fea,
|
||||
'unknown': unknown,
|
||||
'offset_1': offset}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
from matting.networks.encoders.resnet_enc import BasicBlock
|
||||
m = ResGuidedCxtAtten(BasicBlock, [3, 4, 4, 2])
|
||||
for m in m.modules():
|
||||
print(m)
|
||||
@@ -0,0 +1,51 @@
|
||||
import torch.nn as nn
|
||||
# from utils import CONFIG
|
||||
from hair_matting.matting.networks.encoders.resnet_enc import ResNet_D
|
||||
from hair_matting.matting.networks.ops import SpectralNorm
|
||||
|
||||
class ResShortCut_D(ResNet_D):
|
||||
|
||||
def __init__(self, block, layers, norm_layer=None, late_downsample=False):
|
||||
super(ResShortCut_D, self).__init__(block, layers, norm_layer, late_downsample=late_downsample)
|
||||
first_inplane = 3 + 3
|
||||
self.shortcut_inplane = [first_inplane, self.midplanes, 64, 128, 256]
|
||||
self.shortcut_plane = [32, self.midplanes, 64, 128, 256]
|
||||
|
||||
self.shortcut = nn.ModuleList()
|
||||
for stage, inplane in enumerate(self.shortcut_inplane):
|
||||
self.shortcut.append(self._make_shortcut(inplane, self.shortcut_plane[stage]))
|
||||
|
||||
def _make_shortcut(self, inplane, planes):
|
||||
return nn.Sequential(
|
||||
SpectralNorm(nn.Conv2d(inplane, planes, kernel_size=3, padding=1, bias=False)),
|
||||
nn.ReLU(inplace=True),
|
||||
self._norm_layer(planes),
|
||||
SpectralNorm(nn.Conv2d(planes, planes, kernel_size=3, padding=1, bias=False)),
|
||||
nn.ReLU(inplace=True),
|
||||
self._norm_layer(planes)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.activation(out)
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
x1 = self.activation(out) # N x 32 x 256 x 256
|
||||
out = self.conv3(x1)
|
||||
out = self.bn3(out)
|
||||
out = self.activation(out)
|
||||
|
||||
x2 = self.layer1(out) # N x 64 x 128 x 128
|
||||
x3= self.layer2(x2) # N x 128 x 64 x 64
|
||||
x4 = self.layer3(x3) # N x 256 x 32 x 32
|
||||
out = self.layer_bottleneck(x4) # N x 512 x 16 x 16
|
||||
|
||||
fea1 = self.shortcut[0](x) # input image and trimap
|
||||
fea2 = self.shortcut[1](x1)
|
||||
fea3 = self.shortcut[2](x2)
|
||||
fea4 = self.shortcut[3](x3)
|
||||
fea5 = self.shortcut[4](x4)
|
||||
|
||||
return out, {'shortcut':(fea1, fea2, fea3, fea4, fea5), 'image':x[:,:3,...]}
|
||||
@@ -0,0 +1,150 @@
|
||||
import logging
|
||||
import torch.nn as nn
|
||||
from hair_matting.matting.networks.ops import SpectralNorm
|
||||
|
||||
def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
|
||||
"""3x3 convolution with padding"""
|
||||
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
|
||||
padding=dilation, groups=groups, bias=False, dilation=dilation)
|
||||
|
||||
|
||||
def conv1x1(in_planes, out_planes, stride=1):
|
||||
"""1x1 convolution"""
|
||||
return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
|
||||
|
||||
|
||||
class BasicBlock(nn.Module):
|
||||
expansion = 1
|
||||
|
||||
def __init__(self, inplanes, planes, stride=1, downsample=None, norm_layer=None):
|
||||
super(BasicBlock, self).__init__()
|
||||
if norm_layer is None:
|
||||
norm_layer = nn.BatchNorm2d
|
||||
# Both self.conv1 and self.downsample layers downsample the input when stride != 1
|
||||
self.conv1 = SpectralNorm(conv3x3(inplanes, planes, stride))
|
||||
self.bn1 = norm_layer(planes)
|
||||
self.activation = nn.ReLU(inplace=True)
|
||||
self.conv2 = SpectralNorm(conv3x3(planes, planes))
|
||||
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.activation(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
|
||||
if self.downsample is not None:
|
||||
identity = self.downsample(x)
|
||||
|
||||
out += identity
|
||||
out = self.activation(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class ResNet_D(nn.Module):
|
||||
"""
|
||||
Implement and pre-train on ImageNet with the tricks from
|
||||
https://arxiv.org/abs/1812.01187
|
||||
without the mix-up part.
|
||||
"""
|
||||
|
||||
def __init__(self, block, layers, norm_layer=None, late_downsample=False):
|
||||
super(ResNet_D, self).__init__()
|
||||
self.logger = logging.getLogger("Logger")
|
||||
if norm_layer is None:
|
||||
norm_layer = nn.BatchNorm2d
|
||||
self._norm_layer = norm_layer
|
||||
|
||||
self.inplanes = 64
|
||||
self.late_downsample = late_downsample
|
||||
self.midplanes = 64 if late_downsample else 32
|
||||
self.start_stride = [1, 2, 1, 2] if late_downsample else [2, 1, 2, 1]
|
||||
self.conv1 = SpectralNorm(nn.Conv2d(3 + 3, 32, kernel_size=3,
|
||||
stride=self.start_stride[0], padding=1, bias=False))
|
||||
self.conv2 = SpectralNorm(nn.Conv2d(32, self.midplanes, kernel_size=3, stride=self.start_stride[1], padding=1,
|
||||
bias=False))
|
||||
self.conv3 = SpectralNorm(nn.Conv2d(self.midplanes, self.inplanes, kernel_size=3, stride=self.start_stride[2],
|
||||
padding=1, bias=False))
|
||||
self.bn1 = norm_layer(32)
|
||||
self.bn2 = norm_layer(self.midplanes)
|
||||
self.bn3 = norm_layer(self.inplanes)
|
||||
self.activation = nn.ReLU(inplace=True)
|
||||
self.layer1 = self._make_layer(block, 64, layers[0], stride=self.start_stride[3])
|
||||
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
|
||||
self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
|
||||
self.layer_bottleneck = self._make_layer(block, 512, layers[3], stride=2)
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
nn.init.xavier_uniform_(m.weight_bar)
|
||||
elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
|
||||
nn.init.constant_(m.weight, 1)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
# Zero-initialize the last BN in each residual branch,
|
||||
# so that the residual branch starts with zeros, and each residual block behaves like an identity.
|
||||
# This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, BasicBlock):
|
||||
nn.init.constant_(m.bn2.weight, 0)
|
||||
|
||||
self.logger.debug("encoder conv1 weight shape: {}".format(str(self.conv1.module.weight_bar.data.shape)))
|
||||
self.conv1.module.weight_bar.data[:,3:,:,:] = 0
|
||||
|
||||
self.logger.debug(self)
|
||||
|
||||
def _make_layer(self, block, planes, blocks, stride=1):
|
||||
if blocks == 0:
|
||||
return nn.Sequential(nn.Identity())
|
||||
norm_layer = self._norm_layer
|
||||
downsample = None
|
||||
if stride != 1:
|
||||
downsample = nn.Sequential(
|
||||
nn.AvgPool2d(2, stride),
|
||||
SpectralNorm(conv1x1(self.inplanes, planes * block.expansion)),
|
||||
norm_layer(planes * block.expansion),
|
||||
)
|
||||
elif self.inplanes != planes * block.expansion:
|
||||
downsample = nn.Sequential(
|
||||
SpectralNorm(conv1x1(self.inplanes, planes * block.expansion, stride)),
|
||||
norm_layer(planes * block.expansion),
|
||||
)
|
||||
|
||||
layers = [block(self.inplanes, planes, stride, downsample, norm_layer)]
|
||||
self.inplanes = planes * block.expansion
|
||||
for _ in range(1, blocks):
|
||||
layers.append(block(self.inplanes, planes, norm_layer=norm_layer))
|
||||
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.activation(x)
|
||||
x = self.conv2(x)
|
||||
x = self.bn2(x)
|
||||
x1 = self.activation(x) # N x 32 x 256 x 256
|
||||
x = self.conv3(x1)
|
||||
x = self.bn3(x)
|
||||
x2 = self.activation(x) # N x 64 x 128 x 128
|
||||
|
||||
x3 = self.layer1(x2) # N x 64 x 128 x 128
|
||||
x4 = self.layer2(x3) # N x 128 x 64 x 64
|
||||
x5 = self.layer3(x4) # N x 256 x 32 x 32
|
||||
x = self.layer_bottleneck(x5) # N x 512 x 16 x 16
|
||||
|
||||
return x, (x1, x2, x3, x4, x5)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
m = ResNet_D(BasicBlock, [3, 4, 4, 2])
|
||||
for m in m.modules():
|
||||
print(m._get_name())
|
||||
@@ -0,0 +1,60 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
# from utils import CONFIG
|
||||
from hair_matting.matting.networks import encoders, decoders
|
||||
|
||||
|
||||
class Generator(nn.Module):
|
||||
def __init__(self, encoder, decoder, num_class=1):
|
||||
|
||||
super(Generator, self).__init__()
|
||||
|
||||
if encoder not in encoders.__all__:
|
||||
raise NotImplementedError("Unknown Encoder {}".format(encoder))
|
||||
self.encoder = encoders.__dict__[encoder]()
|
||||
|
||||
if decoder not in decoders.__all__:
|
||||
raise NotImplementedError("Unknown Decoder {}".format(decoder))
|
||||
self.decoder = decoders.__dict__[decoder](num_class)
|
||||
|
||||
def forward(self, image, trimap):
|
||||
inp = torch.cat((image, trimap), dim=1)
|
||||
embedding, mid_fea = self.encoder(inp)
|
||||
alpha, info_dict = self.decoder(embedding, mid_fea)
|
||||
|
||||
return alpha, info_dict
|
||||
|
||||
|
||||
def get_generator(encoder, decoder, num_class=1):
|
||||
generator = Generator(encoder=encoder, decoder=decoder, num_class=num_class)
|
||||
return generator
|
||||
|
||||
|
||||
if __name__=="__main__":
|
||||
import time
|
||||
# generator = get_generator(encoder=CONFIG.model.arch.encoder, decoder=CONFIG.model.arch.decoder).cuda().train()
|
||||
batch_size = 12
|
||||
# generator.eval()
|
||||
n_eval = 10
|
||||
# pre run the model
|
||||
# with torch.no_grad():
|
||||
# for i in range(2):
|
||||
# x = torch.rand(batch_size, 3, 512, 512, device=device)
|
||||
# y = torch.rand(batch_size, 3, 512, 512, device=device)
|
||||
# z = generator(x,y)
|
||||
# test without GPU IO
|
||||
|
||||
# x = torch.zeros(batch_size, 3, 512, 512, device=device)
|
||||
# y = torch.zeros(batch_size, 1, 512, 512, device=device)
|
||||
x = torch.randn(batch_size, 3, 512, 512)
|
||||
y = torch.randn(batch_size, 3, 512, 512)
|
||||
t = time.time()
|
||||
# with torch.no_grad():
|
||||
# for i in range(n_eval):
|
||||
# a = generator(x.cuda(),y.cuda())
|
||||
# torch.cuda.synchronize()
|
||||
# print(generator.__class__.__name__, 'With IO \t', f'{(time.time() - t)/n_eval/batch_size:.5f} s')
|
||||
# print(generator.__class__.__name__, 'FPS \t\t', f'{1/((time.time() - t)/n_eval/batch_size):.5f} s')
|
||||
# for n, p in generator.named_parameters():
|
||||
# print(n)
|
||||
@@ -0,0 +1,256 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn import Parameter
|
||||
from torch.autograd import Variable
|
||||
from torch.nn import functional as F
|
||||
|
||||
|
||||
def l2normalize(v, eps=1e-12):
|
||||
return v / (v.norm() + eps)
|
||||
|
||||
|
||||
class SpectralNorm(nn.Module):
|
||||
"""
|
||||
Based on https://github.com/heykeetae/Self-Attention-GAN/blob/master/spectral.py
|
||||
and add _noupdate_u_v() for evaluation
|
||||
"""
|
||||
def __init__(self, module, name='weight', power_iterations=1):
|
||||
super(SpectralNorm, self).__init__()
|
||||
self.module = module
|
||||
self.name = name
|
||||
self.power_iterations = power_iterations
|
||||
if not self._made_params():
|
||||
self._make_params()
|
||||
|
||||
def _update_u_v(self):
|
||||
u = getattr(self.module, self.name + "_u")
|
||||
v = getattr(self.module, self.name + "_v")
|
||||
w = getattr(self.module, self.name + "_bar")
|
||||
|
||||
height = w.data.shape[0]
|
||||
for _ in range(self.power_iterations):
|
||||
v.data = l2normalize(torch.mv(torch.t(w.view(height,-1).data), u.data))
|
||||
u.data = l2normalize(torch.mv(w.view(height,-1).data, v.data))
|
||||
|
||||
sigma = u.dot(w.view(height, -1).mv(v))
|
||||
setattr(self.module, self.name, w / sigma.expand_as(w))
|
||||
|
||||
def _noupdate_u_v(self):
|
||||
u = getattr(self.module, self.name + "_u")
|
||||
v = getattr(self.module, self.name + "_v")
|
||||
w = getattr(self.module, self.name + "_bar")
|
||||
|
||||
height = w.data.shape[0]
|
||||
sigma = u.dot(w.view(height, -1).mv(v))
|
||||
setattr(self.module, self.name, w / sigma.expand_as(w))
|
||||
|
||||
def _made_params(self):
|
||||
try:
|
||||
u = getattr(self.module, self.name + "_u")
|
||||
v = getattr(self.module, self.name + "_v")
|
||||
w = getattr(self.module, self.name + "_bar")
|
||||
return True
|
||||
except AttributeError:
|
||||
return False
|
||||
|
||||
def _make_params(self):
|
||||
w = getattr(self.module, self.name)
|
||||
|
||||
height = w.data.shape[0]
|
||||
width = w.view(height, -1).data.shape[1]
|
||||
|
||||
u = Parameter(w.data.new(height).normal_(0, 1), requires_grad=False)
|
||||
v = Parameter(w.data.new(width).normal_(0, 1), requires_grad=False)
|
||||
u.data = l2normalize(u.data)
|
||||
v.data = l2normalize(v.data)
|
||||
w_bar = Parameter(w.data)
|
||||
|
||||
del self.module._parameters[self.name]
|
||||
|
||||
self.module.register_parameter(self.name + "_u", u)
|
||||
self.module.register_parameter(self.name + "_v", v)
|
||||
self.module.register_parameter(self.name + "_bar", w_bar)
|
||||
|
||||
def forward(self, *args):
|
||||
# if torch.is_grad_enabled() and self.module.training:
|
||||
if self.module.training:
|
||||
self._update_u_v()
|
||||
else:
|
||||
self._noupdate_u_v()
|
||||
return self.module.forward(*args)
|
||||
|
||||
|
||||
class GuidedCxtAtten(nn.Module):
|
||||
# based on https://github.com/nbei/Deep-Flow-Guided-Video-Inpainting/blob/a6fe298fec502bfd9cbc64eb01e39f78a3262a59/models/DeepFill_Models/ops.py#L210
|
||||
def __init__(self, out_channels, guidance_channels, rate=2):
|
||||
super(GuidedCxtAtten, self).__init__()
|
||||
self.rate = rate
|
||||
self.padding = nn.ReflectionPad2d(1)
|
||||
self.up_sample = nn.Upsample(scale_factor=self.rate, mode='nearest')
|
||||
|
||||
self.guidance_conv = nn.Conv2d(in_channels=guidance_channels, out_channels=guidance_channels//2,
|
||||
kernel_size=1, stride=1, padding=0)
|
||||
|
||||
self.W = nn.Sequential(
|
||||
nn.Conv2d(in_channels=out_channels, out_channels=out_channels,
|
||||
kernel_size=1, stride=1, padding=0, bias=False),
|
||||
nn.BatchNorm2d(out_channels)
|
||||
)
|
||||
|
||||
nn.init.xavier_uniform_(self.guidance_conv.weight)
|
||||
nn.init.constant_(self.guidance_conv.bias, 0)
|
||||
nn.init.xavier_uniform_(self.W[0].weight)
|
||||
nn.init.constant_(self.W[1].weight, 1e-3)
|
||||
nn.init.constant_(self.W[1].bias, 0)
|
||||
|
||||
def forward(self, f, alpha, unknown=None, ksize=3, stride=1, fuse_k=3, softmax_scale=1., training=True):
|
||||
|
||||
f = self.guidance_conv(f)
|
||||
# get shapes
|
||||
raw_int_fs = list(f.size()) # N x 64 x 64 x 64
|
||||
raw_int_alpha = list(alpha.size()) # N x 128 x 64 x 64
|
||||
|
||||
# extract patches from background with stride and rate
|
||||
kernel = 2*self.rate
|
||||
alpha_w = self.extract_patches(alpha, kernel=kernel, stride=self.rate)
|
||||
alpha_w = alpha_w.permute(0, 2, 3, 4, 5, 1)
|
||||
alpha_w = alpha_w.contiguous().view(raw_int_alpha[0], raw_int_alpha[2] // self.rate, raw_int_alpha[3] // self.rate, -1)
|
||||
alpha_w = alpha_w.contiguous().view(raw_int_alpha[0], -1, kernel, kernel, raw_int_alpha[1])
|
||||
alpha_w = alpha_w.permute(0, 1, 4, 2, 3)
|
||||
|
||||
f = F.interpolate(f, scale_factor=1/self.rate, mode='nearest')
|
||||
|
||||
fs = f.size() # B x 64 x 32 x 32
|
||||
f_groups = torch.split(f, 1, dim=0) # Split tensors by batch dimension; tuple is returned
|
||||
|
||||
# from b(B*H*W*C) to w(b*k*k*c*h*w)
|
||||
int_fs = list(fs)
|
||||
w = self.extract_patches(f)
|
||||
w = w.permute(0, 2, 3, 4, 5, 1)
|
||||
w = w.contiguous().view(raw_int_fs[0], raw_int_fs[2] // self.rate, raw_int_fs[3] // self.rate, -1)
|
||||
w = w.contiguous().view(raw_int_fs[0], -1, ksize, ksize, raw_int_fs[1])
|
||||
w = w.permute(0, 1, 4, 2, 3)
|
||||
# process mask
|
||||
|
||||
if unknown is not None:
|
||||
unknown = unknown.clone()
|
||||
unknown = F.interpolate(unknown, scale_factor=1/self.rate, mode='nearest')
|
||||
assert unknown.size(2) == f.size(2), "mask should have same size as f at dim 2,3"
|
||||
unknown_mean = unknown.mean(dim=[2,3])
|
||||
known_mean = 1 - unknown_mean
|
||||
unknown_scale = torch.clamp(torch.sqrt(unknown_mean / known_mean), 0.1, 10).to(alpha)
|
||||
known_scale = torch.clamp(torch.sqrt(known_mean / unknown_mean), 0.1, 10).to(alpha)
|
||||
softmax_scale = torch.cat([unknown_scale, known_scale], dim=1)
|
||||
else:
|
||||
unknown = torch.ones([fs[0], 1, fs[2], fs[3]]).to(alpha)
|
||||
softmax_scale = torch.FloatTensor([softmax_scale, softmax_scale]).view(1,2).repeat(fs[0],1).to(alpha)
|
||||
|
||||
m = self.extract_patches(unknown)
|
||||
|
||||
m = m.permute(0, 2, 3, 4, 5, 1)
|
||||
m = m.contiguous().view(raw_int_fs[0], raw_int_fs[2]//self.rate, raw_int_fs[3]//self.rate, -1)
|
||||
m = m.contiguous().view(raw_int_fs[0], -1, ksize, ksize)
|
||||
|
||||
m = self.reduce_mean(m) # smoothing, maybe
|
||||
# mask out the
|
||||
mm = m.gt(0.).float() # (N, 32*32, 1, 1)
|
||||
|
||||
# the correlation with itself should be 0
|
||||
self_mask = F.one_hot(torch.arange(fs[2] * fs[3]).view(fs[2], fs[3]).contiguous().to(alpha).long(),
|
||||
num_classes=int_fs[2] * int_fs[3])
|
||||
self_mask = self_mask.permute(2, 0, 1).view(1, fs[2] * fs[3], fs[2], fs[3]).float() * (-1e4)
|
||||
|
||||
w_groups = torch.split(w, 1, dim=0) # Split tensors by batch dimension; tuple is returned
|
||||
alpha_w_groups = torch.split(alpha_w, 1, dim=0) # Split tensors by batch dimension; tuple is returned
|
||||
mm_groups = torch.split(mm, 1, dim=0)
|
||||
scale_group = torch.split(softmax_scale, 1, dim=0)
|
||||
y = []
|
||||
offsets = []
|
||||
k = fuse_k
|
||||
y_test = []
|
||||
for xi, wi, alpha_wi, mmi, scale in zip(f_groups, w_groups, alpha_w_groups, mm_groups, scale_group):
|
||||
# conv for compare
|
||||
wi = wi[0]
|
||||
escape_NaN = Variable(torch.FloatTensor([1e-4])).to(alpha)
|
||||
wi_normed = wi / torch.max(self.l2_norm(wi), escape_NaN)
|
||||
xi = F.pad(xi, (1,1,1,1), mode='reflect')
|
||||
yi = F.conv2d(xi, wi_normed, stride=1, padding=0) # yi => (B=1, C=32*32, H=32, W=32)
|
||||
y_test.append(yi)
|
||||
# conv implementation for fuse scores to encourage large patches
|
||||
yi = yi.permute(0, 2, 3, 1)
|
||||
yi = yi.contiguous().view(1, fs[2], fs[3], fs[2] * fs[3])
|
||||
yi = yi.permute(0, 3, 1, 2) # (B=1, C=32*32, H=32, W=32)
|
||||
|
||||
# softmax to match
|
||||
# scale the correlation with predicted scale factor for known and unknown area
|
||||
yi = yi * (scale[0,0] * mmi.gt(0.).float() + scale[0,1] * mmi.le(0.).float()) # mmi => (1, 32*32, 1, 1)
|
||||
# mask itself, self-mask only applied to unknown area
|
||||
yi = yi + self_mask * mmi # self_mask: (1, 32*32, 32, 32)
|
||||
# for small input inference
|
||||
yi = F.softmax(yi, dim=1)
|
||||
|
||||
_, offset = torch.max(yi, dim=1) # argmax; index
|
||||
offset = torch.stack([offset // fs[3], offset % fs[3]], dim=1)
|
||||
|
||||
wi_center = alpha_wi[0]
|
||||
|
||||
if self.rate == 1:
|
||||
left = (kernel) // 2
|
||||
right = (kernel - 1) // 2
|
||||
yi = F.pad(yi, (left, right, left, right), mode='reflect')
|
||||
wi_center = wi_center.permute(1, 0, 2, 3)
|
||||
yi = F.conv2d(yi, wi_center, padding=0) / 4. # (B=1, C=128, H=64, W=64)
|
||||
else:
|
||||
yi = F.conv_transpose2d(yi, wi_center, stride=self.rate, padding=1) / 4. # (B=1, C=128, H=64, W=64)
|
||||
y.append(yi)
|
||||
offsets.append(offset)
|
||||
|
||||
y = torch.cat(y, dim=0) # back to the mini-batch
|
||||
y.contiguous().view(raw_int_alpha)
|
||||
offsets = torch.cat(offsets, dim=0)
|
||||
offsets = offsets.view([int_fs[0]] + [2] + int_fs[2:])
|
||||
|
||||
# # case1: visualize optical flow: minus current position
|
||||
# h_add = Variable(torch.arange(0,float(fs[2]))).to(alpha).view([1, 1, fs[2], 1])
|
||||
# h_add = h_add.expand(fs[0], 1, fs[2], fs[3])
|
||||
# w_add = Variable(torch.arange(0,float(fs[3]))).to(alpha).view([1, 1, 1, fs[3]])
|
||||
# w_add = w_add.expand(fs[0], 1, fs[2], fs[3])
|
||||
#
|
||||
# offsets = offsets - torch.cat([h_add, w_add], dim=1).long()
|
||||
|
||||
# case2: visualize absolute position
|
||||
offsets = offsets - torch.Tensor([fs[2]//2, fs[3]//2]).view(1,2,1,1).to(alpha).long()
|
||||
|
||||
y = self.W(y) + alpha
|
||||
|
||||
return y, (offsets, softmax_scale)
|
||||
|
||||
@staticmethod
|
||||
def extract_patches(x, kernel=3, stride=1):
|
||||
left =(kernel - stride + 1) // 2
|
||||
right =(kernel - stride) // 2
|
||||
x = F.pad(x, (left, right, left, right), mode='reflect')
|
||||
all_patches = x.unfold(2, kernel, stride).unfold(3, kernel, stride)
|
||||
|
||||
return all_patches
|
||||
|
||||
@staticmethod
|
||||
def reduce_mean(x):
|
||||
for i in range(4):
|
||||
if i <= 1:
|
||||
continue
|
||||
x = torch.mean(x, dim=i, keepdim=True)
|
||||
return x
|
||||
|
||||
@staticmethod
|
||||
def l2_norm(x):
|
||||
def reduce_sum(x):
|
||||
for i in range(4):
|
||||
if i == 0:
|
||||
continue
|
||||
x = torch.sum(x, dim=i, keepdim=True)
|
||||
return x
|
||||
|
||||
x = x**2
|
||||
x = reduce_sum(x)
|
||||
return torch.sqrt(x)
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user