def _get_outputs(multiplier, img, heatmaps, pafs, numkeypoints=CF.NUM_KEYPOINTS, numlims=CF.NUM_LIMBS): """Computes the averaged heatmap and paf for the given image :param multiplier: :param origImg: numpy array, the image being processed :param model: pytorch model :returns: numpy arrays, the averaged paf and heatmap """ heatmap_avg = np.zeros((img.shape[0], img.shape[1], numkeypoints + 1)) paf_avg = np.zeros((img.shape[0], img.shape[1], numlims * 2)) heatmaps = heatmaps.transpose(0, 2, 3, 1) pafs = pafs.transpose(0, 2, 3, 1) # print('len(multiplier)', len(multiplier)) for m in range(len(multiplier)): scale = multiplier[m] inp_size = scale * img.shape[0] # padding im_cropped, im_scale, real_shape = im_transform.crop_with_factor( img, inp_size, factor=8, is_ceil=True) # print('im_cropped', im_cropped.shape, heatmaps.shape) heatmap = heatmaps[m, :int(im_cropped.shape[0] / 8), :int(im_cropped.shape[1] / 8), :] heatmap = cv2.resize(heatmap, None, fx=8, fy=8, interpolation=cv2.INTER_CUBIC) heatmap = heatmap[0:real_shape[0], 0:real_shape[1], :] heatmap = cv2.resize(heatmap, (img.shape[1], img.shape[0]), interpolation=cv2.INTER_CUBIC) paf = pafs[m, :int(im_cropped.shape[0] / 8), :int(im_cropped.shape[1] / 8), :] paf = cv2.resize(paf, None, fx=8, fy=8, interpolation=cv2.INTER_CUBIC) paf = paf[0:real_shape[0], 0:real_shape[1], :] paf = cv2.resize(paf, (img.shape[1], img.shape[0]), interpolation=cv2.INTER_CUBIC) heatmap_avg = heatmap_avg + heatmap / len(multiplier) paf_avg = paf_avg + paf / len(multiplier) return paf_avg, heatmap_avg
def get_outputs(multiplier, img, model, preprocess): """Computes the averaged heatmap and paf for the given image :param multiplier: :param origImg: numpy array, the image being processed :param model: pytorch model :returns: numpy arrays, the averaged paf and heatmap """ heatmap_avg = np.zeros((img.shape[0], img.shape[1], 19)) paf_avg = np.zeros((img.shape[0], img.shape[1], 38)) max_scale = multiplier[-1] max_size = max_scale * img.shape[0] # padding max_cropped, _, _ = im_transform.crop_with_factor(img, max_size, factor=8, is_ceil=True) batch_images = np.zeros( (len(multiplier), 3, max_cropped.shape[0], max_cropped.shape[1])) for m in range(len(multiplier)): scale = multiplier[m] inp_size = scale * img.shape[0] # padding im_croped, im_scale, real_shape = im_transform.crop_with_factor( img, inp_size, factor=8, is_ceil=True) if preprocess == 'rtpose': im_data = rtpose_preprocess(im_croped) elif preprocess == 'vgg': im_data = vgg_preprocess(im_croped) elif preprocess == 'inception': im_data = inception_preprocess(im_croped) elif preprocess == 'ssd': im_data = ssd_preprocess(im_croped) batch_images[m, :, :im_data.shape[1], :im_data.shape[2]] = im_data # several scales as a batch batch_var = torch.from_numpy(batch_images).cuda().float() predicted_outputs, _ = model(batch_var) output1, output2 = predicted_outputs[-2], predicted_outputs[-1] heatmaps = output2.cpu().data.numpy().transpose(0, 2, 3, 1) pafs = output1.cpu().data.numpy().transpose(0, 2, 3, 1) for m in range(len(multiplier)): scale = multiplier[m] inp_size = scale * img.shape[0] # padding im_cropped, im_scale, real_shape = im_transform.crop_with_factor( img, inp_size, factor=8, is_ceil=True) heatmap = heatmaps[m, :int(im_cropped.shape[0] / 8), :int(im_cropped.shape[1] / 8), :] heatmap = cv2.resize(heatmap, None, fx=8, fy=8, interpolation=cv2.INTER_CUBIC) heatmap = heatmap[0:real_shape[0], 0:real_shape[1], :] heatmap = cv2.resize(heatmap, (img.shape[1], img.shape[0]), interpolation=cv2.INTER_CUBIC) paf = pafs[m, :int(im_cropped.shape[0] / 8), :int(im_cropped.shape[1] / 8), :] paf = cv2.resize(paf, None, fx=8, fy=8, interpolation=cv2.INTER_CUBIC) paf = paf[0:real_shape[0], 0:real_shape[1], :] paf = cv2.resize(paf, (img.shape[1], img.shape[0]), interpolation=cv2.INTER_CUBIC) heatmap_avg = heatmap_avg + heatmap / len(multiplier) paf_avg = paf_avg + paf / len(multiplier) return paf_avg, heatmap_avg
feature_extractor5_1 = FeatureExtractor(model.model5_1) feature_extractor5_2 = FeatureExtractor(model.model5_2) feature_extractor6_1 = FeatureExtractor(model.model6_1) feature_extractor6_2 = FeatureExtractor(model.model6_2) images = os.listdir(image_dir) # iterate all val images feed_size = 368 for a_path in images: one_path = image_dir + a_path print(one_path) oriImg = cv2.imread(one_path) im_croped, im_scale, real_shape = im_transform.crop_with_factor( oriImg, feed_size, factor=8, is_ceil=True) print('size of im crop', im_croped.shape) im_data = vgg_preprocess(im_croped) #im_data = im_data.transpose([2, 0, 1]).astype(np.float32) batch_images = np.expand_dims(im_data, 0) batch_var = torch.from_numpy(batch_images).cuda().float() all_outputs = [] outputs0, out1 = feature_extractor(batch_var) outputs1_1, out1_1 = feature_extractor1_1(out1)