# Modified from: # https://github.com/anibali/pytorch-stacked-hourglass # https://github.com/bearpaw/pytorch-pose import matplotlib as mpl mpl.use('Agg') import matplotlib.pyplot as plt import numpy as np import cv2 import torch # import stacked_hourglass.datasets.utils_stanext as utils_stanext # COLORS, labels = utils_stanext.load_keypoint_labels_and_colours() COLORS = ['#d82400', '#d82400', '#d82400', '#fcfc00', '#fcfc00', '#fcfc00', '#48b455', '#48b455', '#48b455', '#0090aa', '#0090aa', '#0090aa', '#d848ff', '#d848ff', '#fc90aa', '#006caa', '#d89000', '#d89000', '#fc90aa', '#006caa', '#ededed', '#ededed', '#a9d08e', '#a9d08e'] RGB_MEAN = [0.4404, 0.4440, 0.4327] RGB_STD = [0.2458, 0.2410, 0.2468] def get_img_from_fig(fig, dpi=180): buf = io.BytesIO() fig.savefig(buf, format="png", dpi=dpi) buf.seek(0) img_arr = np.frombuffer(buf.getvalue(), dtype=np.uint8) buf.close() img = cv2.imdecode(img_arr, 1) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) return img def save_input_image_with_keypoints(img, tpts, out_path='./test_input.png', colors=COLORS, rgb_mean=RGB_MEAN, rgb_std=RGB_STD, ratio_in_out=4., threshold=0.3, print_scores=False): """ img has shape (3, 256, 256) and is a torch tensor pts has shape (20, 3) and is a torch tensor -> this function is tested with the mpii dataset and the results look ok """ # reverse color normalization for t, m, s in zip(img, rgb_mean, rgb_std): t.add_(m) # inverse to transforms.color_normalize() img_np = img.detach().cpu().numpy().transpose(1, 2, 0) # tpts_np = tpts.detach().cpu().numpy() # plot image fig, ax = plt.subplots() plt.imshow(img_np) # plt.imshow(im) plt.gca().set_axis_off() plt.subplots_adjust(top = 1, bottom = 0, right = 1, left = 0, hspace = 0, wspace = 0) plt.margins(0,0) # plot all visible keypoints #import pdb; pdb.set_trace() for idx, (x, y, v) in enumerate(tpts): if v > threshold: x = int(x*ratio_in_out) y = int(y*ratio_in_out) plt.scatter([x], [y], c=[colors[idx]], marker="x", s=50) if print_scores: txt = '{:2.2f}'.format(v.item()) plt.annotate(txt, (x, y)) # , c=colors[idx]) plt.savefig(out_path, bbox_inches='tight', pad_inches=0) plt.close() return def save_input_image(img, out_path, colors=COLORS, rgb_mean=RGB_MEAN, rgb_std=RGB_STD): for t, m, s in zip(img, rgb_mean, rgb_std): t.add_(m) # inverse to transforms.color_normalize() img_np = img.detach().cpu().numpy().transpose(1, 2, 0) plt.imsave(out_path, img_np) return ###################################################################### def get_bodypart_colors(): # body colors n_body = 8 c = np.arange(1, n_body + 1) norm = mpl.colors.Normalize(vmin=c.min(), vmax=c.max()) cmap = mpl.cm.ScalarMappable(norm=norm, cmap=mpl.cm.gist_rainbow) cmap.set_array([]) body_cols = [] for i in range(0, n_body): body_cols.append(cmap.to_rgba(i + 1)) # head colors n_blue = 5 c = np.arange(1, n_blue + 1) norm = mpl.colors.Normalize(vmin=c.min()-1, vmax=c.max()+1) cmap = mpl.cm.ScalarMappable(norm=norm, cmap=mpl.cm.Blues) cmap.set_array([]) head_cols = [] for i in range(0, n_body): head_cols.append(cmap.to_rgba(i + 1)) # torso colors n_blue = 2 c = np.arange(1, n_blue + 1) norm = mpl.colors.Normalize(vmin=c.min()-1, vmax=c.max()+1) cmap = mpl.cm.ScalarMappable(norm=norm, cmap=mpl.cm.Greens) cmap.set_array([]) torso_cols = [] for i in range(0, n_body): torso_cols.append(cmap.to_rgba(i + 1)) return body_cols, head_cols, torso_cols body_cols, head_cols, torso_cols = get_bodypart_colors() tbp_dict = {'full_body': [0, 8], 'head': [8, 13], 'torso': [13, 15]} def save_image_with_part_segmentation(partseg_big, seg_big, input_image_np, ind_img, out_path_seg=None, out_path_seg_overlay=None, thr=0.3): soft_max = torch.nn.Softmax(dim=0) # create dit with results tbp_dict_res = {} for ind_tbp, part in enumerate(['full_body', 'head', 'torso']): partseg_tbp = partseg_big[:, tbp_dict[part][0]:tbp_dict[part][1], :, :] segm_img_pred = soft_max((partseg_tbp[ind_img, :, :, :])) # [1, :, :] m_v, m_i = segm_img_pred.max(axis=0) tbp_dict_res[part] = { 'inds': tbp_dict[part], 'seg_probs': segm_img_pred, 'seg_max_inds': m_i, 'seg_max_values': m_v} # create output_image partseg_image = np.zeros((256, 256, 3)) for ind_sp in range(0, 5): # partseg_image[tbp_dict_res['head']['seg_max_inds']==ind_sp, :] = head_cols[ind_sp][0:3] mask_a = tbp_dict_res['full_body']['seg_max_inds']==1 mask_b = tbp_dict_res['head']['seg_max_inds']==ind_sp partseg_image[mask_a*mask_b, :] = head_cols[ind_sp][0:3] for ind_sp in range(0, 2): # partseg_image[tbp_dict_res['torso']['seg_max_inds']==ind_sp, :] = torso_cols[ind_sp][0:3] mask_a = tbp_dict_res['full_body']['seg_max_inds']==2 mask_b = tbp_dict_res['torso']['seg_max_inds']==ind_sp partseg_image[mask_a*mask_b, :] = torso_cols[ind_sp][0:3] for ind_sp in range(0, 8): if (not ind_sp == 1) and (not ind_sp == 2): # head and torso partseg_image[tbp_dict_res['full_body']['seg_max_inds']==ind_sp, :] = body_cols[ind_sp][0:3] partseg_image[soft_max((seg_big[ind_img, :, :, :]))[1, :, :]