def detect(img, network):
orig_img = img.copy()
orig_img_h, orig_img_w, _ = orig_img.shape
input_w, input_h = compute_optimal_size(
orig_img, params['inference_img_size']) # 368
# map_w, map_h = compute_optimal_size(orig_img, params['heatmap_size']) # 320
map_w, map_h = compute_optimal_size(orig_img, params['inference_img_size'])
# print("image size is: ", input_w, input_h)
resized_image = cv2.resize(orig_img, (input_w, input_h))
x_data = preprocess(resized_image)
x_data = Tensor(x_data, mstype.float32)
x_data.requires_grad = False
logit_pafs, logit_heatmap = network(x_data)
logit_pafs = logit_pafs[-1].asnumpy()[0]
logit_heatmap = logit_heatmap[-1].asnumpy()[0]
pafs = np.zeros((logit_pafs.shape[0], map_h, map_w))
for i in range(logit_pafs.shape[0]):
pafs[i] = cv2.resize(logit_pafs[i], (map_w, map_h))
if show_gt:
save_path = "./test_output/" + str(i) + "pafs.png"
cv2.imwrite(save_path, pafs[i] * 255)
heatmaps = np.zeros((logit_heatmap.shape[0], map_h, map_w))
for i in range(logit_heatmap.shape[0]):
heatmaps[i] = cv2.resize(logit_heatmap[i], (map_w, map_h))
if show_gt:
save_path = "./test_output/" + str(i) + "heatmap.png"
cv2.imwrite(save_path, heatmaps[i] * 255)
all_peaks = compute_peaks_from_heatmaps(heatmaps)
if all_peaks.shape[0] == 0:
return np.empty((0, len(JointType), 3)), np.empty(0)
all_connections = compute_connections(pafs, all_peaks, map_w, params)
subsets = grouping_key_points(all_connections, all_peaks, params)
all_peaks[:, 1] *= orig_img_w / map_w
all_peaks[:, 2] *= orig_img_h / map_h
poses = subsets_to_pose_array(subsets, all_peaks)
scores = subsets[:, -2]
return poses, scores
