def processing(image_path, file_title):
image = cv2.imread(image_path)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
face_rect = face_detect(image)
if face_rect is None:
return
frameOpencvDnn = image.copy()
cv2.rectangle(frameOpencvDnn, (face_rect[0], face_rect[1]), (face_rect[2], face_rect[3]), (0, 255, 0), int(round(image.shape[0] / 150)), 8)
#cv2.imwrite("result/detect.png", cv2.cvtColor(frameOpencvDnn, cv2.COLOR_BGR2RGB))
face_image, image_rect, image_len = face_crop(image, face_rect)
resize_image = cv2.resize(face_image, (512, 512), cv2.INTER_LINEAR).astype(np.uint8)
#cv2.imwrite("result/normalize.png", cv2.cvtColor(resize_image, cv2.COLOR_BGR2RGB))
seg_image = segmentation(resize_image)
background, body_image, hat_image, hair_image, trip_image = pre_processing(seg_image)
matting_image = inference_img_whole(args, matting_model, resize_image, trip_image)
#cv2.imwrite("result/matting.png", matting_image * 255)
alpha_image = np.zeros(resize_image.shape)
alpha_image[:, :, 0] = matting_image * 255 + hat_image * 255 * (1 - matting_image) + hair_image * 255 * (
1 - matting_image)
alpha_image[:, :, 1] = matting_image * 255 + body_image * 255 * (1 - matting_image) + hair_image * 255 * (
1 - matting_image)
alpha_image[:, :, 2] = matting_image * 255 + body_image * 255 * (1 - matting_image) + hair_image * 255 * (
1 - matting_image)
bk_image = (resize_image + alpha_image) * 0.5
alpha_image = cv2.cvtColor(alpha_image.astype(np.uint8), cv2.COLOR_BGR2RGB)
resize_image = cv2.cvtColor(resize_image.astype(np.uint8), cv2.COLOR_BGR2RGB)
bk_image = cv2.cvtColor(bk_image.astype(np.uint8), cv2.COLOR_BGR2RGB)
alpha_image = cv2.resize(alpha_image, (image_len, image_len), cv2.INTER_LINEAR)
resize_image = cv2.resize(resize_image, (image_len, image_len), cv2.INTER_LINEAR)
bk_image = cv2.resize(bk_image, (image_len, image_len), cv2.INTER_LINEAR)
alpha_image = alpha_image[image_rect[1]: image_rect[3], image_rect[0]: image_rect[2]]
resize_image = resize_image[image_rect[1]: image_rect[3], image_rect[0]: image_rect[2]]
bk_image = bk_image[image_rect[1]: image_rect[3], image_rect[0]: image_rect[2]]
cv2.imwrite("result/{}_!real.png".format(file_title), resize_image)
cv2.imwrite("result/{}_alpha.png".format(file_title), alpha_image)
cv2.imwrite("result/{}_composition.png".format(file_title), bk_image)
wid = 600
hei = wid * alpha_image.shape[0] // alpha_image.shape[1]
cv2.imshow('alpha_image', cv2.resize(alpha_image, (wid, hei)))
cv2.imshow('bk_image', cv2.resize(bk_image, (wid, hei)))
cv2.waitKey(0)
def matting_result(pic_input, tri_input):
model_input = model
original_im = np.array(pic_input)[:, :, :3]
trimap_im = np.array(tri_input)
if len(trimap_im.shape)>2:
trimap_im = trimap_im[:, :, 0]
with torch.no_grad():
alpha = inference_img_whole(model_input, original_im, trimap_im)
alpha[trimap_im == 0] = 0.0
alpha[trimap_im == 255] = 1.0
h, w = original_im.shape[:2]
new_bg = np.array(np.full((h,w,3), 255), dtype='uint8')
im = composite4(original_im, new_bg, alpha, w, h)
im = Image.fromarray(im)
return im
def predictImg(image,trimap,row,col):
result_dir = os.getcwd()+"/pred"
if not os.path.exists(result_dir):
os.makedirs(result_dir)
# parameters setting
parser = argparse.ArgumentParser()
args = parser.parse_args()
args.cuda = False
args.resume = "models/deep_image_matting/stage_sad.pth"
args.stage = 1
args.crop_or_resize = "whole"
args.max_size = 1600
# init model
model = net.VGG16(args)
ckpt = torch.load(args.resume, map_location='cpu')
model.load_state_dict(ckpt['state_dict'], strict=True)
# model = model.cuda()
# infer one by one
torch.cuda.empty_cache()
with torch.no_grad():
pred_mattes = inference_img_whole(args, model, image, trimap)
pred_mattes = (pred_mattes * 255).astype(np.uint8)
pred_mattes[trimap == 255] = 255
pred_mattes[trimap == 0] = 0
# target = image * pred_mattes[:,:,np.newaxis]
# cv2.imshow('target', cv2.cvtColor(target,cv2.COLOR_BGR2RGB))
# cv2.waitKey(0)
# bg = cv2.imread(os.getcwd() + '/composebg/input_image.jpg')
bg = create_bgimage(row,col)
# print('00000000')
# cv2.imshow('img33',image)
im, bg = composite4(image, bg, pred_mattes, col, row)
timestamp = str(int(time.time()))
cv2.imwrite(os.getcwd()+'/pred/test'+timestamp+'.png', pred_mattes)
fileNameAndSuffix=os.path.split(image_path)[1].split('.')
newfileNameAndSuffix=os.getcwd()+'/pred/'+fileNameAndSuffix[0]+'_replace_bg'+timestamp+'.'+fileNameAndSuffix[1]
print('=============:'+newfileNameAndSuffix)
cv2.imwrite(newfileNameAndSuffix,im)
if isDebug:
cv2.waitKey(0)
return newfileNameAndSuffix
def getnewalpha(image, mask, cFlag):
if image.shape[2] == 4: # get rid of alpha channel
image = image[:, :, 0:3]
if mask.shape[2] == 4: # get rid of alpha channel
mask = mask[:, :, 0:3]
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
trimap = mask[:, :, 0]
cudaFlag = False
if torch.cuda.is_available() and not cFlag:
cudaFlag = True
args = Namespace(crop_or_resize='whole',
cuda=cudaFlag,
max_size=1600,
resume=baseLoc +
'weights/deepmatting/stage1_sad_57.1.pth',
stage=1)
model = net.VGG16(args)
if cudaFlag:
ckpt = torch.load(args.resume)
else:
ckpt = torch.load(args.resume, map_location=torch.device("cpu"))
model.load_state_dict(ckpt['state_dict'], strict=True)
if cudaFlag:
model = model.cuda()
# ckpt = torch.load(args.resume)
# model.load_state_dict(ckpt['state_dict'], strict=True)
# model = model.cuda()
torch.cuda.empty_cache()
with torch.no_grad():
pred_mattes = inference_img_whole(args, model, image, trimap)
pred_mattes = (pred_mattes * 255).astype(np.uint8)
pred_mattes[trimap == 255] = 255
pred_mattes[trimap == 0] = 0
# pred_mattes = np.repeat(pred_mattes[:, :, np.newaxis], 3, axis=2)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
pred_mattes = np.dstack((image, pred_mattes))
return pred_mattes
def test(args, model):
model.eval()
sample_set = []
img_ids = os.listdir(args.testImgDir)
img_ids.sort()
cnt = len(img_ids)
mse_diffs = 0.
sad_diffs = 0.
cur = 0
t0 = time.time()
for img_id in img_ids:
img_path = os.path.join(args.testImgDir, img_id)
trimap_path = os.path.join(args.testTrimapDir, img_id)
assert(os.path.exists(img_path))
assert(os.path.exists(trimap_path))
img = cv2.imread(img_path)
trimap = cv2.imread(trimap_path)[:, :, 0]
assert(img.shape[:2] == trimap.shape[:2])
img_info = (img_path.split('/')[-1], img.shape[0], img.shape[1])
cur += 1
print('[{}/{}] {}'.format(cur, cnt, img_info[0]))
with torch.no_grad():
torch.cuda.empty_cache()
if args.crop_or_resize == "whole":
origin_pred_mattes = inference_img_whole(args, model, img, trimap)
elif args.crop_or_resize == "crop":
origin_pred_mattes = inference_img_by_crop(args, model, img, trimap)
else:
origin_pred_mattes = inference_img_by_resize(args, model, img, trimap)
# only attention unknown region
origin_pred_mattes[trimap == 255] = 1.
origin_pred_mattes[trimap == 0 ] = 0.
# origin trimap
pixel = float((trimap == 128).sum())
# eval if gt alpha is given
if args.testAlphaDir != '':
alpha_name = os.path.join(args.testAlphaDir, img_info[0])
assert(os.path.exists(alpha_name))
alpha = cv2.imread(alpha_name)[:, :, 0] / 255.
assert(alpha.shape == origin_pred_mattes.shape)
mse_diff = ((origin_pred_mattes - alpha) ** 2).sum() / pixel
sad_diff = np.abs(origin_pred_mattes - alpha).sum()
mse_diffs += mse_diff
sad_diffs += sad_diff
print("sad:{} mse:{}".format(sad_diff, mse_diff))
origin_pred_mattes = (origin_pred_mattes * 255).astype(np.uint8)
if not os.path.exists(args.testResDir):
os.makedirs(args.testResDir)
cv2.imwrite(os.path.join(args.testResDir, img_info[0]), origin_pred_mattes)
print("Avg-Cost: {} s/image".format((time.time() - t0) / cnt))
if args.testAlphaDir != '':
print("Eval-MSE: {}".format(mse_diffs / cnt))
print("Eval-SAD: {}".format(sad_diffs / cnt))
return sad_diffs / cnt
# input file list
image_path = "boy-1518482_1920_12_img.png"
trimap_path = "boy-1518482_1920_12.png"
image = cv2.imread(image_path)
trimap = cv2.imread(trimap_path)
# print(trimap.shape)
trimap = trimap[:, :, 0]
# init model
args = Namespace(crop_or_resize='whole',
cuda=True,
max_size=1600,
resume='model/stage1_sad_57.1.pth',
stage=1)
model = net.VGG16(args)
ckpt = torch.load(args.resume)
model.load_state_dict(ckpt['state_dict'], strict=True)
model = model.cuda()
torch.cuda.empty_cache()
with torch.no_grad():
pred_mattes = inference_img_whole(args, model, image, trimap)
pred_mattes = (pred_mattes * 255).astype(np.uint8)
pred_mattes[trimap == 255] = 255
pred_mattes[trimap == 0] = 0
# print(pred_mattes)
# cv2.imwrite('out.png', pred_mattes)
# import matplotlib.pyplot as plt
# plt.imshow(image)
# plt.show()