{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"\n\nNOTE: Turn on Internet and GPU","metadata":{"papermill":{"duration":0.016848,"end_time":"2020-09-29T03:23:48.564827","exception":false,"start_time":"2020-09-29T03:23:48.547979","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"The code hidden below handles all the imports and function definitions (the heavy lifting). If you're a beginner I'd advice you skip this for now. When you are able to understand the rest of the code, come back here and understand each function to get a deeper knowledge.","metadata":{"papermill":{"duration":0.014462,"end_time":"2020-09-29T03:23:48.595054","exception":false,"start_time":"2020-09-29T03:23:48.580592","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# !/usr/bin/env python3\n# coding=utf-8\n# author=dave.fang@outlook.com\n# create=20171225\n\nimport os\nimport cv2\nimport sys\nimport math\nimport time\nimport numpy as np\nimport matplotlib.pyplot as plt\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.parallel\nimport torch.backends.cudnn as cudnn\nimport torch.optim as optim\nimport torchvision.transforms as transforms\nimport torchvision.datasets as datasets\nimport torchvision.models as models\n\nfrom torch.autograd import Variable\n\nfrom scipy.ndimage.filters import gaussian_filter\n\n%matplotlib inline\n%config InlineBackend.figure_format = 'retina'\n\n# find connection in the specified sequence, center 29 is in the position 15\nlimb_seq = [[2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], [2, 9], [9, 10],\n            [10, 11], [2, 12], [12, 13], [13, 14], [2, 1], [1, 15], [15, 17],\n            [1, 16], [16, 18], [3, 17], [6, 18]]\n\n# the middle joints heatmap correpondence\nmap_ids = [[31, 32], [39, 40], [33, 34], [35, 36], [41, 42], [43, 44], [19, 20], [21, 22],\n           [23, 24], [25, 26], [27, 28], [29, 30], [47, 48], [49, 50], [53, 54], [51, 52],\n           [55, 56], [37, 38], [45, 46]]\n\n# these are the colours for the 18 body points\ncolors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0],\n          [0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255],\n          [170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]]\n\n\nclass PoseEstimation(nn.Module):\n    def __init__(self, model_dict):\n        super(PoseEstimation, self).__init__()\n\n        self.model0 = model_dict['block_0']\n        self.model1_1 = model_dict['block1_1']\n        self.model2_1 = model_dict['block2_1']\n        self.model3_1 = model_dict['block3_1']\n        self.model4_1 = model_dict['block4_1']\n        self.model5_1 = model_dict['block5_1']\n        self.model6_1 = model_dict['block6_1']\n\n        self.model1_2 = model_dict['block1_2']\n        self.model2_2 = model_dict['block2_2']\n        self.model3_2 = model_dict['block3_2']\n        self.model4_2 = model_dict['block4_2']\n        self.model5_2 = model_dict['block5_2']\n        self.model6_2 = model_dict['block6_2']\n\n    def forward(self, x):\n        out1 = self.model0(x)\n\n        out1_1 = self.model1_1(out1)\n        out1_2 = self.model1_2(out1)\n        out2 = torch.cat([out1_1, out1_2, out1], 1)\n\n        out2_1 = self.model2_1(out2)\n        out2_2 = self.model2_2(out2)\n        out3 = torch.cat([out2_1, out2_2, out1], 1)\n\n        out3_1 = self.model3_1(out3)\n        out3_2 = self.model3_2(out3)\n        out4 = torch.cat([out3_1, out3_2, out1], 1)\n\n        out4_1 = self.model4_1(out4)\n        out4_2 = self.model4_2(out4)\n        out5 = torch.cat([out4_1, out4_2, out1], 1)\n\n        out5_1 = self.model5_1(out5)\n        out5_2 = self.model5_2(out5)\n        out6 = torch.cat([out5_1, out5_2, out1], 1)\n\n        out6_1 = self.model6_1(out6)\n        out6_2 = self.model6_2(out6)\n\n        return out6_1, out6_2\n\n\ndef make_layers(layer_dict):\n    layers = []\n\n    for i in range(len(layer_dict) - 1):\n        layer = layer_dict[i]\n        for k in layer:\n            v = layer[k]\n            if 'pool' in k:\n                layers += [nn.MaxPool2d(kernel_size=v[0], stride=v[1], padding=v[2])]\n            else:\n                conv2d = nn.Conv2d(in_channels=v[0], out_channels=v[1], kernel_size=v[2], stride=v[3], padding=v[4])\n                layers += [conv2d, nn.ReLU(inplace=True)]\n    layer = list(layer_dict[-1].keys())\n    k = layer[0]\n    v = layer_dict[-1][k]\n\n    conv2d = nn.Conv2d(in_channels=v[0], out_channels=v[1], kernel_size=v[2], stride=v[3], padding=v[4])\n    layers += [conv2d]\n\n    return nn.Sequential(*layers)\n\n\ndef get_pose_model():\n    blocks = {}\n\n    block_0 = [{'conv1_1': [3, 64, 3, 1, 1]}, {'conv1_2': [64, 64, 3, 1, 1]}, {'pool1_stage1': [2, 2, 0]},\n               {'conv2_1': [64, 128, 3, 1, 1]}, {'conv2_2': [128, 128, 3, 1, 1]}, {'pool2_stage1': [2, 2, 0]},\n               {'conv3_1': [128, 256, 3, 1, 1]}, {'conv3_2': [256, 256, 3, 1, 1]}, {'conv3_3': [256, 256, 3, 1, 1]},\n               {'conv3_4': [256, 256, 3, 1, 1]}, {'pool3_stage1': [2, 2, 0]}, {'conv4_1': [256, 512, 3, 1, 1]},\n               {'conv4_2': [512, 512, 3, 1, 1]}, {'conv4_3_CPM': [512, 256, 3, 1, 1]},\n               {'conv4_4_CPM': [256, 128, 3, 1, 1]}]\n\n    blocks['block1_1'] = [{'conv5_1_CPM_L1': [128, 128, 3, 1, 1]}, {'conv5_2_CPM_L1': [128, 128, 3, 1, 1]},\n                          {'conv5_3_CPM_L1': [128, 128, 3, 1, 1]}, {'conv5_4_CPM_L1': [128, 512, 1, 1, 0]},\n                          {'conv5_5_CPM_L1': [512, 38, 1, 1, 0]}]\n\n    blocks['block1_2'] = [{'conv5_1_CPM_L2': [128, 128, 3, 1, 1]}, {'conv5_2_CPM_L2': [128, 128, 3, 1, 1]},\n                          {'conv5_3_CPM_L2': [128, 128, 3, 1, 1]}, {'conv5_4_CPM_L2': [128, 512, 1, 1, 0]},\n                          {'conv5_5_CPM_L2': [512, 19, 1, 1, 0]}]\n\n    for i in range(2, 7):\n        blocks['block%d_1' % i] = [{'Mconv1_stage%d_L1' % i: [185, 128, 7, 1, 3]},\n                                   {'Mconv2_stage%d_L1' % i: [128, 128, 7, 1, 3]},\n                                   {'Mconv3_stage%d_L1' % i: [128, 128, 7, 1, 3]},\n                                   {'Mconv4_stage%d_L1' % i: [128, 128, 7, 1, 3]},\n                                   {'Mconv5_stage%d_L1' % i: [128, 128, 7, 1, 3]},\n                                   {'Mconv6_stage%d_L1' % i: [128, 128, 1, 1, 0]},\n                                   {'Mconv7_stage%d_L1' % i: [128, 38, 1, 1, 0]}]\n        blocks['block%d_2' % i] = [{'Mconv1_stage%d_L2' % i: [185, 128, 7, 1, 3]},\n                                   {'Mconv2_stage%d_L2' % i: [128, 128, 7, 1, 3]},\n                                   {'Mconv3_stage%d_L2' % i: [128, 128, 7, 1, 3]},\n                                   {'Mconv4_stage%d_L2' % i: [128, 128, 7, 1, 3]},\n                                   {'Mconv5_stage%d_L2' % i: [128, 128, 7, 1, 3]},\n                                   {'Mconv6_stage%d_L2' % i: [128, 128, 1, 1, 0]},\n                                   {'Mconv7_stage%d_L2' % i: [128, 19, 1, 1, 0]}]\n\n    layers = []\n    for block in block_0:\n        # print(block)\n        for key in block:\n            v = block[key]\n            if 'pool' in key:\n                layers += [nn.MaxPool2d(kernel_size=v[0], stride=v[1], padding=v[2])]\n            else:\n                conv2d = nn.Conv2d(in_channels=v[0], out_channels=v[1], kernel_size=v[2], stride=v[3], padding=v[4])\n                layers += [conv2d, nn.ReLU(inplace=True)]\n\n    models = {\n        'block_0': nn.Sequential(*layers)\n    }\n\n    for k in blocks:\n        v = blocks[k]\n        models[k] = make_layers(v)\n\n    return PoseEstimation(models)\n\n\ndef get_paf_and_heatmap(model, img_raw, scale_search, param_stride=8, box_size=368):\n    multiplier = [scale * box_size / img_raw.shape[0] for scale in scale_search]\n\n    heatmap_avg = torch.zeros((len(multiplier), 19, img_raw.shape[0], img_raw.shape[1])).cuda()\n    paf_avg = torch.zeros((len(multiplier), 38, img_raw.shape[0], img_raw.shape[1])).cuda()\n\n    for i, scale in enumerate(multiplier):\n        img_test = cv2.resize(img_raw, (0, 0), fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)\n        img_test_pad, pad = pad_right_down_corner(img_test, param_stride, param_stride)\n        img_test_pad = np.transpose(np.float32(img_test_pad[:, :, :, np.newaxis]), (3, 2, 0, 1)) / 256 - 0.5\n\n        feed = Variable(torch.from_numpy(img_test_pad)).cuda()\n        output1, output2 = model(feed)\n\n        print(output1.size())\n        print(output2.size())\n\n        heatmap = nn.UpsamplingBilinear2d((img_raw.shape[0], img_raw.shape[1])).cuda()(output2)\n\n        paf = nn.UpsamplingBilinear2d((img_raw.shape[0], img_raw.shape[1])).cuda()(output1)\n\n        heatmap_avg[i] = heatmap[0].data\n        paf_avg[i] = paf[0].data\n\n    heatmap_avg = torch.transpose(torch.transpose(torch.squeeze(torch.mean(heatmap_avg, 0)), 0, 1), 1, 2).cuda()\n    heatmap_avg = heatmap_avg.cpu().numpy()\n\n    paf_avg = torch.transpose(torch.transpose(torch.squeeze(torch.mean(paf_avg, 0)), 0, 1), 1, 2).cuda()\n    paf_avg = paf_avg.cpu().numpy()\n\n    return paf_avg, heatmap_avg\n\n\ndef extract_heatmap_info(heatmap_avg, param_thre1=0.1):\n    all_peaks = []\n    peak_counter = 0\n\n    for part in range(18):\n        map_ori = heatmap_avg[:, :, part]\n        map_gau = gaussian_filter(map_ori, sigma=3)\n\n        map_left = np.zeros(map_gau.shape)\n        map_left[1:, :] = map_gau[:-1, :]\n        map_right = np.zeros(map_gau.shape)\n        map_right[:-1, :] = map_gau[1:, :]\n        map_up = np.zeros(map_gau.shape)\n        map_up[:, 1:] = map_gau[:, :-1]\n        map_down = np.zeros(map_gau.shape)\n        map_down[:, :-1] = map_gau[:, 1:]\n\n        peaks_binary = np.logical_and.reduce(\n            (map_gau >= map_left, map_gau >= map_right, map_gau >= map_up,\n             map_gau >= map_down, map_gau > param_thre1))\n\n        peaks = zip(np.nonzero(peaks_binary)[1], np.nonzero(peaks_binary)[0])  # note reverse\n        peaks = list(peaks)\n        peaks_with_score = [x + (map_ori[x[1], x[0]],) for x in peaks]\n        ids = range(peak_counter, peak_counter + len(peaks))\n        peaks_with_score_and_id = [peaks_with_score[i] + (ids[i],) for i in range(len(ids))]\n\n        all_peaks.append(peaks_with_score_and_id)\n        peak_counter += len(peaks)\n\n    return all_peaks\n\n\ndef extract_paf_info(img_raw, paf_avg, all_peaks, param_thre2=0.05, param_thre3=0.5):\n    connection_all = []\n    special_k = []\n    mid_num = 10\n\n    for k in range(len(map_ids)):\n        score_mid = paf_avg[:, :, [x - 19 for x in map_ids[k]]]\n        candA = all_peaks[limb_seq[k][0] - 1]\n        candB = all_peaks[limb_seq[k][1] - 1]\n        nA = len(candA)\n        nB = len(candB)\n        if nA != 0 and nB != 0:\n            connection_candidate = []\n            for i in range(nA):\n                for j in range(nB):\n                    vec = np.subtract(candB[j][:2], candA[i][:2])\n                    norm = math.sqrt(vec[0] * vec[0] + vec[1] * vec[1])\n                    vec = np.divide(vec, norm)\n\n                    startend = zip(np.linspace(candA[i][0], candB[j][0], num=mid_num),\n                                   np.linspace(candA[i][1], candB[j][1], num=mid_num))\n                    startend = list(startend)\n\n                    vec_x = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 0]\n                                      for I in range(len(startend))])\n                    vec_y = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 1]\n                                      for I in range(len(startend))])\n\n                    score_midpts = np.multiply(vec_x, vec[0]) + np.multiply(vec_y, vec[1])\n                    score_with_dist_prior = sum(score_midpts) / len(score_midpts)\n                    score_with_dist_prior += min(0.5 * img_raw.shape[0] / norm - 1, 0)\n\n                    criterion1 = len(np.nonzero(score_midpts > param_thre2)[0]) > 0.8 * len(score_midpts)\n                    criterion2 = score_with_dist_prior > 0\n                    if criterion1 and criterion2:\n                        connection_candidate.append(\n                            [i, j, score_with_dist_prior, score_with_dist_prior + candA[i][2] + candB[j][2]])\n\n            connection_candidate = sorted(connection_candidate, key=lambda x: x[2], reverse=True)\n            connection = np.zeros((0, 5))\n            for c in range(len(connection_candidate)):\n                i, j, s = connection_candidate[c][0:3]\n                if i not in connection[:, 3] and j not in connection[:, 4]:\n                    connection = np.vstack([connection, [candA[i][3], candB[j][3], s, i, j]])\n                    if len(connection) >= min(nA, nB):\n                        break\n\n            connection_all.append(connection)\n        else:\n            special_k.append(k)\n            connection_all.append([])\n\n    return special_k, connection_all\n\n\ndef get_subsets(connection_all, special_k, all_peaks):\n    # last number in each row is the total parts number of that person\n    # the second last number in each row is the score of the overall configuration\n    subset = -1 * np.ones((0, 20))\n    candidate = np.array([item for sublist in all_peaks for item in sublist])\n\n    for k in range(len(map_ids)):\n        if k not in special_k:\n            partAs = connection_all[k][:, 0]\n            partBs = connection_all[k][:, 1]\n            indexA, indexB = np.array(limb_seq[k]) - 1\n\n            for i in range(len(connection_all[k])):  # = 1:size(temp,1)\n                found = 0\n                subset_idx = [-1, -1]\n                for j in range(len(subset)):  # 1:size(subset,1):\n                    if subset[j][indexA] == partAs[i] or subset[j][indexB] == partBs[i]:\n                        subset_idx[found] = j\n                        found += 1\n\n                if found == 1:\n                    j = subset_idx[0]\n                    if (subset[j][indexB] != partBs[i]):\n                        subset[j][indexB] = partBs[i]\n                        subset[j][-1] += 1\n                        subset[j][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]\n                elif found == 2:  # if found 2 and disjoint, merge them\n                    j1, j2 = subset_idx\n                    print(\"found = 2\")\n                    membership = ((subset[j1] >= 0).astype(int) + (subset[j2] >= 0).astype(int))[:-2]\n                    if len(np.nonzero(membership == 2)[0]) == 0:  # merge\n                        subset[j1][:-2] += (subset[j2][:-2] + 1)\n                        subset[j1][-2:] += subset[j2][-2:]\n                        subset[j1][-2] += connection_all[k][i][2]\n                        subset = np.delete(subset, j2, 0)\n                    else:  # as like found == 1\n                        subset[j1][indexB] = partBs[i]\n                        subset[j1][-1] += 1\n                        subset[j1][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]\n\n                # if find no partA in the subset, create a new subset\n                elif not found and k < 17:\n                    row = -1 * np.ones(20)\n                    row[indexA] = partAs[i]\n                    row[indexB] = partBs[i]\n                    row[-1] = 2\n                    row[-2] = sum(candidate[connection_all[k][i, :2].astype(int), 2]) + connection_all[k][i][2]\n                    subset = np.vstack([subset, row])\n    return subset, candidate\n\n\ndef draw_key_point(subset, all_peaks, img_raw):\n    del_ids = []\n    for i in range(len(subset)):\n        if subset[i][-1] < 4 or subset[i][-2] / subset[i][-1] < 0.4:\n            del_ids.append(i)\n    subset = np.delete(subset, del_ids, axis=0)\n\n    img_canvas = img_raw.copy()  # B,G,R order\n\n    for i in range(18):\n        for j in range(len(all_peaks[i])):\n            cv2.circle(img_canvas, all_peaks[i][j][0:2], 4, colors[i], thickness=-1)\n\n    return subset, img_canvas\n\n\ndef link_key_point(img_canvas, candidate, subset, stickwidth=4):\n    for i in range(17):\n        for n in range(len(subset)):\n            index = subset[n][np.array(limb_seq[i]) - 1]\n            if -1 in index:\n                continue\n            cur_canvas = img_canvas.copy()\n            Y = candidate[index.astype(int), 0]\n            X = candidate[index.astype(int), 1]\n            mX = np.mean(X)\n            mY = np.mean(Y)\n            length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5\n            angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))\n            polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth), int(angle), 0, 360, 1)\n            cv2.fillConvexPoly(cur_canvas, polygon, colors[i])\n            img_canvas = cv2.addWeighted(img_canvas, 0.4, cur_canvas, 0.6, 0)\n\n    return img_canvas\n\ndef pad_right_down_corner(img, stride, pad_value):\n    h = img.shape[0]\n    w = img.shape[1]\n\n    pad = 4 * [None]\n    pad[0] = 0  # up\n    pad[1] = 0  # left\n    pad[2] = 0 if (h % stride == 0) else stride - (h % stride)  # down\n    pad[3] = 0 if (w % stride == 0) else stride - (w % stride)  # right\n\n    img_padded = img\n    pad_up = np.tile(img_padded[0:1, :, :] * 0 + pad_value, (pad[0], 1, 1))\n    img_padded = np.concatenate((pad_up, img_padded), axis=0)\n    pad_left = np.tile(img_padded[:, 0:1, :] * 0 + pad_value, (1, pad[1], 1))\n    img_padded = np.concatenate((pad_left, img_padded), axis=1)\n    pad_down = np.tile(img_padded[-2:-1, :, :] * 0 + pad_value, (pad[2], 1, 1))\n    img_padded = np.concatenate((img_padded, pad_down), axis=0)\n    pad_right = np.tile(img_padded[:, -2:-1, :] * 0 + pad_value, (1, pad[3], 1))\n    img_padded = np.concatenate((img_padded, pad_right), axis=1)\n\n    return img_padded, pad\n\n\nif __name__ == '__main__':\n    print(get_pose_model())","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"papermill":{"duration":2.754634,"end_time":"2020-09-29T03:23:51.364957","exception":false,"start_time":"2020-09-29T03:23:48.610323","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"First let's download the pre-trained model.","metadata":{"papermill":{"duration":0.014861,"end_time":"2020-09-29T03:23:51.396920","exception":false,"start_time":"2020-09-29T03:23:51.382059","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# Using gdown to download the model directly from Google Drive\n\n! conda install -y gdown\nimport gdown","metadata":{"_kg_hide-output":true,"papermill":{"duration":70.21864,"end_time":"2020-09-29T03:25:01.631100","exception":false,"start_time":"2020-09-29T03:23:51.412460","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"url = 'https://drive.google.com/u/0/uc?export=download&confirm=f_Ix&id=0B1asvDK18cu_MmY1ZkpaOUhhRHM'\nmodel = 'coco_pose_iter_440000.pth.tar'\ngdown.download(url, model, quiet=False)","metadata":{"papermill":{"duration":2.716847,"end_time":"2020-09-29T03:25:04.456135","exception":false,"start_time":"2020-09-29T03:25:01.739288","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"state_dict = torch.load('./coco_pose_iter_440000.pth.tar')['state_dict']   # getting the pre-trained model's parameters\n# A state_dict is simply a Python dictionary object that maps each layer to its parameter tensor.\n\nmodel_pose = get_pose_model()   # building the model (see fn. defn. above). To see the architecture, see below cell.\nmodel_pose.load_state_dict(state_dict)   # Loading the parameters (weights, biases) into the model.\n\nmodel_pose.float()   # I'm not sure why this is used. No difference if you remove it.","metadata":{"_kg_hide-output":true,"papermill":{"duration":0.788917,"end_time":"2020-09-29T03:25:05.366889","exception":false,"start_time":"2020-09-29T03:25:04.577972","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"arch_image = '../input/indonesian-traditional-dance/tgagrakanyar/tga_0000.jpg'\nimg_ori = cv2.imread(arch_image)\nplt.figure(figsize=(15, 8))\nplt.imshow(img_ori[...,::-1])","metadata":{"papermill":{"duration":0.882163,"end_time":"2020-09-29T03:25:06.368274","exception":false,"start_time":"2020-09-29T03:25:05.486111","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Notice, the first 10 layers are from VGG-19. But here instead of downloading the model and loading the layers from there, we simply hardcoaded it in get_pose_model()","metadata":{"papermill":{"duration":0.128709,"end_time":"2020-09-29T03:25:06.629074","exception":false,"start_time":"2020-09-29T03:25:06.500365","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# Run this to view the model's architecture\n#model_pose.eval()","metadata":{"papermill":{"duration":0.162978,"end_time":"2020-09-29T03:25:06.919564","exception":false,"start_time":"2020-09-29T03:25:06.756586","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"use_gpu = True\n\nif use_gpu:\n    model_pose.cuda()\n    model_pose = torch.nn.DataParallel(model_pose, device_ids=range(torch.cuda.device_count()))\n    cudnn.benchmark = True","metadata":{"papermill":{"duration":6.01287,"end_time":"2020-09-29T03:25:13.066731","exception":false,"start_time":"2020-09-29T03:25:07.053861","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def estimate_pose(img_ori):\n    \n    # People might be at different scales in the image, perform inference at multiple scales to boost results\n    scale_param = [0.5, 1.0, 1.5, 2.0]\n    \n    # Predict Heatmaps for approximate joint position\n    # Use Part Affinity Fields (PAF's) as guidance to link joints to form skeleton\n    # PAF's are just unit vectors along the limb encoding the direction of the limb\n    # A dot product of possible joint connection will be high if actual limb else low\n    \n    paf_info, heatmap_info = get_paf_and_heatmap(model_pose, img_ori, scale_param)\n    peaks = extract_heatmap_info(heatmap_info)\n    sp_k, con_all = extract_paf_info(img_ori, paf_info, peaks)\n    \n    subsets, candidates = get_subsets(con_all, sp_k, peaks)\n    subsets, img_points = draw_key_point(subsets, peaks, img_ori)\n    \n    # After predicting Heatmaps and PAF's, proceeed to link joints correctly\n    img_canvas = link_key_point(img_points, candidates, subsets)\n    \n    \n    plt.figure(figsize=(15, 10))\n    \n    plt.subplot(1, 2, 1)\n    plt.imshow(img_points[...,::-1])\n    \n    plt.subplot(1, 2, 2)\n    plt.imshow(img_canvas[...,::-1])","metadata":{"papermill":{"duration":0.141263,"end_time":"2020-09-29T03:25:13.336611","exception":false,"start_time":"2020-09-29T03:25:13.195348","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_image = '../input/indonesian-traditional-dance/tgagrakanyar/tga_0000.jpg'\nimg_ori = cv2.imread(test_image)\nestimate_pose(img_ori)","metadata":{"papermill":{"duration":3.306725,"end_time":"2020-09-29T03:25:16.825272","exception":false,"start_time":"2020-09-29T03:25:13.518547","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_image = '../input/indonesian-traditional-dance/tgagrakanyar/tga_0010.jpg'\nimg_ori = cv2.imread(test_image)\nestimate_pose(img_ori)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_image = '../input/indonesian-traditional-dance/tgagrakanyar/tga_0020.jpg'\nimg_ori = cv2.imread(test_image)\nestimate_pose(img_ori)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_image = '../input/indonesian-traditional-dance/tgagrakanyar/tga_0030.jpg'\nimg_ori = cv2.imread(test_image)\nestimate_pose(img_ori)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_image = '../input/indonesian-traditional-dance/tgagrakanyar/tga_0040.jpg'\nimg_ori = cv2.imread(test_image)\nestimate_pose(img_ori)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_image = '../input/indonesian-traditional-dance/tgagrakanyar/tga_0050.jpg'\nimg_ori = cv2.imread(test_image)\nestimate_pose(img_ori)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}