def imgs_generator(): camNetAll, camNetLeft, camNetRight = _getCamNets(args) camNetLeft.triangulate() camNetLeft.bundle_adjust(cam_id_list=(0, 1, 2), unique=False, prior=True) camNetRight.triangulate() camNetRight.bundle_adjust(cam_id_list=(0, 1, 2), unique=False, prior=True) camNetAll.triangulate() camNetAll.points3d_m = procrustes_seperate(camNetAll.points3d_m) camNetAll.points3d_m = normalize_pose_3d(camNetAll.points3d_m, rotate=True) camNetAll.points3d_m = filter_batch(camNetAll.points3d_m) for cam in camNetAll: cam.points2d = smooth_pose2d(cam.points2d) def stack(img_id): row1 = np.hstack([ _compute_2d_img(camNetLeft, img_id, cam_id) for cam_id in (0, 1, 2) ]) row2 = np.hstack([ _compute_2d_img(camNetRight, img_id, cam_id) for cam_id in (0, 1, 2) ]) row3 = np.hstack([ _compute_3d_img(camNetAll, img_id, cam_id) for cam_id in (2, 3, 4) ]) img = np.vstack([row1, row2, row3]) return img for img_id in range(args.num_images): print("img_id:" + str(img_id)) yield stack(img_id)
def save_pose(self): pts2d = np.zeros((7, self.state.num_images, config["num_joints"], 2), dtype=float) # pts3d = np.zeros((self.cfg.num_images, self.cfg.num_joints, 3), dtype=float) for cam in self.camNetAll: pts2d[cam.cam_id, :] = cam.points2d.copy() # overwrite by manual correction count = 0 for cam_id in range(config["num_cameras"]): for img_id in range(self.state.num_images): if self.state.db.has_key(cam_id, img_id): pt = self.state.db.read(cam_id, img_id) * config["image_shape"] pts2d[cam_id, img_id, :] = pt count += 1 if "fly" in config["name"]: # some post-processing for body-coxa for cam_id in range(len(self.camNetAll.cam_list)): for j in range(config["skeleton"].num_joints): if config["skeleton"].camera_see_joint( cam_id, j) and config["skeleton"].is_tracked_point( j, config["skeleton"].Tracked.BODY_COXA): pts2d[cam_id, :, j, 0] = np.median(pts2d[cam_id, :, j, 0]) pts2d[cam_id, :, j, 1] = np.median(pts2d[cam_id, :, j, 1]) dict_merge = self.camNetAll.save_network(path=None) dict_merge["points2d"] = pts2d # take a copy of the current points2d pts2d_orig = np.zeros( (7, self.state.num_images, config["num_joints"], 2), dtype=float) for cam_id in range(config["num_cameras"]): pts2d_orig[cam_id, :] = self.camNetAll[cam_id].points2d.copy() # ugly hack to temporarly incorporate manual corrections c = 0 for cam_id in range(config["num_cameras"]): for img_id in range(self.state.num_images): if self.state.db.has_key(cam_id, img_id): pt = self.state.db.read(cam_id, img_id) * config["image_shape"] self.camNetAll[cam_id].points2d[img_id, :] = pt c += 1 print("Replaced points2d with {} manual correction".format(count)) # do the triangulation if we have the calibration if self.camNetLeft.has_calibration() and self.camNetLeft.has_pose(): self.camNetAll.triangulate() pts3d = self.camNetAll.points3d_m dict_merge["points3d"] = pts3d # apply procrustes if config["procrustes_apply"]: print("Applying Procrustes on 3D Points") dict_merge["points3d"] = procrustes_seperate( dict_merge["points3d"]) # put old values back for cam_id in range(config["num_cameras"]): self.camNetAll[cam_id].points2d = pts2d_orig[cam_id, :].copy() pickle.dump( dict_merge, open( os.path.join( self.folder_output, "pose_result_{}.pkl".format(self.folder.replace("/", "_")), ), "wb", ), ) print("Saved the pose at: {}".format( os.path.join( self.folder_output, "pose_result_{}.pkl".format(self.folder.replace("/", "_")), )))