Reshape displacements and offsets array for cleaner access, few lines less code, sligh perf increase

Author: rwightman

posenet/decode.py

@@ -3,47 +3,31 @@
 from posenet.constants import *
 
 
-def get_offset_point(coord, keypoint_id, offsets):
-    return np.array((
-        offsets[coord[0], coord[1], keypoint_id],
-        offsets[coord[0], coord[1], keypoint_id + NUM_KEYPOINTS])).astype(np.int32)
-
-
-def get_image_coords(heatmap_coord, keypoint_id, output_stride, offsets):
-    return heatmap_coord * output_stride + get_offset_point(heatmap_coord, keypoint_id, offsets)
-
-
 def traverse_to_targ_keypoint(
         edge_id, source_keypoint, target_keypoint_id, scores, offsets, output_stride, displacements
 ):
     height = scores.shape[0]
     width = scores.shape[1]
-    num_edges = displacements.shape[2] // 2
 
     source_keypoint_indices = np.clip(
         np.round(source_keypoint / output_stride), a_min=0, a_max=[height - 1, width - 1]).astype(np.int32)
 
-    displacement = np.array((
-        displacements[source_keypoint_indices[0], source_keypoint_indices[1], edge_id],
-        displacements[source_keypoint_indices[0], source_keypoint_indices[1], edge_id + num_edges]
-    ))
-
-    displaced_point = source_keypoint + displacement
+    displaced_point = source_keypoint + displacements[
+        source_keypoint_indices[0], source_keypoint_indices[1], edge_id]
 
     displaced_point_indices = np.clip(
         np.round(displaced_point / output_stride), a_min=0, a_max=[height - 1, width - 1]).astype(np.int32)
 
-    offset_point = get_offset_point(displaced_point_indices, target_keypoint_id, offsets)
-
     score = scores[displaced_point_indices[0], displaced_point_indices[1], target_keypoint_id]
 
-    position = displaced_point_indices * output_stride + offset_point
+    image_coord = displaced_point_indices * output_stride + offsets[
+        displaced_point_indices[0], displaced_point_indices[1], target_keypoint_id]
 
-    return score, position
+    return score, image_coord
 
 
 def decode_pose(
-        root_score, root_id, root_coord,
+        root_score, root_id, root_image_coord,
         scores,
         offsets,
         output_stride,
@@ -55,11 +39,8 @@ def decode_pose(
 
     instance_keypoint_scores = np.zeros(num_parts)
     instance_keypoint_coords = np.zeros((num_parts, 2))
-
-    root_point = get_image_coords(root_coord, root_id, output_stride, offsets)
-
     instance_keypoint_scores[root_id] = root_score
-    instance_keypoint_coords[root_id] = root_point
+    instance_keypoint_coords[root_id] = root_image_coord
 
     # FIXME can we vectorize these loops cleanly?
     for edge in reversed(range(num_edges)):

posenet/decode_multi.py

@@ -126,15 +126,22 @@ def decode_multiple_poses(
     scored_parts = build_part_with_score_fast(score_threshold, LOCAL_MAXIMUM_RADIUS, scores)
     scored_parts = sorted(scored_parts, key=lambda x: x[0], reverse=True)
 
+    # change dimensions from (h, w, x) to (h, w, x//2, 2) to allow return of complete coord array
+    height = scores.shape[0]
+    width = scores.shape[1]
+    offsets = offsets.reshape(height, width, 2, -1).swapaxes(2, 3)
+    displacements_fwd = displacements_fwd.reshape(height, width, 2, -1).swapaxes(2, 3)
+    displacements_bwd = displacements_bwd.reshape(height, width, 2, -1).swapaxes(2, 3)
+
     for root_score, root_id, root_coord in scored_parts:
-        root_image_coords = get_image_coords(root_coord, root_id, output_stride, offsets)
+        root_image_coords = root_coord * output_stride + offsets[root_coord[0], root_coord[1], root_id]
 
         if within_nms_radius_fast(
                 pose_keypoint_coords[:pose_count, root_id, :], squared_nms_radius, root_image_coords):
             continue
 
         keypoint_scores, keypoint_coords = decode_pose(
-            root_score, root_id, root_coord,
+            root_score, root_id, root_image_coords,
             scores, offsets, output_stride,
             displacements_fwd, displacements_bwd)