├── .gitignore ├── colormaps ├── tab.txt ├── tab20c.txt └── anliang_render.txt ├── pics ├── BamaPig2D.jpg ├── keypoint.jpg └── labelorder_for_3dlabel.jpg ├── LICENSE ├── assemble_BamaPig3D.py ├── utils.py ├── undistortion.py ├── visualize_BamaPig2D.py ├── bodymodel_np.py ├── README.md └── visualize_BamaPig3D.py /.gitignore: -------------------------------------------------------------------------------- 1 | *__pycache__* 2 | output/ 3 | figs/ -------------------------------------------------------------------------------- /colormaps/tab.txt: -------------------------------------------------------------------------------- 1 | 0 0 127 2 | 0 28 255 3 | 0 188 255 4 | 95 255 150 -------------------------------------------------------------------------------- /colormaps/tab20c.txt: -------------------------------------------------------------------------------- 1 | 49 130 189 2 | 107 174 214 3 | 158 202 225 4 | 198 219 239 -------------------------------------------------------------------------------- /pics/BamaPig2D.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/anl13/MAMMAL_datasets/HEAD/pics/BamaPig2D.jpg -------------------------------------------------------------------------------- /pics/keypoint.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/anl13/MAMMAL_datasets/HEAD/pics/keypoint.jpg -------------------------------------------------------------------------------- /pics/labelorder_for_3dlabel.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/anl13/MAMMAL_datasets/HEAD/pics/labelorder_for_3dlabel.jpg -------------------------------------------------------------------------------- /colormaps/anliang_render.txt: -------------------------------------------------------------------------------- 1 | 246 205 97 2 | 173 203 227 3 | 254 138 113 4 | 26 153 0 5 | 250 217 193 6 | 170 0 204 7 | 176 191 26 8 | 175 0 42 9 | 239 22 205 10 | 124 185 232 11 | 241 156 187 12 | 59 122 87 13 | 30 105 210 14 | 80 127 255 15 | 237 149 100 16 | 220 248 255 -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 Liang An 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. -------------------------------------------------------------------------------- /assemble_BamaPig3D.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | # This function shows how to use "label_mesh" and "label_3d" annotations to get the 4 | # final label_mix gt annotation. 5 | # This function will help you to understand Supplementary Fig. 8 in the paper. 6 | # 7 | BamaPig3D_path = "H:/examples/BamaPig3D/" 8 | def assemble(): 9 | folder1 = BamaPig3D_path + "label_mesh" 10 | order1 = [0,1,2,3] 11 | folder2 = BamaPig3D_path + "label_3d" 12 | order2 = [0,1,2,3] 13 | folder3 = BamaPig3D_path + "label_mix" 14 | 15 | for i in range(65,66): 16 | frameid = 25 * i 17 | all_data1 = np.zeros((4, 23, 3)) 18 | all_data2 = np.zeros((4, 23, 3)) 19 | all_data3 = np.zeros((4, 23, 3)) 20 | for pid in range(4): 21 | filename1 = folder1 + "/pig_{}_frame_{:06d}.txt".format(pid, frameid) 22 | data = np.loadtxt(filename1) 23 | all_data1[order1[pid]] = data 24 | filename2 = folder2 + "/pig_{}_frame_{:06d}.txt".format(pid, frameid) 25 | data = np.loadtxt(filename2) 26 | all_data2[pid] = data 27 | all_data3 = all_data2.copy() 28 | for pid in range(4): 29 | for jid in range(23): 30 | if np.linalg.norm(all_data3[pid, jid]) == 0 and np.linalg.norm(all_data1[pid,jid]) > 0: 31 | all_data3[pid, jid] = all_data1[pid, jid] 32 | 33 | for pid in range(4): 34 | filename = folder3 +"/pig_{}_frame_{:06d}.txt".format(pid, frameid) 35 | np.savetxt(filename, all_data3[pid]) 36 | 37 | if __name__ == "__main__": 38 | assemble() -------------------------------------------------------------------------------- /utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import os 3 | 4 | # 23 keypoints names 5 | # only 19 have not `none` name 6 | # This is an issue rooted in history 7 | # in fact, one can slightly modify the dataset to remove these 4 non-used keypoints. 8 | g_jointnames = [ 9 | "nose", 10 | "l_eye", 11 | "r_eye", 12 | "l_ear", 13 | "r_ear", 14 | "l_shoulder", 15 | "r_shoulder", 16 | "l_elbow", 17 | "r_elbow", 18 | "l_paw", 19 | "r_paw", 20 | "l_hip", 21 | "r_hip", 22 | "l_knee", 23 | "r_knee", 24 | "l_foot", 25 | "r_foot", 26 | "none", # not used. 27 | "tail", 28 | "none", # not used 29 | "center", 30 | "none", # not used 31 | "none" # not used 32 | ] 33 | 34 | g_bones = [ # bone structure for 23 keypoints 35 | [0,1], 36 | [0,2], 37 | [1,2], 38 | [1,3], 39 | [2,4], 40 | [0,20], 41 | [20,18], 42 | [20,5], 43 | [5,7], 44 | [7,9], 45 | [20,6], 46 | [6,8], 47 | [8,10], 48 | [18,11], 49 | [11,13], 50 | [13,15], 51 | [18,12], 52 | [12,14], 53 | [14,16] 54 | ] 55 | 56 | g_bones_19 = [ # bone structure for the final 19 valid keypoints 57 | [0,1], 58 | [0,2], 59 | [1,2], 60 | [1,3], 61 | [2,4], 62 | [0,18], 63 | [18,17], 64 | [18,5], 65 | [5,7], 66 | [7,9], 67 | [18,6], 68 | [6,8], 69 | [8,10], 70 | [17,11], 71 | [11,13], 72 | [13,15], 73 | [17,12], 74 | [12,14], 75 | [14,16] 76 | ] 77 | 78 | # group name of each keypoint 79 | # e.g. nose, eyes, ears are all Head part. 80 | g_groupnames = [ 81 | "Head", 82 | "Head", 83 | "Head", 84 | "Head", 85 | "Head", 86 | "L_arm", 87 | "R_arm", 88 | "L_arm", 89 | "R_arm", 90 | "L_arm", 91 | "R_arm", 92 | "L_leg", 93 | "R_leg", 94 | "L_leg", 95 | "R_leg", 96 | "L_leg", 97 | "R_leg", 98 | "none", 99 | "Tail", 100 | "none", 101 | "Center", 102 | "none", 103 | "none" 104 | ] 105 | 106 | # indices of not `none` keypoints 107 | g_all_parts = [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,18,20] 108 | 109 | # some other part divide for flexible usage. 110 | g_head = [0,1,2,3,4] 111 | g_left_front_leg = [5,7,9] 112 | g_right_front_leg = [6,8,10] 113 | g_left_hind_leg = [11,13,15] 114 | g_right_hind_leg = [12,14,16] 115 | g_legs = g_left_front_leg + g_left_hind_leg + g_right_front_leg + g_right_hind_leg 116 | g_leg_level1 = [5,6,11,12] 117 | g_leg_level2 = [7,8,13,14] 118 | g_leg_level3 = [9,10,15,16] 119 | g_trunk = [20,18] 120 | g_pig_ids_for_eval = [0,1,2,3] 121 | 122 | # This function is used to load all 3D labeled data. 123 | def load_joint23(folder, start = 0, step = 25, num = 70, order=[0,1,2,3]): 124 | all_data = [] 125 | for i in range(num): 126 | frameid = start + step * i 127 | single_frame = [0,1,2,3] 128 | for pid in range(4): 129 | filename = folder + "/pig_{}_frame_{:06d}.txt".format(pid, frameid) 130 | data = np.loadtxt(filename) 131 | index = order[pid] 132 | single_frame[index] = data 133 | all_data.append(single_frame) 134 | all_data = np.asarray(all_data) 135 | 136 | return all_data -------------------------------------------------------------------------------- /undistortion.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import cv2 3 | import os 4 | import pickle 5 | from time import time 6 | ''' 7 | This function shows how to generate inverse mapping from original mapping. 8 | Given a calibration folder `calib_folder`, this function writes `inverse_map_dict.pkl` file to it. 9 | This function may run 80 seconds. 10 | 11 | input: 12 | mapx: mapx [h,w] 13 | mapy: mapy [h,w] 14 | return: 15 | result: [w*h, 5], nearest neighbor indices 16 | dists : [w*h, 5], nearest neighbor distances 17 | ''' 18 | 19 | def inverse_dist_weighting(xs, ds): 20 | assert(xs.shape == ds.shape) 21 | ws = 1 / ds 22 | interp = (xs * ws).sum() / (ws.sum()) 23 | return interp 24 | 25 | def get_inverse_remap(result, dists, w, h): 26 | inv_mapx = np.zeros((h,w), dtype=np.float32) 27 | inv_mapy = np.zeros((h,w), dtype=np.float32) 28 | for i in range(h): 29 | for j in range(w): 30 | index = i * w + j 31 | nn_inds = result[index] 32 | nn_inds_x = (nn_inds % w).astype(np.int) 33 | nn_inds_y = (nn_inds / w).astype(np.int) 34 | nn_dists = dists[index] 35 | new_x = inverse_dist_weighting(nn_inds_x, nn_dists) 36 | new_y = inverse_dist_weighting(nn_inds_y, nn_dists) 37 | inv_mapx[i,j] = new_x 38 | inv_mapy[i,j] = new_y 39 | return inv_mapx, inv_mapy 40 | 41 | def linear_interp(y,x,map): 42 | assert(y>=0 and x>=0 and y<=map.shape[0]-1 and x<=map.shape[1]-1) 43 | y1 = int(y) 44 | y2 = y1 + 1 45 | x1 = int(x) 46 | x2 = x1 + 1 47 | dx = x - x1 48 | dy = y - y1 49 | value = map[y1,x1] * (1-dy) * (1-dx) \ 50 | + map[y1,x2] * (1-dy) * dx \ 51 | + map[y2,x1] * dy * (1-dx) \ 52 | + map[y2,x2] * dy * dx 53 | return value 54 | 55 | def undist_point(point, inv_mapx, inv_mapy): 56 | [x,y] = point 57 | new_x = linear_interp(y,x,inv_mapx) 58 | new_y = linear_interp(y,x,inv_mapy) 59 | return np.array([new_x, new_y]) 60 | 61 | # points are in [x,y] format, as used in cv2 62 | def undist_points(points, inv_mapx, inv_mapy): 63 | new_points = points.copy() 64 | for i in range(points.shape[0]): 65 | new_points[i] = undist_point(points[i], inv_mapx, inv_mapy) 66 | return new_points 67 | 68 | 69 | ''' 70 | This function is used to undistort 2D points. 71 | points: [N,2], np.float32 72 | K : pin-hole projection matrix, [3,3], np.float32 73 | coeff : [5], np.float32, is non-linear distortion coefficients, in format [k1,k2,p1,p2,k3]. 74 | Please refer to https://docs.opencv.org/4.x/dc/dbb/tutorial_py_calibration.html for details. 75 | newcameramtx: [3,3], np.float32, K matrix after getOptimalNewCameraMatrix. 76 | ''' 77 | def undist_points_cv2(points, K, coeff, newcameramtx): 78 | points_cv = points.copy() 79 | points_cv = points_cv.reshape([points_cv.shape[0], 1, points_cv.shape[1]]) 80 | new_points_2 = cv2.undistortPoints(points_cv, K, coeff, P=newcameramtx) 81 | new_points_2 = new_points_2.squeeze() 82 | return new_points_2 # [N,2] 83 | 84 | # You can find the intrinsic parameters used by BamaPig2D and BamaPig3D datasets here. 85 | # img : image used. [1920,1080,3], np.uint8 array is required. 86 | # calib_folder: Folder containing intrinsic calibration files. 87 | # If an empty folder is set, it will takes a little time to write calibration pkl files to it. 88 | # The K and coeff are two key intrinsic parameters. 89 | def undist_image_demo(img, calib_folder): 90 | ''' 91 | basic intrinsic calibration information 92 | ''' 93 | K = [[1625.30923, 0, 963.88710], 94 | [0, 1625.34802, 523.45901], 95 | [0, 0, 1]] 96 | K = np.array(K, dtype=np.float32) 97 | coeff = [-0.35582, 0.14595, -0.00031, -0.00004, 0.00000] 98 | coeff = np.array(coeff, dtype=np.float32) 99 | w = 1920 100 | h = 1080 101 | newcameramtx, roi = cv2.getOptimalNewCameraMatrix(K, coeff, (w,h), 1, (w,h)) 102 | mapx, mapy = cv2.initUndistortRectifyMap(K, coeff, None, newcameramtx, (w,h), 5) 103 | 104 | if not os.path.exists("output"): 105 | os.makedirs("output") 106 | ''' 107 | test undistort an empty image 108 | ''' 109 | empty = np.zeros([1080, 1920, 3], np.uint8) 110 | undist_empty = cv2.remap(empty, mapx, mapy, cv2.INTER_LINEAR, borderValue=(255,255,255)) 111 | cv2.imwrite("output/undist.png", undist_empty) 112 | ''' 113 | test undistortion mapping 114 | ''' 115 | undist = cv2.remap(img, mapx, mapy, cv2.INTER_LINEAR, borderValue=(255,255,255)) 116 | cv2.imwrite("output/intrinsic_calib_demo.png", undist) 117 | 118 | if __name__ == "__main__": 119 | # Change BamaPig3D_folder to your own BamaPig3D path. 120 | BamaPig3D_folder = "H:/examples/BamaPig3D/" 121 | 122 | calib_folder = BamaPig3D_folder + "intrinsic_camera_params/" 123 | demo_img = cv2.imread(BamaPig3D_folder + "images/cam0/000000.jpg") 124 | undist_image_demo(demo_img, calib_folder) -------------------------------------------------------------------------------- /visualize_BamaPig2D.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | import numpy as np 4 | import json 5 | import cv2 6 | import os 7 | from utils import g_bones 8 | import pickle 9 | import matplotlib as mpl 10 | import matplotlib.pyplot as plt 11 | 12 | COLORS = np.loadtxt("colormaps/anliang_render.txt") 13 | COLORS = COLORS[:,(2,1,0)].astype(np.int32).tolist() 14 | 15 | 16 | def draw_mask(img, mask, color): 17 | # cv2.fillPoly(img, np.int32([mask + 0.5]), color) 18 | cv2.polylines(img, np.int32([mask + 0.5]), isClosed=True, color=color, thickness=5) 19 | 20 | ## points : [N, 3] 21 | ## if N == 23: set bone=g_bones 22 | ## if N == 19: set bone=g_bones_19 23 | def draw_keypoints(img, points, id, bone=None): 24 | base_color = np.asarray(COLORS[id]) 25 | for i in range(points.shape[0]): 26 | if points.shape[1] == 3 and points[i,2] == 2: # visible 27 | color = base_color 28 | color_tuple = (int(color[0]), int(color[1]), int(color[2])) 29 | cv2.circle(img, tuple(points[i,0:2].astype(np.int32)), 9, color_tuple, thickness=-1) 30 | elif points.shape[1] == 2: 31 | color = base_color 32 | color_tuple = (int(color[0]), int(color[1]), int(color[2])) 33 | cv2.circle(img, tuple(points[i,0:2].astype(np.int32)), 9, color_tuple, thickness=-1) 34 | if bone is None: 35 | return 36 | 37 | for pair in bone: 38 | i, dad_id = pair 39 | sun = points[i,0:2].astype(np.int32) 40 | if points.shape[1] == 3: 41 | v_sun = points[i,2] 42 | else: 43 | v_sun = 1 44 | if dad_id < 0: 45 | continue 46 | dad = points[dad_id,0:2].astype(np.int32) 47 | if points.shape[1] == 3: 48 | v_dad = points[dad_id,2].astype(np.int32) 49 | else: 50 | v_dad = 1 51 | if v_sun == 0 or v_dad == 0: 52 | continue 53 | color = base_color 54 | color_tuple = (int(color[0]), int(color[1]), int(color[2])) 55 | cv2.line(img, tuple(sun), tuple(dad), color_tuple, thickness=4) 56 | 57 | def draw_box(img, box, colorid): 58 | x1, y1, x2, y2 = box[0], box[1], box[0] + box[2], box[1] + box[3] 59 | cv2.rectangle(img, (int(x1), int(y1)), (int(x2),int(y2)), tuple(COLORS[colorid]), thickness=3) 60 | 61 | 62 | # This function is used to reproduce BamaPig2D visualization results in the 63 | # Supplementary Fig. 4b of the paper. 64 | def visualize_BamaPig2D(datafolder, output_folder,image_ids_to_rend=[0]): 65 | with open(datafolder + "/annotations/eval_pig_cocostyle.json", 'r') as f: 66 | annos_eval = json.load(f) 67 | with open(datafolder + "/annotations/train_pig_cocostyle.json", 'r') as f: 68 | annos_train = json.load(f) 69 | annotations = annos_eval["annotations"] + annos_train["annotations"] 70 | if not os.path.exists(output_folder): 71 | os.makedirs(output_folder) 72 | 73 | for imgid in image_ids_to_rend: 74 | img = cv2.imread(datafolder + "/images/{:06d}.png".format(imgid)) 75 | colorid = 0 76 | for A in annotations: 77 | if A['image_id'] != imgid: 78 | continue 79 | box = np.asarray(A['bbox']) 80 | keypoints = np.asarray(A['keypoints']) 81 | masks = A['segmentation'] 82 | draw_box(img, box, colorid) 83 | for m in masks: 84 | m = np.asarray(m) 85 | m = m.reshape([-1,2]) 86 | draw_mask(img, m, COLORS[colorid]) 87 | draw_keypoints(img, keypoints.reshape([-1,3]), colorid, g_bones) 88 | colorid += 1 89 | cv2.imwrite(output_folder + "/{:06d}.png".format(imgid), img) 90 | print("write image ", imgid) 91 | 92 | 93 | def draw_vis_2D(BamaPig2D_folder): 94 | with open(BamaPig2D_folder + "/annotations/eval_pig_cocostyle.json", 'r') as f: 95 | annos_eval = json.load(f) 96 | with open(BamaPig2D_folder + "/annotations/train_pig_cocostyle.json", 'r') as f: 97 | annos_train = json.load(f) 98 | Total = 11504 * 19 99 | part_levels = [ 100 | [0,17,18], # trunk 101 | [1,2,3,4], # head 102 | [5,6,7,8,9,10,11,12,13,14,15,16] # limbs 103 | ] 104 | 105 | Vis = np.zeros(3) 106 | for A in (annos_train['annotations'] + annos_eval['annotations']): 107 | keypoints = np.asarray(A['keypoints']) 108 | keypoints = keypoints.reshape([-1,3]) 109 | for k in range(3): 110 | visible = (keypoints[part_levels[k],2] > 1).sum() 111 | Vis[k] += visible 112 | Vis[0] /= (11504 * 3) 113 | Vis[1] /= (11504 * 4) 114 | Vis[2] /= (11504 * 12) 115 | 116 | mpl.rc('font', family='Arial') 117 | fig = plt.figure(figsize=(1.2, 1.4)) 118 | colormaps = np.loadtxt("colormaps/tab20c.txt") / 255 119 | part_level_names = ["Trunk", "Head", "Limbs"] 120 | xs = np.asarray([0,1,2]) 121 | plt.bar(xs, Vis, color=colormaps[0], width=0.6, edgecolor=(0,0,0), lw=0.5) 122 | plt.ylim(0,0.5) 123 | plt.xticks(xs, part_level_names, fontsize=7) 124 | plt.yticks([0,0.1,0.2,0.3,0.4,0.5], [0,10,20,30,40,50], fontsize=7) 125 | plt.ylabel("Percentage of Visible\n Keypoints (%)", fontsize=7) 126 | ax = fig.get_axes()[0] 127 | ax.spines['right'].set_visible(False) 128 | ax.spines['top'].set_visible(False) 129 | for line in ["bottom", "left", "right"]: 130 | ax.spines[line].set_linewidth(0.5) 131 | ax.xaxis.set_tick_params(width=0.5) 132 | ax.yaxis.set_tick_params(width=0.5) 133 | 134 | plt.savefig("output/supp_fig_4d.png", dpi=1000, bbox_inches='tight', pad_inches=0) 135 | # plt.savefig("output/supp_fig_4d.svg", dpi=1000, bbox_inches='tight', pad_inches=0) # uncomment this line to generate vector image file. 136 | 137 | 138 | if __name__ == "__main__": 139 | # You may change the BamaPig2D_folder to your own path of BamaPig2D 140 | BamaPig2D_folder = "H:/examples/BamaPig2D/" 141 | 142 | # [0,100,1000,2000,3000,3009,3100,3300] These images are used for generating Supplementary Fig. 4b. 143 | visualize_BamaPig2D(BamaPig2D_folder, "output", [0,100,1000,2000,3000,3009,3100,3300]) 144 | 145 | draw_vis_2D(BamaPig2D_folder) 146 | 147 | -------------------------------------------------------------------------------- /bodymodel_np.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import os 3 | from scipy.spatial.transform import Rotation 4 | 5 | ''' 6 | To use this model, please first download the model files from (TODO link) and 7 | put them as the folder `./PIG_Model/` (you can rename it as you like and change the 8 | `model_folder` variable below. ) 9 | ''' 10 | class BodyModelNumpy(object): 11 | def __init__(self, model_folder = "./PIG_Model/"): 12 | self.model_folder = model_folder 13 | self.vertices, self.parents, self.joints, self.weights, self.faces = \ 14 | self.readmodel(self.model_folder) 15 | self.joint_num = self.joints.shape[0] 16 | 17 | # [target joint, type, source index, weight] 18 | self.optimize_pair = [ 19 | [ 0, 1, 10895, 1 ], # nose 20 | [ 1, 1, 938, 1 ], # left eye 21 | [ 2, 1, 6053, 1 ], # right eye 22 | [ 3, 1, 1368, 1 ], # left ear 23 | [ 4, 1, 6600, 1 ], # right ear 24 | [ 5, 0, 15, 1 ], # left shouder 25 | [ 6, 0, 7, 1 ], # right shoulder 26 | [ 7, 0, 16, 1 ], # left elbow 27 | [ 8, 0, 8, 1 ], # right elbow 28 | [ 9, 0, 17, 1 ], # left paw 29 | [ 10, 0, 9, 1 ], # right paw 30 | [ 11, 0, 56, 1 ], # left hip 31 | [ 12, 0, 40, 1 ], # right hip 32 | [ 13, 0, 57, 1 ], # left knee 33 | [ 14, 0, 41, 1 ], # right knee 34 | [ 15, 0, 58, 1 ], # left foot 35 | [ 16, 0, 42, 1 ], # right foot 36 | [ 17, -1, 0, 0], # neck(not use) 37 | [ 18, 1, 7903, 1 ], # tail 38 | [ 19, -1, 0,0], #wither (not use) 39 | [ 20, 0, 2, 1 ], # center 40 | [ 21, -1, 0,0], # tail middle (not use) 41 | [ 22, -1, 0,0] # tail end (not use) 42 | ] 43 | 44 | self.translation = np.zeros(3, dtype=np.float32) 45 | self.poseparam = np.zeros([self.joint_num, 3], dtype=np.float32) 46 | self.scale = 1 47 | self.posed_joints = self.joints.copy() 48 | self.posed_vertices = self.vertices.copy() 49 | 50 | # function to parse the pose param txt file. 51 | def readstate(self,filename): 52 | states = np.loadtxt(filename) 53 | translation = states[0:3] # the first 3 dim is global translation 54 | scale = states[-1] # the last param is global scale 55 | poseparam = states[3:-1].reshape([-1,3]) # others are pose params, in which the first 3-dim rotation is global rotation. 56 | 57 | self.translation = translation 58 | self.poseparam = poseparam 59 | self.scale = scale 60 | return translation, poseparam, scale 61 | 62 | 63 | def readmodel(self, model_folder): 64 | vertices_np = np.loadtxt(os.path.join(model_folder, "vertices.txt")) 65 | parents_np = np.loadtxt(os.path.join(model_folder, "parents.txt")).squeeze() 66 | joints_np = np.loadtxt(os.path.join(model_folder, "t_pose_joints.txt")) 67 | weights = np.zeros((vertices_np.shape[0], parents_np.shape[0])) 68 | _weights = np.loadtxt(os.path.join(model_folder, "skinning_weights.txt")) 69 | for i in range(_weights.shape[0]): 70 | jointid = int(_weights[i,0]) 71 | vertexid = int(_weights[i,1]) 72 | value = _weights[i,2] 73 | weights[vertexid, jointid] = value 74 | faces_vert = np.loadtxt(os.path.join(model_folder, "faces_vert.txt")) 75 | 76 | return vertices_np, parents_np.astype(np.int), joints_np, weights, faces_vert.astype(np.int) 77 | 78 | 79 | def poseparam2Rot(self, poseparam): 80 | Rot = np.zeros((poseparam.shape[0], 3, 3), dtype=np.float32) 81 | r_tmp1 = Rotation.from_euler('ZYX', poseparam[0], degrees=False) 82 | Rot[0] = r_tmp1.as_matrix() 83 | r_tmp2 = Rotation.from_rotvec(poseparam[1:]) 84 | Rot[1:] = r_tmp2.as_matrix() 85 | 86 | return Rot 87 | 88 | 89 | def write_obj(self, filename): 90 | with open(filename, 'w') as fp: 91 | for v in self.posed_vertices: 92 | fp.write('v %f %f %f\n' % (v[0], v[1], v[2])) 93 | 94 | for f in self.faces + 1: 95 | fp.write('f %d %d %d\n' % (f[0], f[1], f[2])) 96 | 97 | 98 | def joint_Rot(self, Rot): 99 | skinmat = np.repeat(np.eye(4, dtype=np.float32).reshape(1, 4, 4), repeats=self.joints.shape[0], axis=0) 100 | skinmat[0, :3, :3] = Rot[0] 101 | skinmat[0, :3, 3] = self.joints[0] 102 | for jind in range(1, self.joints.shape[0]): 103 | skinmat[jind, :3, :3] = Rot[jind] 104 | skinmat[jind, :3, 3] = self.joints[jind] - self.joints[self.parents[jind]] 105 | skinmat[jind] = np.matmul(skinmat[self.parents[jind]], skinmat[jind]) 106 | 107 | joints_final = skinmat[:, :3, 3].copy() 108 | joints_deformed = np.zeros((self.joints.shape[0], 4), dtype=np.float32) 109 | for jind in range(self.joints.shape[0]): 110 | joints_deformed[jind, :3] = np.matmul(skinmat[jind, :3, :3], self.joints[jind]) 111 | skinmat[:, :, 3] = skinmat[:, :, 3] - joints_deformed 112 | return skinmat[:, :3, :], joints_final 113 | 114 | def regress_verts(self, skinmat): 115 | vertsmat = np.tensordot(self.weights, skinmat, axes=([1], [0])) 116 | verts_final = np.zeros((self.vertices.shape[0], 3), dtype=np.float32) 117 | for vind in range(self.vertices.shape[0]): 118 | verts_final[vind] = np.matmul(vertsmat[vind, :, :3], self.vertices[vind]) + vertsmat[vind, :, 3] 119 | 120 | return verts_final 121 | 122 | def forward(self, pose, trans=np.zeros(3, dtype=np.float32), scale=1): 123 | rot = self.poseparam2Rot(pose) 124 | skinmat, joints_final = self.joint_Rot(rot) 125 | self.posed_joints = joints_final * scale + trans 126 | verts = self.regress_verts(skinmat) 127 | self.posed_vertices = verts * scale + trans 128 | return self.posed_vertices, self.posed_joints 129 | 130 | # This function regress all the 23 keypoints with non-used ones set as zero. 131 | # return: np.ndarray (float32), [23,3] 132 | def regress_keypoints(self): 133 | keynum = len(self.optimize_pair) 134 | keypoints = np.zeros((keynum, 3), dtype=np.float32) 135 | for i in range(keynum): 136 | if self.optimize_pair[i][1] == 0: 137 | keypoints[i] = self.posed_joints[self.optimize_pair[i][2]] 138 | elif self.optimize_pair[i][1] == 1: 139 | keypoints[i] = self.posed_vertices[self.optimize_pair[i][2]] 140 | return keypoints 141 | 142 | # This function regress 19 used keypoints from the model. 143 | # return: np.ndarray (float32), [19,3] 144 | def regress_keypoints_pack(self): 145 | keynum = 19 146 | keypoints = np.zeros((keynum, 3), dtype=np.float32) 147 | non_zero = 0 148 | for i in range(len(self.optimize_pair)): 149 | if self.optimize_pair[i][1] == 0: 150 | keypoints[non_zero] = self.posed_joints[self.optimize_pair[i][2]] 151 | non_zero += 1 152 | elif self.optimize_pair[i][1] == 1: 153 | keypoints[non_zero] = self.posed_vertices[self.optimize_pair[i][2]] 154 | non_zero += 1 155 | return keypoints 156 | 157 | if __name__ == "__main__": 158 | bm = BodyModelNumpy() 159 | BamaPig3D_path = "H:/examples/BamaPig3D/" 160 | for k in range(25): 161 | frameid = k * 25 162 | for pid in range(4): 163 | filename = BamaPig3D_path + "label_pose_params/pig_{}_frame_{:06d}.txt".format(pid, frameid) 164 | trans, poseparam, scale = bm.readstate(filename) 165 | # V: [11239, 3]; 166 | # J: [62, 3] 167 | V, J = bm.forward(poseparam, trans=trans, scale = scale) 168 | # keypoints: [23, 3] 169 | keypoints = bm.regress_keypoints() 170 | 171 | # you can do something here 172 | 173 | 174 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Datasets proposed by MAMMAL 2 | This repository presents how to download and use BamaPig2D and BamaPig3D datasets proposed by [MAMMAL](https://github.com/anl13/MAMMAL_core) system. 3 | 4 | ![img](pics/BamaPig2D.jpg) 5 | 6 | ## Download 7 | BamaPig2D (8.02GB for zipflie. 9.23G after unzip, yet occupy 10.7G space on windows) can be downloaded from [Google Drive](https://drive.google.com/file/d/1yWBtNpYpkUdGKDqUAE7ya5m_fwinn0HN/view?usp=sharing) or [Baidu Drive](https://pan.baidu.com/s/1vTwipVuXHNhBFc91tNXteQ) (extract code: vj9n). 8 | 9 | BamaPig2D_sleap (33.7M, only contains `.json` files) can be downloaded from [Google Drive](https://drive.google.com/file/d/1XRFvUM8iBtzkIzr83rMkpQ4NonTKY9dk/view?usp=sharing) or [Baidu Yun](https://pan.baidu.com/s/1PC_n9_nqRsduw5JYuyTOVA) (extract code: qtb9). 10 | 11 | BamaPig3D (8.86GB for zipfile. 9.62G after unzip yet occupy 24.9G space on windows because it contains many small files) can be downloaded from [Google Drive](https://drive.google.com/file/d/1rZtuR9B9ojxQKkps0j5duwAJ-v3iM1k7/view?usp=sharing) or [Baidu Drive](https://pan.baidu.com/s/1tOgf5icIt0GKI4zpV_TE4Q) (extract code: z8u6). 12 | 13 | BamaPig3D_pure_pickle (481M for zipfile, 579M after unzip, yet occupy 941M space on Windows. ). [Google Drive](https://drive.google.com/file/d/17-jiZh4D8cYUNkzsZMSfjLL_5gyr-3i_/view?usp=sharing) or [Baidu Drive](https://pan.baidu.com/s/1ZrAdLHwDDm1ZWqUpz94P7Q) (extract code: jams). This is a concise version containing only labeled images and labels. 14 | ## Description 15 | 16 | ### BamaPig2D 17 | When you download `BamaPig2D.zip` and unzip it, you will get two folders: `images` and `annotations`. 18 | 1. `images`: contains 3340 images used for training. 19 | 2. `annotation`: contains two files. `train_pig_cocostyle.json` is for training and `eval_pig_cocostyle.json` is used for testing. Both are in COCO style, and you can read them using COCO PythonAPI (see also [pig_pose_det]. Train split contains 3008 images and 10356 instances, while eval split contains 332 images and 1148 instances. 20 | ![dataset](pics/keypoint.jpg) 21 | ### BamaPig3D 22 | BamaPig3D dataset contains 1750 images with 70 ones annotated. The contents in each folder are described below. The annotations here are mainly `.json` file or `.txt` file which are more friendly to MATALB or C++ users. 23 | 24 | 1. `image` folder contains uncalibrated synchronized images. It has 10 folders. Each folder contains 1750 images of a single view. The camera names are `0`, `1`, `2`, `5`, `6`, `7`, `8`, `9`, `10`, `11`. Images are in `xxxxxx.jpg` name style. 25 | 2. `label_images` folder contains calibrated images organized the same to `image` folder. `label_images` also contains 2D annotations as `xxxxxx.json` in the same folder to `xxxxxx.jpg`. Each `xxxxxx.json` file is the output of [LabelMe software](https://pypi.org/project/labelme/) and you can use `code/demo_readBamaPig3D.py` to check how to parse these 2D information and visualize them together with uncalibrated images. The pigs are labeled in the following order: 26 | 33 | 34 | 35 | 3. `label_3d` folder contains 3D keypoints annotation. For pig `i` (`i=0,1,2,3`) and frame `k` (`k=0,25,...,1725`), the 3D keypoint file is `pig_{i}_frame_{k}.txt`. The 3D pig annotation follows the same order to 2D. 36 | Each txt file is a `23*3` matrix, with the 18, 20, 22, 23 rows always set zero. Invisible keypoints without 3D labels are set to zero. Therefore, only 19 keypoints are valid which names are defined in the order. See also BamaPig2D. 37 | 38 | 4. `label_mesh` is organized same to `label_3d`. The difference is, its keypoints totally come from the labeled mesh (i.e. the PIG model), whose pose parameters are stored in `label_pose_params`. You can use `bodymodel_np.py` and the PIG model files (see [PIG model]) to read these pose params and regress the keypoints from pose parameters. 39 | 40 | 5. `label_mix` is organized same to `label_3d`. It is the final 3D keypoint labeling combining `label_3d` and `label_mesh`. All the experiments in the paper are performed on this labeling. Please refer to the paper for detailed decription. 41 | 42 | 6. `boxes_pr` and `masks_pr` are detection results from [pig_silhouette_det] (a modified [PointRend](https://github.com/facebookresearch/detectron2/tree/main/projects/PointRend)). 43 | 44 | 7. `keypoints_hrnet` are keypoint detection results from [pig_pose_det] (a modified HRNet) using our weights pre-trained on BamaPig2D dataset. Note that, `boxes_pr`, `masks_pr` and `keypoints_hrnet` are the detection results used to generate evaluation results in Fig.2 and the video in Supplementary Video 1 of the paper. You can test other 3D reconstruction methods fairly based on these baseline results, or just use your own detection methods to generate another detection results. 45 | 46 | 8. `extrinsic_camera_params` contains 10 camera extrinsic paramters in `{camid}.txt` file. For example, for `00.txt`, it contains 6 float number, with the first three are camera rotation in axis-angle format, the last three are translation in xyz order. Unit is meter. `marker_3dpositions.txt` contains the 3d positions of 75 scene points for extrinsic camera parameter solving with PnP algorithm (see Supplementary Fig. 1 in the paper). `markerid_for_extrinsic_pnp.ppt` shows how these 75 points correspond to the scene. `markers{camid}.png` shows the projection of 3d scene points (red) and labeled 2d points on the image (green). It indicates how well the extrinsic parameters are solved. 47 | 48 | 9. `intrinsic_camera_params` contains two pickle files. You can also find the intrinsic parameters in `undistortion.py` file. 49 | 50 | ### BamaPig3D_pure_pickle 51 | This is a slim version of BamaPig3D, in which we remove `images`, `boxes_pr`, `keypoints_hrnet`, `masks_pr` folders. Only labeled images and labels are reserved. To save space, all label data are in `.pkl` format. `read_2dlabel_to_pickle` function in `visualize_BamaPig3D.py` shows how to encode 2D labels to pickle file. `label_mesh.pkl`, `label_3d.pkl` and `label_mix.pkl` are 70x4x19x3 matrices. `label_pose_params.pkl` is a dict seperating pose parameters to different parts, see information in the dict. 52 | 53 | ## Demo code requirements 54 | These functions are tested on Python 3.7 with conda virtual environment. The following python packages are necessary to run the codes in `code/` folder. Simply install the newest version. 55 | * scipy 56 | * numpy 57 | * opencv-python 58 | * videoio 59 | * ipython 60 | * tqdm 61 | * matplotlib 62 | * pyflann 63 | 64 | Specifically, after install anaconda (follow https://www.anaconda.com/ to install the newest version), you can create a conda virtual environment by running 65 | ```shell 66 | conda create -n MAMMAL python=3.7.9 67 | conda activate MAMMAL 68 | pip install scipy numpy opencv-python videoio 69 | pip install ipython 70 | pip install tqdm matplotlib pyflann-py3 71 | ``` 72 | It works for both windows 10 and ubuntu 20.04 (other mainstream windows and ubuntu version may work as well). If the installation of some packages fail, just try to install them again. If always fail, you may need to google the solution. 73 | 74 | ## Demo code description 75 | `utils.py` contains some keypoint structure definitions. 76 | 77 | `visualize_BamaPig2D.py` tells how to load and visualize 2d labels onto images, and generate Supplementary Fig. 4b and 4d. 78 | 79 | `visualize_BamaPig3D.py` tells how to load 2d keypoints in BamaPig3D dataset and generate Supplementary Fig. 8c and 8d. 80 | 81 | `bodymodel_np.py` is used to drive the PIG model. You should have prepared model files of the PIG model before run this file. 82 | 83 | `assemble_BamaPig3D.py` shows the procedure of Supplementary Fig. 8b. 84 | 85 | `undistortion.py` contains the intrinsic calibration parameters 86 | 87 | ## Train SLEAP using BamaPig2D 88 | [SLEAP](https://sleap.ai/) and [DeepLabCut](https://github.innominds.com/DeepLabCut) are the most popular multiple-animal pose estimation methods. Here, we provide an instruction on how to use BamaPig2D dataset to train SLEAP. Note that, we train SLEAP instead of DeepLabCut because only SLEAP supports COCO style dataset currently. 89 | 90 | 1. Install SLEAP v1.2.6 following their official instructions. 91 | 2. Download `BamaPig2D_sleap.zip`, unzip it, and you will get `train_pig_cocostyle_sleap.json`, `eval_pig_cocostyle_sleap.json` and `full_pig_cocostyle_sleap.json` files. We recommend to train SLEAP using `train_pig_cocostyle_sleap.json` or `full_pig_cocostyle_sleap.json`. Let's take `train_pig_cocostyle_sleap.json` as example. Put it under `{BamaPig2D_path}/images/` folder first. We put it under `images` instead of `annotations` folder because SLEAP load images from the folder where `.json` file lives. 92 | 3. After preparing the data, open SLEAP software, click `File->Import->COCO dataset`, wait for about half a minute before SLEAP load all the images. 93 | 4. Click `Predict->Run training` to open the training setting panel. 94 | 5. We recommend to use "top-down" structure. Set `Sigma for Centroids` as 8.00 and `Sigma for Nodes` as 4.00. Choose `unet` as model backbone, set `Max Stride` as 64. 95 | 6. Click `Run` to start the training process. It may take half a day to finish the training. 96 | 97 | ATTENTION! Because BamaPig2D dataset contains 3003 images for training (3340 for full dataset), SLEAP requires at least 17.6GB GPU memory to train the top-down model. Therefore, I trained SLEAP using a single NVIDIA RTX 3090Ti GPU which has 24GB memory. If you could not access such a high-end GPU, you could remove some annotation samples from the json file. For example, for NVIDIA RTX 2080Ti (11GB), you may reduce the training images to 1000. 98 | 99 | 100 | ## Citation 101 | If you use these datasets in your research, please cite the paper 102 | 103 | ```BibTex 104 | @article{MAMMAL, 105 | author = {An, Liang and Ren, Jilong and Yu, Tao and Hai, Tang and Jia, Yichang and Liu, Yebin}, 106 | title = {Three-dimensional surface motion capture of multiple freely moving pigs using MAMMAL}, 107 | booktitle = {}, 108 | month = {July}, 109 | year = {2022} 110 | } 111 | ``` -------------------------------------------------------------------------------- /visualize_BamaPig3D.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import cv2 3 | import json 4 | from tqdm import tqdm 5 | import pickle 6 | from utils import * 7 | from matplotlib.patches import Patch 8 | import matplotlib as mpl 9 | import matplotlib.pyplot as plt 10 | import os 11 | from visualize_BamaPig2D import COLORS, draw_keypoints, draw_mask 12 | 13 | # this function is used to transfer 2d labeled json files to unified pickle file for easier usage. 14 | # if you load BamaPig3D_pure_pickle, it already contains these two files. 15 | def read_2dlabel_to_pickle(BamaPig3D_folder): 16 | camids = [0,1,2,5,6,7,8,9,10,11] 17 | all_2D_points = np.zeros([70,10,4,19,3]) # store all 2d keypoints 18 | mask_all_frames = [] # store all silhouettes 19 | for i in tqdm(range(70)): 20 | mask_all_views = {} 21 | for index, camid in enumerate(camids): 22 | frameid = 25 * i 23 | outfilename = BamaPig3D_folder + "/label_images/cam{}/{:06d}.json".format(camid, frameid) 24 | mask_4 = [[], [], [], []] 25 | with open(outfilename, 'r') as f: 26 | data = json.load(f) 27 | for part in data['shapes']: 28 | if part['shape_type'] == 'point': 29 | point = np.asarray(part['points'][0], dtype=np.int32) 30 | group_id = part["group_id"] 31 | label = int(part['label']) 32 | if label == 18: 33 | label = 17 34 | elif label == 20: 35 | label = 18 36 | elif label > 18: 37 | print("error") 38 | return 39 | try: 40 | all_2D_points[i, index, group_id, label, 0:2] = point 41 | all_2D_points[i, index, group_id, label, 2] = 1 42 | except: 43 | from IPython import embed; embed() # debug 44 | exit() 45 | elif part['shape_type'] == 'polygon': 46 | group_id = part["group_id"] 47 | if len(part["points"]) > 0: 48 | mask_4[group_id].append(part["points"]) 49 | mask_all_views.update({camid:mask_4}) 50 | mask_all_frames.append(mask_all_views) 51 | with open(BamaPig3D_folder + "label_keypoints2d.pkl", 'wb') as f: 52 | pickle.dump(all_2D_points, f) 53 | with open(BamaPig3D_folder + "label_silhouettes2d.pkl", 'wb') as f: 54 | pickle.dump(mask_all_frames, f) 55 | 56 | 57 | ## This function load all 2D keypoint annotations and combine them into `output/points2d.pkl` file for 58 | def count_visibility(BamaPig3D_folder): 59 | camids = [0,1,2,5,6,7,8,9,10,11] 60 | all_2D_points = np.zeros([70,10,4,19,3]) 61 | if os.path.exists(BamaPig3D_folder + "/label_keypoints2d.pkl"): 62 | with open(BamaPig3D_folder + "/label_keypoints2d.pkl", 'rb') as f: 63 | all_2D_points = pickle.load(f) 64 | else: 65 | print("please run read_2dlabel_to_pickle() function first!") 66 | return 67 | cam_level1 = [3] 68 | cam_level2 = [0, 5, 6] 69 | cam_level3 = [1,2,7,8] 70 | cam_level4 = [4, 9] 71 | total_visible_point_num = all_2D_points[:,:,:,:,2].sum() 72 | total_point_num = 70 * 10 * 4 * 19 73 | print("visible_ratio of all keypoints : ", (float(total_visible_point_num) / float(total_point_num) ) ) 74 | print("visible_ratio for each part : ") 75 | for k in range(19): 76 | print("{:10s}".format(g_jointnames[g_all_parts[k]]), all_2D_points[:,:,:,k,2].sum() / (70 * 10 * 4) ) 77 | 78 | part_levels = [ 79 | [0,17,18], 80 | [1,2,3,4], 81 | [5,6,7,8,9,10,11,12,13,14,15,16] 82 | ] 83 | part_level_names = ["Trunk", "Head", "Limb"] 84 | for index, part_level in enumerate(part_levels): 85 | N = len(part_level) 86 | print(part_level_names[index]) 87 | print(" cam level1 : ", all_2D_points[:, cam_level1, :, :, 2][:,:,:,part_level].sum() / (70 * 4 * N * 1) ) 88 | print(" cam level2 : ", all_2D_points[:, cam_level2, :, :, 2][:,:,:,part_level].sum() / (70 * 4 * N * 3) ) 89 | print(" cam level3 : ", all_2D_points[:, cam_level3, :, :, 2][:,:,:,part_level].sum() / (70 * 4 * N * 4) ) 90 | print(" cam level4 : ", all_2D_points[:, cam_level4, :, :, 2][:,:,:,part_level].sum() / (70 * 4 * N * 2) ) 91 | 92 | print("ratio of visible to more than 1 views: ") 93 | for k in range(19): 94 | part_sum = all_2D_points[:,:,:,k,2].sum(axis=1) 95 | print("{:10s}".format(g_jointnames[g_all_parts[k]]), (part_sum > 1).sum() / (70 * 4) ) 96 | 97 | # This function draws Supplementary Fig. 8c in the paper. 98 | def draw_visibility_level(BamaPig3D_folder): 99 | if not os.path.exists(BamaPig3D_folder + "/label_keypoints2d.pkl"): 100 | print("Please run read_2dlabel_to_pickle() function to generate label_keypoints2d.pkl") 101 | return 102 | with open(BamaPig3D_folder + "/label_keypoints2d.pkl", 'rb') as f: 103 | all_2D_points = pickle.load(f) 104 | 105 | part_levels = [ 106 | [0,17,18], # trunk 107 | [1,2,3,4], # head 108 | [5,6,7,8,9,10,11,12,13,14,15,16] # limbs 109 | ] 110 | 111 | cam_level1 = [3] 112 | cam_level2 = [0, 5, 6] 113 | cam_level3 = [1,2,7,8] 114 | cam_level4 = [4, 9] 115 | 116 | data = np.zeros([4,3]) 117 | for index, part_level in enumerate(part_levels): 118 | N = len(part_level) 119 | data[0,index] = all_2D_points[:, cam_level1, :, :, 2][:,:,:,part_level].sum() / (70 * 4 * N * 1) 120 | data[1,index] = all_2D_points[:, cam_level2, :, :, 2][:,:,:,part_level].sum() / (70 * 4 * N * 3) 121 | data[2,index] = all_2D_points[:, cam_level3, :, :, 2][:,:,:,part_level].sum() / (70 * 4 * N * 4) 122 | data[3,index] = all_2D_points[:, cam_level4, :, :, 2][:,:,:,part_level].sum() / (70 * 4 * N * 2) 123 | 124 | mpl.rc('font', family='Arial') 125 | fig = plt.figure(figsize=(1.8, 1.4)) 126 | colormaps = np.loadtxt("colormaps/tab.txt") / 255 127 | part_level_names = ["Trunk", "Head", "Limbs"] 128 | xs = np.asarray([0,1,2]) 129 | for x_index in range(3): 130 | plt.bar(xs-0.3, data[0,:], width=0.2, color=colormaps[0], linewidth=0.5) 131 | plt.bar(xs-0.1, data[1,:], width=0.2, color=colormaps[1], linewidth=0.5) 132 | plt.bar(xs+0.1, data[2,:], width=0.2, color=colormaps[2], linewidth=0.5) 133 | plt.bar(xs+0.3, data[3,:], width=0.2, color=colormaps[3], linewidth=0.5) 134 | plt.legend(["Top view", "Corner views", "Middle views", "Side views"], fontsize=6, frameon=False, ncol=1) 135 | plt.ylim(0,1) 136 | plt.xticks(xs, part_level_names, fontsize=7) 137 | plt.yticks([0,0.2,0.4,0.6,0.8,1], [0,20,40,60,80,100], fontsize=7) 138 | plt.xlabel("Body Parts", fontsize=7) 139 | plt.ylabel("Percentage of Visible\n Keypoints (%)", fontsize=7) 140 | ax = fig.get_axes()[0] 141 | ax.spines['right'].set_visible(False) 142 | ax.spines['top'].set_visible(False) 143 | for line in ["bottom", "left", "right"]: 144 | ax.spines[line].set_linewidth(0.5) 145 | ax.xaxis.set_tick_params(width=0.5) 146 | ax.yaxis.set_tick_params(width=0.5) 147 | 148 | plt.savefig("output/supp_fig_8c.png", dpi=1000, bbox_inches='tight', pad_inches=0) 149 | # plt.savefig("output/supp_fig_8c.svg", dpi=1000, bbox_inches='tight', pad_inches=0) 150 | 151 | # This function draws Supplementary Fig. 8d in the paper. 152 | def draw_keypoint_visibility_hist(BamaPig3D_folder): 153 | jointnames = [g_jointnames[k] for k in g_all_parts] 154 | if not os.path.exists(BamaPig3D_folder + "/label_keypoints2d.pkl"): 155 | print("Please run read_2dlabel_to_pickle() function first!") 156 | return 157 | with open(BamaPig3D_folder + "/label_keypoints2d.pkl", 'rb') as f: 158 | all_2D_points = pickle.load(f) 159 | data = np.zeros([19,4]) 160 | for k in range(19): 161 | part_sum = all_2D_points[:,:,:,k,2].sum(axis=1) 162 | data[k,2] = (part_sum > 1).sum() / (70 * 4) 163 | data[k,3] = (part_sum > 4).sum() / (70 * 4) 164 | data[k,1] = (part_sum == 1).sum() / (70 * 4) 165 | data[k,0] = 1 166 | mpl.rc('font', family='Arial') 167 | 168 | fig = plt.figure(figsize=(4,1.4)) 169 | 170 | colormaps = np.loadtxt("colormaps/tab20c.txt") / 255 171 | 172 | xs = np.arange(0,19,1) 173 | plt.bar(xs, data[:,0], color=colormaps[3], edgecolor=(0,0,0), lw=0.5) 174 | plt.bar(xs, data[:,1] + data[:,2], color=colormaps[2], edgecolor=(0,0,0), lw=0.5) 175 | plt.bar(xs, data[:,2], color=colormaps[1], edgecolor=(0,0,0), lw=0.5) 176 | plt.bar(xs, data[:,3], color=colormaps[0], edgecolor=(0,0,0), lw=0.5) 177 | 178 | 179 | plt.xticks(xs,jointnames, rotation=45, ha='right', fontsize=7) 180 | plt.yticks([0,0.2,0.4,0.6,0.8,1],[0,20,40,60,80,100], fontsize=7) 181 | legend_elements = [ 182 | Patch(facecolor=colormaps[3], edgecolor='black', label="Visible to 0 view", linewidth=0.5), 183 | Patch(facecolor=colormaps[2], edgecolor='black', label="Visible to 1 view", linewidth=0.5), 184 | Patch(facecolor=colormaps[1], edgecolor='black', label="Visible to 2~4 views", linewidth=0.5), 185 | Patch(facecolor=colormaps[0], edgecolor='black', label="Visible to 5~10 views", linewidth=0.5), 186 | ] 187 | plt.legend(handles=legend_elements, fontsize=6, ncol=2, loc='upper left', bbox_to_anchor=(0.0, 1.3), frameon=False) 188 | ax = fig.get_axes()[0] 189 | ax.spines['right'].set_visible(False) 190 | ax.spines['top'].set_visible(False) 191 | for line in ["bottom", "left", "right", "top"]: 192 | ax.spines[line].set_linewidth(0.5) 193 | ax.xaxis.set_tick_params(width=0.5) 194 | ax.yaxis.set_tick_params(width=0.5) 195 | plt.xlim(-1, 19) 196 | 197 | plt.xlabel("", fontsize=7) 198 | plt.ylabel("Percentage of Visible Keypoints (%)", fontsize=7) 199 | plt.ylim(0,1) 200 | 201 | plt.savefig("output/supp_fig_8d.png", dpi=1000, bbox_inches='tight', pad_inches=0.01) 202 | # plt.savefig("output/supp_fig_8d.svg", dpi=1000, bbox_inches='tight', pad_inches=0.01) 203 | 204 | # demo for how to load and draw silhouettes used in BamaPig3D dataset. 205 | def demo_draw_mask(BamaPig3D_folder): 206 | camid = 0 207 | frameid = 0 208 | imgfile = BamaPig3D_folder + "/label_images/cam{}/{:06d}.jpg".format(camid, frameid) 209 | img = cv2.imread(imgfile) 210 | 211 | camids = [0,1,2,5,6,7,8,9,10,11] 212 | with open(BamaPig3D_folder + "/label_silhouettes2d.pkl", 'rb') as f: 213 | mask_label = pickle.load(f) 214 | 215 | mask_label_current = mask_label[frameid][camid] 216 | 217 | for pid in range(4): 218 | draw_mask(img, np.asarray(mask_label_current[pid]), COLORS[pid]) 219 | if not os.path.exists("output"): 220 | os.makedirs("output") 221 | cv2.imwrite("output/demo_sil_BamaPig3D_frame0_cam0.png", img) 222 | 223 | # 224 | def demo_how_to_project_points_with_extrinsic_params(BamaPig3D_folder): 225 | frameid = 0 226 | camids = [0,1,2,5,6,7,8,9,10,11] 227 | 228 | with open(BamaPig3D_folder + "/intrinsic_camera_params/distortion_info.pkl", 'rb') as f: 229 | intrinsic_params = pickle.load(f) 230 | K = intrinsic_params["newcameramtx"] 231 | 232 | for camid in camids: 233 | undist_image = cv2.imread(BamaPig3D_folder + "/label_images/cam{}/{:06d}.jpg".format(camid, frameid)) 234 | 235 | extrinsic_params = np.loadtxt(BamaPig3D_folder + "/extrinsic_camera_params/{:02d}.txt".format(camid)).squeeze() 236 | # extrinsic_params = np.loadtxt("H:/MAMMAL_core/data/calibdata/adjust/{:02d}.txt".format(camid)).squeeze() 237 | R = cv2.Rodrigues(extrinsic_params[0:3])[0] 238 | T = extrinsic_params[3:] 239 | 240 | # if you use BamaPig3D_pure_pickle, just load pickle files for all 3D keypoints GT 241 | for pid in range(4): 242 | points3d = np.loadtxt(BamaPig3D_folder + "/label_mix/pig_{}_frame_{:06d}.txt".format(pid, frameid)) 243 | points3d = points3d[g_all_parts] # remove invalid ones to get 19 keypoints only 244 | ## KEY process 245 | points2d = (points3d @ R.T + T) @ K.T 246 | points2d = points2d[:,0:2] / points2d[:,2:] 247 | draw_keypoints(undist_image, points2d, pid, g_bones_19) 248 | 249 | cv2.imwrite("output/demo_BamaPig3D_proj_3d_keypoints_{}.png".format(camid), undist_image) 250 | 251 | if __name__ == "__main__": 252 | # To run this file, you should change this folder to your own BamaPig3D dataset path 253 | BamaPig3D_folder = "H:/examples/BamaPig3D/" 254 | # count_visibility(BamaPig3D_folder) 255 | 256 | # output supp_fig_8c.png 257 | draw_visibility_level(BamaPig3D_folder) 258 | 259 | # output supp_fig_8d.png 260 | draw_keypoint_visibility_hist(BamaPig3D_folder) 261 | 262 | # demo for how to load and draw sihouettes 263 | demo_draw_mask(BamaPig3D_folder) 264 | 265 | # this function shows how to project 3D keypoints to 2D 266 | demo_how_to_project_points_with_extrinsic_params(BamaPig3D_folder) --------------------------------------------------------------------------------