diff --git a/contrib/HumanSeg/bg_replace.py b/contrib/HumanSeg/bg_replace.py index 06f2e550e9875ed0ce0cb7723d707f3036fd9806..3eee82fb7e865bd1371b3c907a5ca55a43de18ed 100644 --- a/contrib/HumanSeg/bg_replace.py +++ b/contrib/HumanSeg/bg_replace.py @@ -196,8 +196,8 @@ def infer(args): score_map, im_info = predict(frame, model, test_transforms) cur_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) cur_gray = cv2.resize(cur_gray, (resize_w, resize_h)) - scoremap = 255 * score_map[:, :, 1] - optflow_map = postprocess(cur_gray, scoremap, prev_gray, prev_cfd, \ + score_map = 255 * score_map[:, :, 1] + optflow_map = postprocess(cur_gray, score_map, prev_gray, prev_cfd, \ disflow, is_init) prev_gray = cur_gray.copy() prev_cfd = optflow_map.copy() @@ -252,8 +252,8 @@ def infer(args): score_map, im_info = predict(frame, model, test_transforms) cur_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) cur_gray = cv2.resize(cur_gray, (resize_w, resize_h)) - scoremap = 255 * score_map[:, :, 1] - optflow_map = postprocess(cur_gray, scoremap, prev_gray, prev_cfd, \ + score_map = 255 * score_map[:, :, 1] + optflow_map = postprocess(cur_gray, score_map, prev_gray, prev_cfd, \ disflow, is_init) prev_gray = cur_gray.copy() prev_cfd = optflow_map.copy() diff --git a/contrib/HumanSeg/video_infer.py b/contrib/HumanSeg/video_infer.py index f37809366d06ee4553bb10d848b3e387a18797bd..63a31c30c7a0de649b453e3368f22c0a9b302d7d 100644 --- a/contrib/HumanSeg/video_infer.py +++ b/contrib/HumanSeg/video_infer.py @@ -22,6 +22,7 @@ import numpy as np from utils.humanseg_postprocess import postprocess, threshold_mask import models import transforms +import time def parse_args(): @@ -109,7 +110,7 @@ def video_infer(args): fps = cap.get(cv2.CAP_PROP_FPS) if args.video_path: - + print('Please waite. It is computing......') # 用于保存预测结果视频 if not osp.exists(args.save_dir): os.makedirs(args.save_dir) @@ -123,8 +124,8 @@ def video_infer(args): score_map, im_info = predict(frame, model, test_transforms) cur_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) cur_gray = cv2.resize(cur_gray, (resize_w, resize_h)) - scoremap = 255 * score_map[:, :, 1] - optflow_map = postprocess(cur_gray, scoremap, prev_gray, prev_cfd, \ + score_map = 255 * score_map[:, :, 1] + optflow_map = postprocess(cur_gray, score_map, prev_gray, prev_cfd, \ disflow, is_init) prev_gray = cur_gray.copy() prev_cfd = optflow_map.copy() @@ -132,10 +133,11 @@ def video_infer(args): optflow_map = cv2.GaussianBlur(optflow_map, (3, 3), 0) optflow_map = threshold_mask( optflow_map, thresh_bg=0.2, thresh_fg=0.8) - img_mat = np.repeat(optflow_map[:, :, np.newaxis], 3, axis=2) - img_mat = recover(img_mat, im_info) - bg_im = np.ones_like(img_mat) * 255 - comb = (img_mat * frame + (1 - img_mat) * bg_im).astype( + img_matting = np.repeat( + optflow_map[:, :, np.newaxis], 3, axis=2) + img_matting = recover(img_matting, im_info) + bg_im = np.ones_like(img_matting) * 255 + comb = (img_matting * frame + (1 - img_matting) * bg_im).astype( np.uint8) out.write(comb) else: @@ -150,20 +152,20 @@ def video_infer(args): score_map, im_info = predict(frame, model, test_transforms) cur_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) cur_gray = cv2.resize(cur_gray, (resize_w, resize_h)) - scoremap = 255 * score_map[:, :, 1] - optflow_map = postprocess(cur_gray, scoremap, prev_gray, prev_cfd, \ + score_map = 255 * score_map[:, :, 1] + optflow_map = postprocess(cur_gray, score_map, prev_gray, prev_cfd, \ disflow, is_init) prev_gray = cur_gray.copy() prev_cfd = optflow_map.copy() is_init = False - # optflow_map = optflow_map/255.0 optflow_map = cv2.GaussianBlur(optflow_map, (3, 3), 0) optflow_map = threshold_mask( optflow_map, thresh_bg=0.2, thresh_fg=0.8) - img_mat = np.repeat(optflow_map[:, :, np.newaxis], 3, axis=2) - img_mat = recover(img_mat, im_info) - bg_im = np.ones_like(img_mat) * 255 - comb = (img_mat * frame + (1 - img_mat) * bg_im).astype( + img_matting = np.repeat( + optflow_map[:, :, np.newaxis], 3, axis=2) + img_matting = recover(img_matting, im_info) + bg_im = np.ones_like(img_matting) * 255 + comb = (img_matting * frame + (1 - img_matting) * bg_im).astype( np.uint8) cv2.imshow('HumanSegmentation', comb) if cv2.waitKey(1) & 0xFF == ord('q'):