update humanseg_postprocess

contrib/HumanSeg/utils/humanseg_postprocess.py

@@ -14,13 +14,6 @@
 # limitations under the License.
 
 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 human_seg_tracking(pre_gray, cur_gray, prev_cfd, dl_weights, disflow):
@@ -39,28 +32,32 @@ def human_seg_tracking(pre_gray, cur_gray, prev_cfd, dl_weights, disflow):
     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)
-    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 + c
-            dy_fw = get_round(fxy_fw[1])
-            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])
-            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[r, c]
+    # cur_position =
+    flow_fw = np.round(flow_fw).astype(np.int)
+    flow_bw = np.round(flow_bw).astype(np.int)
+    y_list = np.array(range(h))
+    x_list = np.array(range(w))
+    yv, xv = np.meshgrid(y_list, x_list)
+    yv, xv = yv.T, xv.T
+    cur_x = xv + flow_fw[:, :, 0]
+    cur_y = yv + flow_fw[:, :, 1]
+
+    # 超出边界不跟踪
+    not_track = (cur_x < 0) + (cur_x >= w) + (cur_y < 0) + (cur_y >= h)
+    flow_bw[~not_track] = flow_bw[cur_y[~not_track], cur_x[~not_track]]
+    not_track += (np.square(flow_fw[:, :, 0] + flow_bw[:, :, 0]) +
+                  np.square(flow_fw[:, :, 1] + flow_bw[:, :, 1])) >= check_thres
+    track_cfd[cur_y[~not_track], cur_x[~not_track]] = prev_cfd[~not_track]
+
+    is_track[cur_y[~not_track], cur_x[~not_track]] = 1
+
+    not_flow = np.all(
+        np.abs(flow_fw) == 0, axis=-1) * np.all(
+            np.abs(flow_bw) == 0, axis=-1)
+    dl_weights[cur_y[not_flow], cur_x[not_flow]] = 0.05
     return track_cfd, is_track, dl_weights
 
 
@@ -75,24 +72,27 @@ def human_seg_track_fuse(track_cfd, dl_cfd, dl_weights, is_track):
         cur_cfd: 光流跟踪图和人像分割结果融合图
     """
     fusion_cfd = dl_cfd.copy()
-    idxs = np.where(is_track > 0)
-    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:
-                fusion_cfd[x, y] = 0.3 * dl_score + 0.7 * track_score
-            else:
-                fusion_cfd[x, y] = 0.4 * dl_score + 0.6 * track_score
-        else:
-            fusion_cfd[x, y] = dl_weights[x, y] * dl_score + (
-                1 - dl_weights[x, y]) * track_score
+    is_track = is_track.astype(np.bool)
+    fusion_cfd[is_track] = dl_weights[is_track] * dl_cfd[is_track] + (
+        1 - dl_weights[is_track]) * track_cfd[is_track]
+    # 确定区域
+    index_certain = ((dl_cfd > 0.9) + (dl_cfd < 0.1)) * is_track
+    index_less01 = (dl_weights < 0.1) * index_certain
+    fusion_cfd[index_less01] = 0.3 * dl_cfd[index_less01] + 0.7 * track_cfd[
+        index_less01]
+    index_larger09 = (dl_weights >= 0.1) * index_certain
+    fusion_cfd[index_larger09] = 0.4 * dl_cfd[index_larger09] + 0.6 * track_cfd[
+        index_larger09]
     return fusion_cfd
 
 
+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)
+
+
 def postprocess(cur_gray, scoremap, prev_gray, pre_cfd, disflow, is_init):
     """光流优化
     Args:
@@ -105,8 +105,6 @@ def postprocess(cur_gray, scoremap, prev_gray, pre_cfd, disflow, is_init):
     Returns:
         fusion_cfd : 光流追踪图和预测结果融合图
     """
-    height, width = scoremap.shape[0], scoremap.shape[1]
-    disflow = cv2.DISOpticalFlow_create(cv2.DISOPTICAL_FLOW_PRESET_ULTRAFAST)
     h, w = scoremap.shape
     cur_cfd = scoremap.copy()
 
@@ -120,18 +118,9 @@ def postprocess(cur_gray, scoremap, prev_gray, pre_cfd, disflow, is_init):
             disflow.setFinestScale(3)
         fusion_cfd = cur_cfd
     else:
-        weights = np.ones((w, h), np.float32) * 0.3
+        weights = np.ones((h, w), 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 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)