初始化换发型项目：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,158 @@
|
||||
from hair_matting.matting.gca_matting_hair_single_fg import *
|
||||
from hair_matting.seg.hairseg_single_model import Evaluator
|
||||
from time import time
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
class GenTrimap(object):
|
||||
def __init__(self):
|
||||
self.erosion_kernels = [None] + [cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (size, size)) for size in range(1,30)]
|
||||
|
||||
def __call__(self, alpha):
|
||||
|
||||
fg_mask = np.zeros_like(alpha)
|
||||
bg_mask = np.zeros_like(alpha)
|
||||
fg_mask[alpha == 255] = 1
|
||||
bg_mask[alpha == 0] = 1
|
||||
|
||||
fg_mask = fg_mask.astype(np.int).astype(np.uint8)
|
||||
bg_mask = bg_mask.astype(np.int).astype(np.uint8)
|
||||
|
||||
fg_mask = cv2.erode(fg_mask, self.erosion_kernels[15])
|
||||
bg_mask = cv2.erode(bg_mask, self.erosion_kernels[29])
|
||||
|
||||
trimap = np.ones_like(alpha, dtype=np.uint8) * 128
|
||||
trimap[fg_mask == 1] = 255
|
||||
trimap[bg_mask == 1] = 0
|
||||
return trimap
|
||||
|
||||
class Generator_Matte(object):
|
||||
|
||||
def __init__(self, gpu, device_id):
|
||||
if gpu and torch.cuda.is_available():
|
||||
self.device = torch.device("cuda:%d" % device_id if device_id >= 0 else "cpu")
|
||||
self.cuda = True
|
||||
else:
|
||||
self.device = torch.device("cpu")
|
||||
self.cuda = False
|
||||
|
||||
self.debug_flag = False
|
||||
self.model_dir = os.path.dirname(__file__)
|
||||
self.output_img_size = 512
|
||||
|
||||
triseg_model_path = os.path.join(self.model_dir, 'deeplabv3_hair512_360_0520_wl.pth') # deeplabv3_hair512_520_0121 deeplabv3_hair512_360_0520_wl
|
||||
self.triseg_model = Evaluator(gpu_id=device_id, output_img_size=self.output_img_size, nclass=3, seg_model_path=triseg_model_path)
|
||||
|
||||
hair_matte_model_path = os.path.join(self.model_dir, 'gca-dist-fg-0430-latest_model.pth') # gca-dist-fg-0203-latest_model gca-dist-fg-0430-latest_model
|
||||
self.matte_model = self.load_hair_matte_model(hair_matte_model_path)
|
||||
|
||||
self.gen_trimap = GenTrimap()
|
||||
|
||||
def remove_prefix_state_dict(self, state_dict, prefix="module"):
|
||||
"""
|
||||
remove prefix from the key of pretrained state dict for Data-Parallel
|
||||
"""
|
||||
new_state_dict = {}
|
||||
first_state_name = list(state_dict.keys())[0]
|
||||
if not first_state_name.startswith(prefix):
|
||||
for key, value in state_dict.items():
|
||||
new_state_dict[key] = state_dict[key].float()
|
||||
else:
|
||||
for key, value in state_dict.items():
|
||||
new_state_dict[key[len(prefix) + 1:]] = state_dict[key].float()
|
||||
return new_state_dict
|
||||
|
||||
|
||||
def load_hair_matte_model(self, hair_matte_model_path):
|
||||
# build model
|
||||
model = networks.get_generator(encoder="resnet_gca_encoder_29", decoder="res_gca_decoder_22", num_class=4)
|
||||
|
||||
# load checkpoint
|
||||
checkpoint = torch.load(hair_matte_model_path, map_location=lambda storage, loc: storage)
|
||||
model.load_state_dict(self.remove_prefix_state_dict(checkpoint['state_dict']), strict=True)
|
||||
model.to(self.device)
|
||||
# print("matte_model: ", model)
|
||||
|
||||
# inference
|
||||
model = model.eval()
|
||||
return model
|
||||
|
||||
def matte_inference(self, image, landmark1k):
|
||||
|
||||
trimap = self.triseg_model.eval(image, landmark1k)
|
||||
|
||||
ori_h, ori_w, _ = image.shape
|
||||
|
||||
limit_size = 1600
|
||||
if ori_h > limit_size or ori_w > limit_size:
|
||||
if ori_h > ori_w:
|
||||
new_tri_h = limit_size
|
||||
new_tri_w = int(ori_w * limit_size / ori_h)
|
||||
else:
|
||||
new_tri_w = limit_size
|
||||
new_tri_h = int(ori_h * limit_size / ori_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
|
||||
|
||||
# cv2.imshow("image before: ", image)
|
||||
# cv2.imshow("trimap before: ", trimap)
|
||||
# cv2.waitKey()
|
||||
trimap = self.gen_trimap(trimap_resize[:, :, 0])
|
||||
|
||||
# cv2.imwrite(os.path.join(args.output, image_name.replace(ext, "_trimap.png")), trimap)
|
||||
|
||||
image_dict = generator_tensor_dict(image_resize, trimap)
|
||||
|
||||
pred_fg, pred, offset = single_inference(self.matte_model, image_dict, device=self.device)
|
||||
|
||||
pred_fg[trimap==1] = image_resize[trimap==1]
|
||||
|
||||
if pred.shape[1] != image.shape[1] or pred.shape[0] != image.shape[0]:
|
||||
pred = cv2.resize(pred, (image.shape[1], image.shape[0]))
|
||||
return pred_fg, pred
|
||||
|
||||
if __name__ == '__main__':
|
||||
print('Torch Version: ', torch.__version__)
|
||||
from utils.data_preprocess import *
|
||||
|
||||
generator_matte = Generator_Matte(True, 0)
|
||||
data_dir = "/media/DATA_4T/hair_data/check_fafeng_trimap/pick_test_fafeng/pick_liuhai"
|
||||
for image_name in os.listdir(data_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(data_dir, image_name)
|
||||
|
||||
start = time()
|
||||
|
||||
image = cv2.imread(image_path)
|
||||
|
||||
pt_path = image_path.replace(ext, '_landmark1k.txt')
|
||||
|
||||
if not os.path.exists(pt_path):
|
||||
continue
|
||||
|
||||
landmark1k = np.loadtxt(pt_path)
|
||||
|
||||
matting_start = time()
|
||||
pred = generator_matte.matte_inference(image, landmark1k)
|
||||
matting_end = time()
|
||||
print("matting cost time :{:.4f}".format(matting_end - matting_start))
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
pred = cv2.resize(pred, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_CUBIC)
|
||||
# cv2.imshow("pred:", pred)
|
||||
# cv2.waitKey()
|
||||
|
||||
end = time()
|
||||
print('end img_path: {}, Time cost : {:.4f}'.format(image_path, end-start))
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:9fd1c65a4c002dfb215771a3f2882740e08c49dc974e982a3e9a934e6f4355a5
|
||||
size 233017644
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:9003635a734d030978e8d35e13de3e226a06f56767c443c868562d30a51acdfb
|
||||
size 302856998
|
||||
@@ -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
|
||||
@@ -0,0 +1,95 @@
|
||||
import os
|
||||
import torch
|
||||
|
||||
from hair_matting.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
|
||||
|
||||
pre_start = time.time()
|
||||
|
||||
M1 = landmark_processor.get_transform_mat_hair(pts1k, self.output_img_size, ratio=0.28, h_ratio=0.30)
|
||||
img = cv2.warpAffine(img, M1, (self.output_img_size, self.output_img_size), flags=cv2.INTER_LANCZOS4)
|
||||
img = self.process_data(img)
|
||||
img = img.to(self.device)
|
||||
pre_end = time.time()
|
||||
# print("seg pre cost time : {:.6f} s".format(pre_end - pre_start))
|
||||
|
||||
with torch.no_grad():
|
||||
# torch.cuda.synchronize()
|
||||
start = time.time()
|
||||
outputs = self.model(img)
|
||||
end = time.time()
|
||||
# print("seg forward cost time : {:.6f} s".format(end - start))
|
||||
pred = torch.argmax(outputs[0], 1)
|
||||
# print("argmax cost time : {:.6f} s".format(time.time() - end))
|
||||
|
||||
pred = pred[0].detach().cpu().numpy()
|
||||
predict = pred.astype(np.float32)
|
||||
pred_detach = time.time()
|
||||
# print("pred_detach cost time : {:.6f} s".format(pred_detach - end))
|
||||
|
||||
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()
|
||||
# print("seg post cost time : {:.6f} s".format(time.time() - end))
|
||||
|
||||
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 hair_matting.seg.networks.segbase import SegBaseModel
|
||||
from hair_matting.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 hair_matting.seg.networks.xception import get_xception
|
||||
from hair_matting.seg.networks.deeplabv3 import _ASPP
|
||||
from hair_matting.seg.networks.fcn import _FCNHead
|
||||
from hair_matting.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 hair_matting.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 hair_matting.seg.networks.jpu import JPU
|
||||
from hair_matting.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