From 3402f5b19be221f55da17ceeca346f3a229d801f Mon Sep 17 00:00:00 2001 From: chenguowei01 Date: Tue, 12 May 2020 17:54:34 +0800 Subject: [PATCH] updata optflow --- .../HumanSeg/utils/humanseg_postprocess.py | 133 +++++++----------- contrib/HumanSeg/video_infer.py | 51 +++++-- 2 files changed, 89 insertions(+), 95 deletions(-) diff --git a/contrib/HumanSeg/utils/humanseg_postprocess.py b/contrib/HumanSeg/utils/humanseg_postprocess.py index 6dd73eae..44376726 100644 --- a/contrib/HumanSeg/utils/humanseg_postprocess.py +++ b/contrib/HumanSeg/utils/humanseg_postprocess.py @@ -1,24 +1,14 @@ -# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== - -import os import numpy as np import cv2 +import os + +def get_round(data): + round = 0.5 if data >= 0 else -0.5 + return (int)(data + round) -def humanseg_tracking(pre_gray, cur_gray, prev_cfd, dl_weights, disflow): + +def human_seg_tracking(pre_gray, cur_gray, prev_cfd, dl_weights, disflow): """计算光流跟踪匹配点和光流图 输入参数: pre_gray: 上一帧灰度图 @@ -31,131 +21,102 @@ def humanseg_tracking(pre_gray, cur_gray, prev_cfd, dl_weights, disflow): track_cfd: 光流跟踪图 """ check_thres = 8 - hgt, wdh = pre_gray.shape[:2] + h, w = pre_gray.shape[:2] track_cfd = np.zeros_like(prev_cfd) is_track = np.zeros_like(pre_gray) - # 计算前向光流 flow_fw = disflow.calc(pre_gray, cur_gray, None) - # 计算后向光流 flow_bw = disflow.calc(cur_gray, pre_gray, None) - get_round = lambda data: (int)(data + 0.5) if data >= 0 else (int)(data - - 0.5) - for row in range(hgt): - for col in range(wdh): - # 计算光流处理后对应点坐标 - # (row, col) -> (cur_x, cur_y) - fxy_fw = flow_fw[row, col] + for r in range(h): + for c in range(w): + fxy_fw = flow_fw[r, c] dx_fw = get_round(fxy_fw[0]) - cur_x = dx_fw + col + cur_x = dx_fw + c dy_fw = get_round(fxy_fw[1]) - cur_y = dy_fw + row - if cur_x < 0 or cur_x >= wdh or cur_y < 0 or cur_y >= hgt: + cur_y = dy_fw + r + if cur_x < 0 or cur_x >= w or cur_y < 0 or cur_y >= h: continue fxy_bw = flow_bw[cur_y, cur_x] dx_bw = get_round(fxy_bw[0]) dy_bw = get_round(fxy_bw[1]) - # 光流移动小于阈值 - lmt = ((dy_fw + dy_bw) * (dy_fw + dy_bw) + - (dx_fw + dx_bw) * (dx_fw + dx_bw)) - if lmt >= check_thres: + if ((dy_fw + dy_bw) * (dy_fw + dy_bw) + + (dx_fw + dx_bw) * (dx_fw + dx_bw)) >= check_thres: continue - # 静止点降权 if abs(dy_fw) <= 0 and abs(dx_fw) <= 0 and abs(dy_bw) <= 0 and abs( dx_bw) <= 0: dl_weights[cur_y, cur_x] = 0.05 is_track[cur_y, cur_x] = 1 - track_cfd[cur_y, cur_x] = prev_cfd[row, col] + track_cfd[cur_y, cur_x] = prev_cfd[r, c] return track_cfd, is_track, dl_weights -def humanseg_track_fuse(track_cfd, dl_cfd, dl_weights, is_track): +def human_seg_track_fuse(track_cfd, dl_cfd, dl_weights, is_track): """光流追踪图和人像分割结构融合 输入参数: track_cfd: 光流追踪图 dl_cfd: 当前帧分割结果 dl_weights: 融合权重图 is_track: 光流点匹配二值图 - 返回值: + 返回 cur_cfd: 光流跟踪图和人像分割结果融合图 """ - cur_cfd = dl_cfd.copy() + fusion_cfd = dl_cfd.copy() idxs = np.where(is_track > 0) - for i in range(len(idxs)): + for i in range(len(idxs[0])): x, y = idxs[0][i], idxs[1][i] dl_score = dl_cfd[x, y] track_score = track_cfd[x, y] + fusion_cfd[x, y] = dl_weights[x, y] * dl_score + ( + 1 - dl_weights[x, y]) * track_score if dl_score > 0.9 or dl_score < 0.1: if dl_weights[x, y] < 0.1: - cur_cfd[x, y] = 0.3 * dl_score + 0.7 * track_score + fusion_cfd[x, y] = 0.3 * dl_score + 0.7 * track_score else: - cur_cfd[x, y] = 0.4 * dl_score + 0.6 * track_score + fusion_cfd[x, y] = 0.4 * dl_score + 0.6 * track_score else: - cur_cfd[x, y] = dl_weights[x, y] * dl_score + ( + fusion_cfd[x, y] = dl_weights[x, y] * dl_score + ( 1 - dl_weights[x, y]) * track_score - return cur_cfd - - -def threshold_mask(img, thresh_bg, thresh_fg): - """设置背景和前景阈值mask - 输入参数: - img : 原始图像, np.uint8 类型. - thresh_bg : 背景阈值百分比，低于该值置为0. - thresh_fg : 前景阈值百分比，超过该值置为1. - 返回值: - dst : 原始图像设置完前景背景阈值mask结果, np.float32 类型. - """ - dst = (img / 255.0 - thresh_bg) / (thresh_fg - thresh_bg) - dst[np.where(dst > 1)] = 1 - dst[np.where(dst < 0)] = 0 - return dst.astype(np.float32) + return fusion_cfd -def optflow_handle(cur_gray, scoremap, is_init): +def postprocess(cur_gray, scoremap, prev_gray, pre_cfd, disflow, is_init): """光流优化 Args: cur_gray : 当前帧灰度图 + pre_gray : 前一帧灰度图 + pre_cfd ：前一帧融合结果 scoremap : 当前帧分割结果 + difflow : 光流 is_init : 是否第一帧 Returns: - dst : 光流追踪图和预测结果融合图, 类型为 np.float32 + fusion_cfd : 光流追踪图和预测结果融合图 """ height, width = scoremap.shape[0], scoremap.shape[1] disflow = cv2.DISOpticalFlow_create(cv2.DISOPTICAL_FLOW_PRESET_ULTRAFAST) - prev_gray = np.zeros((height, width), np.uint8) - prev_cfd = np.zeros((height, width), np.float32) + h, w = scoremap.shape cur_cfd = scoremap.copy() + if is_init: is_init = False - if height <= 64 or width <= 64: + if h <= 64 or w <= 64: disflow.setFinestScale(1) - elif height <= 160 or width <= 160: + elif h <= 160 or w <= 160: disflow.setFinestScale(2) else: disflow.setFinestScale(3) fusion_cfd = cur_cfd else: - weights = np.ones((height, width), np.float32) * 0.3 - track_cfd, is_track, weights = humanseg_tracking( - prev_gray, cur_gray, prev_cfd, weights, disflow) - fusion_cfd = humanseg_track_fuse(track_cfd, cur_cfd, weights, is_track) + weights = np.ones((w, h), np.float32) * 0.3 + track_cfd, is_track, weights = human_seg_tracking( + prev_gray, cur_gray, pre_cfd, weights, disflow) + fusion_cfd = human_seg_track_fuse(track_cfd, cur_cfd, weights, is_track) + fusion_cfd = cv2.GaussianBlur(fusion_cfd, (3, 3), 0) + return fusion_cfd -def postprocess(image, output_data): - """对预测结果进行后处理 - Args: - image: 原始图，opencv 图片对象 - output_data: Paddle预测结果原始数据 - Returns: - 原图和预测结果融合并做了光流优化的结果图 - """ - scoremap = output_data[:, :, 1] - scoremap = (scoremap * 255).astype(np.uint8) - # 光流处理 - cur_gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) - optflow_map = optflow_handle(cur_gray, scoremap, False) - optflow_map = cv2.GaussianBlur(optflow_map, (3, 3), 0) - optflow_map = threshold_mask(optflow_map, thresh_bg=0.2, thresh_fg=0.8) - optflow_map = np.repeat(optflow_map[:, :, np.newaxis], 3, axis=2) - return optflow_map +def threshold_mask(img, thresh_bg, thresh_fg): + dst = (img / 255.0 - thresh_bg) / (thresh_fg - thresh_bg) + dst[np.where(dst > 1)] = 1 + dst[np.where(dst < 0)] = 0 + return dst.astype(np.float32) diff --git a/contrib/HumanSeg/video_infer.py b/contrib/HumanSeg/video_infer.py index 45a8189e..bc8ee690 100644 --- a/contrib/HumanSeg/video_infer.py +++ b/contrib/HumanSeg/video_infer.py @@ -4,7 +4,7 @@ import os.path as osp import cv2 import numpy as np -from utils.humanseg_postprocess import postprocess +from utils.humanseg_postprocess import postprocess, threshold_mask import models import transforms @@ -60,8 +60,12 @@ def recover(img, im_info): def video_infer(args): + + resize_h = 192 + resize_w = 192 + test_transforms = transforms.Compose( - [transforms.Resize((192, 192)), + [transforms.Resize((resize_w, resize_h)), transforms.Normalize()]) model = models.load_model(args.model_dir) if not args.video_path: @@ -73,10 +77,18 @@ def video_infer(args): "--video_path whether existing: {}" " or camera whether working".format(args.video_path)) return + + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + + disflow = cv2.DISOpticalFlow_create(cv2.DISOPTICAL_FLOW_PRESET_ULTRAFAST) + prev_gray = np.zeros((resize_h, resize_w), np.uint8) + prev_cfd = np.zeros((resize_h, resize_w), np.float32) + is_init = True + + fps = cap.get(cv2.CAP_PROP_FPS) if args.video_path: - width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) - height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) - fps = cap.get(cv2.CAP_PROP_FPS) + # 用于保存预测结果视频 if not osp.exists(args.save_dir): os.makedirs(args.save_dir) @@ -88,8 +100,18 @@ def video_infer(args): ret, frame = cap.read() if ret: score_map, im_info = predict(frame, model, test_transforms) - img = cv2.resize(frame, (192, 192)) - img_mat = postprocess(img, score_map) + 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, \ + disflow, is_init) + prev_gray = cur_gray.copy() + prev_cfd = optflow_map.copy() + is_init = False + 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( @@ -105,8 +127,19 @@ def video_infer(args): ret, frame = cap.read() if ret: score_map, im_info = predict(frame, model, test_transforms) - img = cv2.resize(frame, (192, 192)) - img_mat = postprocess(img, score_map) + 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, \ + 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( -- GitLab