| 13 | from . import util |
| 14 | |
| 15 | class Hand(object): |
| 16 | def __init__(self, model_path): |
| 17 | self.model = handpose_model() |
| 18 | if torch.cuda.is_available(): |
| 19 | self.model = self.model.cuda() |
| 20 | print('cuda') |
| 21 | model_dict = util.transfer(self.model, torch.load(model_path)) |
| 22 | self.model.load_state_dict(model_dict) |
| 23 | self.model.eval() |
| 24 | |
| 25 | def __call__(self, oriImg): |
| 26 | scale_search = [0.5, 1.0, 1.5, 2.0] |
| 27 | # scale_search = [0.5] |
| 28 | boxsize = 368 |
| 29 | stride = 8 |
| 30 | padValue = 128 |
| 31 | thre = 0.05 |
| 32 | multiplier = [x * boxsize / oriImg.shape[0] for x in scale_search] |
| 33 | heatmap_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 22)) |
| 34 | # paf_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 38)) |
| 35 | |
| 36 | for m in range(len(multiplier)): |
| 37 | scale = multiplier[m] |
| 38 | imageToTest = cv2.resize(oriImg, (0, 0), fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC) |
| 39 | imageToTest_padded, pad = util.padRightDownCorner(imageToTest, stride, padValue) |
| 40 | im = np.transpose(np.float32(imageToTest_padded[:, :, :, np.newaxis]), (3, 2, 0, 1)) / 256 - 0.5 |
| 41 | im = np.ascontiguousarray(im) |
| 42 | |
| 43 | data = torch.from_numpy(im).float() |
| 44 | if torch.cuda.is_available(): |
| 45 | data = data.cuda() |
| 46 | # data = data.permute([2, 0, 1]).unsqueeze(0).float() |
| 47 | with torch.no_grad(): |
| 48 | output = self.model(data).cpu().numpy() |
| 49 | # output = self.model(data).numpy()q |
| 50 | |
| 51 | # extract outputs, resize, and remove padding |
| 52 | heatmap = np.transpose(np.squeeze(output), (1, 2, 0)) # output 1 is heatmaps |
| 53 | heatmap = cv2.resize(heatmap, (0, 0), fx=stride, fy=stride, interpolation=cv2.INTER_CUBIC) |
| 54 | heatmap = heatmap[:imageToTest_padded.shape[0] - pad[2], :imageToTest_padded.shape[1] - pad[3], :] |
| 55 | heatmap = cv2.resize(heatmap, (oriImg.shape[1], oriImg.shape[0]), interpolation=cv2.INTER_CUBIC) |
| 56 | |
| 57 | heatmap_avg += heatmap / len(multiplier) |
| 58 | |
| 59 | all_peaks = [] |
| 60 | for part in range(21): |
| 61 | map_ori = heatmap_avg[:, :, part] |
| 62 | one_heatmap = gaussian_filter(map_ori, sigma=3) |
| 63 | binary = np.ascontiguousarray(one_heatmap > thre, dtype=np.uint8) |
| 64 | # 全部小于阈值 |
| 65 | if np.sum(binary) == 0: |
| 66 | all_peaks.append([0, 0]) |
| 67 | continue |
| 68 | label_img, label_numbers = label(binary, return_num=True, connectivity=binary.ndim) |
| 69 | max_index = np.argmax([np.sum(map_ori[label_img == i]) for i in range(1, label_numbers + 1)]) + 1 |
| 70 | label_img[label_img != max_index] = 0 |
| 71 | map_ori[label_img == 0] = 0 |
| 72 | |