def _make_video(video_path, imgs): """Code used to generate a video using cv2. Parameters: video_path: a path ending with .mp4, for instance: "/results/pose2d.mp4" imgs: an iterable or generator with the images to turn into a video """ first_frame = next(imgs) imgs = itertools.chain([first_frame], imgs) shape = int(first_frame.shape[1]), int(first_frame.shape[0]) logger.debug('Saving video to: ' + video_path) fourcc = cv2.VideoWriter_fourcc(*'mp4v') fps = 30 output_shape = _resize(current_shape=shape, new_width=video_width) logger.debug('Video size is: {}'.format(output_shape)) video_writer = cv2.VideoWriter(video_path, fourcc, fps, output_shape) progress_bar = tqdm if logger.info_enabled() else lambda x: x for img in progress_bar(imgs): resized = cv2.resize(img, output_shape) rgb = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB) video_writer.write(rgb) video_writer.release() logger.info('Video created at {}\n'.format(video_path))
def save_pose(self): """Saves the pose estimation results to a file in the output folder.""" pts2d = self.corrected_points2d_matrix() dict_merge = self.camNetAll.save_network(path=None) pts2d_orig = self.camNetAll.get_points2d_matrix() # temporarily incorporate corrected values self.camNetAll.set_points2d_matrix(pts2d) self.post_process(pts2d) dict_merge["points2d"] = pts2d if self.camNetLeft.has_calibration() and self.camNetLeft.has_pose(): self.camNetAll.triangulate() pts3d = self.camNetAll.points3d_m if config["procrustes_apply"]: print("Applying Procrustes on 3D Points") pts3d = procrustes_seperate(pts3d) dict_merge["points3d"] = pts3d else: logger.debug("Triangulation skipped.") # put uncorrected values back self.camNetAll.set_points2d_matrix(pts2d_orig) save_path = os.path.join( self.output_folder, "pose_result_{}.pkl".format(self.input_folder.replace("/", "_")), ) pickle.dump(dict_merge, open(save_path, "wb")) print(f"Saved the pose at: {save_path}")
def _compute_mean(self): meanstd_file = config["mean"] if isfile(meanstd_file): meanstd = torch.load(meanstd_file) else: raise FileNotFoundError mean = torch.zeros(3) std = torch.zeros(3) for k in self.annotation_key: img_path = os.path.join(self.data_folder, k[FOLDER_NAME], k[IMAGE_NAME] + ".jpg") img = load_image(img_path) # CxHxW mean += img.view(img.size(0), -1).mean(1) std += img.view(img.size(0), -1).std(1) mean /= len(self) std /= len(self) meanstd = {"mean": mean, "std": std} torch.save(meanstd, meanstd_file) if self.is_train: logger.debug( " Mean: %.4f, %.4f, %.4f" % (meanstd["mean"][0], meanstd["mean"][1], meanstd["mean"][2])) logger.debug( " Std: %.4f, %.4f, %.4f" % (meanstd["std"][0], meanstd["std"][1], meanstd["std"][2])) return meanstd["mean"], meanstd["std"]
def find_default_camera_ordering(input_folder): """Uses regexes to infer the correct camera ordering based on folder path. This is useful for Ramdya's Lab as a given data acquisition agent (say CLC) always uses the same camera ordering. Parameters: input_folder: the folder path on which to run the regexes. """ known_users = [ (r"/CLC/", [0, 6, 5, 4, 3, 2, 1]), (r"data/test", [0, 1, 2, 3, 4, 5, 6]), ] # input_folder = str(input_folder) # use `str` in case pathlib.Path instance def match(regex): return re.search(regex, input_folder) candidates = [order for (regex, order) in known_users if match(regex)] if candidates: order = candidates[0] logger.debug(f"Default camera ordering found: {order}") return np.array(order)
def load_heatmap(hm_path, shape): logger.debug("Heatmap shape: {}".format(shape)) heatmap = np.memmap(filename=hm_path, mode="r", shape=shape, dtype="float32") return heatmap
def write_camera_order(folder, cidread2cid): assert os.path.isdir( folder ), "Trying to write_camera_order into {}, which is not a folder".format( folder) path = os.path.join(folder, "cam_order") logger.debug("Writing the camera ordering {} into folder {}".format( cidread2cid, folder)) # print("Saving camera order {}: {}".format(path, cidread2cid)) np.save(path, cidread2cid)
def calc_essential_matrix(points2d_1, points2d_2, intr): E, mask = cv2.findEssentialMat( points1=points2d_1, points2=points2d_2, cameraMatrix=intr, method=cv2.RANSAC, prob=0.9999, threshold=5, ) logger.debug("Essential matrix inlier ratio: {}".format( np.sum(mask) / mask.shape[0])) return E, mask
def solve_belief_propagation(cam_list, img_id, bone_param, num_peak=10, prior=None): # find all the connected parts j_id_list_list = [[ j for j in range(config["skeleton"].num_joints) if config["skeleton"].limb_id[j] == limb_id ] for limb_id in range(config["skeleton"].num_limbs)] chain_list = list() for j_id_l in j_id_list_list: visible = np.zeros(shape=(len(j_id_l), ), dtype=np.int) for cam in cam_list: visible += [ config["skeleton"].camera_see_joint(cam.cam_id, j_id) for j_id in j_id_l ] if np.all(visible >= 2): chain_list.append( LegBP( cam_list=cam_list, img_id=img_id, j_id_list=j_id_l, bone_param=bone_param, num_peak=num_peak, prior=prior, )) else: pass # logger.debug("Joints {} is not visible from at least two cameras".format(j_id_l)) logger.debug([[len(leg[i].candid_list) for i in range(len(leg.jointbp))] for leg in chain_list]) for chain in chain_list: chain.propagate() chain.solve() # read the best 2d locations points2d_list = [ np.zeros((config["skeleton"].num_joints, 2), dtype=float) for _ in range(len(cam_list)) ] for leg in chain_list: for cam_idx in range(len(cam_list)): for idx, j_id in enumerate(leg.j_id_list): points2d_list[cam_idx][j_id] = leg[idx][ leg[idx].argmin].p2d[cam_idx] return points2d_list.copy()
def load_network(self, calib): d = calib if calib is None: return None for cam in self.cam_list: if cam.cam_id in d and d[cam.cam_id]: cam.set_R(d[cam.cam_id]["R"]) cam.set_tvec(d[cam.cam_id]["tvec"]) cam.set_intrinsic(d[cam.cam_id]["intr"]) cam.set_distort(d[cam.cam_id]["distort"]) else: logger.debug("Camera {} is not on the calibration file".format( cam.cam_id)) return d["meta"]
def process_folder(model, loader, unlabeled, output_folder, overwrite, num_classes, acc_joints): save_path_pred, save_path_heatmap = ( get_save_path_pred(unlabeled, output_folder), get_save_path_heatmap(unlabeled, output_folder), ) if os.path.isfile(save_path_pred) and not overwrite: logger.info("Prediction file exists, skipping pose estimation") return None, None elif os.path.isfile(save_path_pred) and overwrite: logger.info("Overwriting existing predictions") save_path_heatmap.parent.mkdir(exist_ok=True, parents=True) save_path_pred.parent.mkdir(exist_ok=True, parents=True) logger.debug(f"creaint heatmap path: {save_path_heatmap}") heatmap = np.memmap( filename=save_path_heatmap, dtype="float32", mode="w+", shape=( config["num_cameras"] + 1, loader.dataset.greatest_image_id() + 1, config["num_predict"], config["heatmap_shape"][0], config["heatmap_shape"][1], ), ) # num_cameras+1 for the mirrored camera 3 logger.debug(f"creating heatmap shape: {heatmap.shape}") pred, heatmap, _, _, _ = step( loader=loader, model=model, optimizer=None, mode=Mode.test, heatmap=heatmap, epoch=0, num_classes=num_classes, acc_joints=acc_joints, ) _, cid2cidread = read_camera_order( get_output_path(unlabeled, output_folder)) cid_to_reverse = config["flip_cameras"] cid_read_to_reverse = [cid2cidread[cid] for cid in cid_to_reverse] pred = flip_pred(pred, cid_read_to_reverse) logger.debug("Flipping heatmaps") heatmap = flip_heatmap(heatmap, cid_read_to_reverse) logger.debug("Flipping heatmaps") save_dict(pred, save_path_pred) if type(heatmap) != np.memmap: save_dict(heatmap, save_path_heatmap) print(pred.shape) return pred, heatmap
def calibrate(self, cam_id_list=None): assert self.cam_list ignore_joint_list = config["skeleton"].ignore_joint_id if cam_id_list is None: cam_id_list = range(self.num_cameras) self.reprojection_error() ( x0, points_2d, n_cameras, n_points, camera_indices, point_indices, ) = self.prepare_bundle_adjust_param(cam_id_list) logger.debug(f"Number of points for calibration: {n_points}") A = bundle_adjustment_sparsity( n_cameras, n_points, camera_indices, point_indices ) res = least_squares( residuals, x0, jac_sparsity=A, verbose=2 if logger.debug_enabled() else 0, x_scale="jac", ftol=1e-4, method="trf", args=( self.cam_list, n_cameras, n_points, camera_indices, point_indices, points_2d, ), max_nfev=1000, ) logger.debug( "Bundle adjustment, Average reprojection error: {}".format( np.mean(np.abs(res.fun)) ) ) self.triangulate() return res
def get_max_img_id(path): print(path) bound_low = 0 bound_high = 100000 curr = (bound_high + bound_low) // 2 while bound_high - bound_low > 1: if image_exists_img_id(path, curr): bound_low = curr else: bound_high = curr curr = (bound_low + bound_high) // 2 if not image_exists_img_id(path, curr): logger.debug("Cannot find image at {} with img_id {}".format(path, curr)) raise FileNotFoundError("No image found.") return curr
def reprojection_error(self, cam_indices=None, ignore_joint_list=None): if ignore_joint_list is None: ignore_joint_list = config["skeleton"].ignore_joint_id if cam_indices is None: cam_indices = range(len(self.cam_list)) err_list = list() for (img_id, j_id, _), _ in np.ndenumerate(self.points3d_m): p3d = self.points3d_m[img_id, j_id].reshape(1, 3) if j_id in ignore_joint_list: continue for cam in self.cam_list: if not config["skeleton"].camera_see_joint(cam.cam_id, j_id): continue err_list.append((cam.project(p3d) - cam[img_id, j_id]).ravel()) err_mean = np.mean(np.abs(err_list)) logger.debug("Ignore_list {}: {:.4f}".format(ignore_joint_list, err_mean)) return err_list
def reprojection_error(self): ignore_joint_list = config["skeleton"].ignore_joint_id s = self.points3d.shape err = np.zeros((len(self.cam_list), s[0], s[1])) for (img_id, j_id, cam_idx) in product( range(s[0]), range(s[1]), range(len(self.cam_list)) ): p3d = self.points3d[img_id, j_id].reshape(1, 3) if ( config["skeleton"].camera_see_joint(self[cam_idx].cam_id, j_id) and not j_id in ignore_joint_list and not np.any(self.points3d[img_id, j_id, :] == 0) and not np.any(self[cam_idx][img_id, j_id, :] == 0) ): err[cam_idx, img_id, j_id] = np.sum( np.abs(self[cam_idx].project(p3d) - self[cam_idx][img_id, j_id]) ) err_mean = np.mean(np.abs(err)) logger.debug("Ignore_list {}: {:.4f}".format(ignore_joint_list, err_mean)) return err
def read_camera_order(folder): assert os.path.isdir( folder ), "Trying to call read_camera_order on {}, which is not a folder".format( folder) path = os.path.join(folder, "./cam_order.npy") if os.path.isfile(path): order = np.load(file=path, allow_pickle=True) else: order = np.arange(config["num_cameras"]) write_camera_order(folder, order) logger.debug( "Could not find camera order under {}. Writing the default ordering {}." .format(folder, order)) cidread2cid = order.copy() cid2cidread = np.zeros(cidread2cid.size, dtype=int) for cidread, cid in enumerate(cidread2cid): cid2cidread[cid] = cidread return cidread2cid, cid2cidread
def __init__( self, image_folder, output_folder, num_images=900, cam_list=None, cam_id_list=range(config["num_cameras"]), cid2cidread=None, pred=None, ): self.image_folder = image_folder self.output_folder = output_folder self.points3d = None self.num_images = num_images self.num_cameras = len(cam_id_list) self.cid2cidread = ( cid2cidread if cid2cidread is not None else read_camera_order(self.output_folder)[0] ) if cam_list: logger.debug("Camera list is already given, skipping loading.") self.cam_list = cam_list return self.cam_list = list() pred_path = find_pred_path(self.output_folder) if pred_path is not None: logger.debug("no pred file under {}".format(self.output_folder)) pred = np.load(file=pred_path, mmap_mode="r", allow_pickle=True) pred = pred[:, : self.num_images] else: pred = None for cam_id in cam_id_list: cam_id_read = self.cid2cidread[cam_id] self.cam_list.append( Camera( cid=cam_id, cid_read=cam_id_read, image_folder=image_folder, points2d=pred2points2d( pred, cam_id, cam_id_read, config["image_shape"] ) if pred is not None else pred, hm=None, ) ) calibration_path = find_calib_path(self.output_folder) if calibration_path is not None: calibration = np.load(file=calibration_path, allow_pickle=True) logger.debug("Reading calibration from {}".format(self.output_folder)) _ = self.load_network(calibration)
def read_manual_corrections(d, output_folder, manual_path_list, cidread2cid_global, num_classes): pose_corr_path_list = [] for root in manual_path_list: logger.debug("Searching recursively: {}".format(root)) pose_corr_path_list.extend(find_pose_corr_recursively(root)) logger.debug("Number of manual correction files: {}".format( len(pose_corr_path_list))) for path in pose_corr_path_list: d = pickle.load(open(path, "rb")) folder_name = d["folder"] key_folder_name = folder_name if folder_name not in cidread2cid_global: cam_folder = os.path.join(folder_name, output_folder) cidread2cid, cid2cidread = read_camera_order(cam_folder) cidread2cid_global[key_folder_name] = cidread2cid for cid in range(config["num_cameras"]): for img_id, points2d in d[cid].items(): cid_read = cidread2cid[key_folder_name].tolist().index(cid) key = (key_folder_name, constr_img_name(cid_read, img_id)) num_heatmaps = points2d.shape[0] pts = np.zeros((2 * num_classes, 2), dtype=np.float) if cid < 3: pts[:num_heatmaps // 2, :] = points2d[:num_heatmaps // 2, :] elif 3 < cid < 7: pts[num_classes:num_classes + (num_heatmaps // 2), :] = points2d[num_heatmaps // 2:, :] elif cid == 3: continue else: raise NotImplementedError d[key] = pts
def load_weights(model, resume: str): if isfile(resume): logger.debug("Loading checkpoint '{}'".format(resume)) checkpoint = (torch.load(resume) if torch.cuda.is_available() else torch.load(resume, map_location=torch.device("cpu"))) if "mpii" in resume: # weights for sh trained on mpii dataset logger.debug("Removing input/output layers") ignore_weight_list_template = [ "module.score.{}.bias", "module.score.{}.weight", "module.score_.{}.weight", ] ignore_weight_list = list() for i in range(8): for template in ignore_weight_list_template: ignore_weight_list.append(template.format(i)) for k in ignore_weight_list: if k in checkpoint["state_dict"]: checkpoint["state_dict"].pop(k) state = model.state_dict() state.update(checkpoint["state_dict"]) logger.debug(model.state_dict()) logger.debug(checkpoint["state_dict"]) model.load_state_dict(state, strict=False) else: pretrained_dict = checkpoint["state_dict"] model.load_state_dict(pretrained_dict, strict=False) logger.debug("Loaded checkpoint '{}' (epoch {})".format( resume, checkpoint["epoch"])) else: logger.debug("=> no checkpoint found at '{}'".format(resume)) raise FileNotFoundError
def main(args): logger.debug("Creating model '{}', stacks={}, blocks={}".format( args.arch, args.stacks, args.blocks)) model = hg( num_stacks=args.stacks, num_blocks=args.blocks, num_classes=args.num_classes, num_feats=args.features, inplanes=args.inplanes, init_stride=args.stride, ) model = on_cuda(torch.nn.DataParallel(model)) optimizer = torch.optim.RMSprop( model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay, ) scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, verbose=True, patience=5) if args.resume: load_weights(model, args.resume) logger.debug("Total params: %.2fM" % (sum(p.numel() for p in model.parameters()) / 1000000.0)) if args.unlabeled: loader = DataLoader( DrosophilaDataset( data_folder=args.data_folder, train=False, sigma=args.sigma, session_id_train_list=None, folder_train_list=None, img_res=args.img_res, hm_res=args.hm_res, augmentation=False, evaluation=True, unlabeled=args.unlabeled, num_classes=args.num_classes, max_img_id=min(get_max_img_id(args.unlabeled), args.max_img_id), output_folder=args.output_folder, ), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=False, drop_last=False, ) pred, heatmap = process_folder( model, loader, args.unlabeled, args.output_folder, args.overwrite, num_classes=args.num_classes, acc_joints=args.acc_joints, ) return pred, heatmap else: train_loader, val_loader = create_dataloader() lr = args.lr best_acc = 0 for epoch in range(args.start_epoch, args.epochs): logger.debug("\nEpoch: %d | LR: %.8f" % (epoch + 1, lr)) _, _, _, _, _ = step( loader=train_loader, model=model, optimizer=optimizer, mode=Mode.train, heatmap=False, epoch=epoch, num_classes=args.num_classes, acc_joints=args.acc_joints, ) val_pred, _, val_loss, val_acc, val_mse = step( loader=val_loader, model=model, optimizer=optimizer, mode=Mode.test, heatmap=False, epoch=epoch, num_classes=args.num_classes, acc_joints=args.acc_joints, ) scheduler.step(val_loss) is_best = val_acc > best_acc best_acc = max(val_acc, best_acc) save_checkpoint( { "epoch": epoch + 1, "arch": args.arch, "state_dict": model.state_dict(), "optimizer_state_dict": optimizer.state_dict(), "best_acc": best_acc, "optimizer": optimizer.state_dict(), "image_shape": args.img_res, "heatmap_shape": args.hm_res, }, val_pred, is_best, checkpoint=args.checkpoint, snapshot=args.snapshot, )
def __init__( self, data_folder, data_corr_folder=None, img_res=None, hm_res=None, train=True, sigma=1, jsonfile="drosophilaimaging-export.json", session_id_train_list=None, folder_train_list=None, augmentation=False, evaluation=False, unlabeled=None, num_classes=config["num_predict"], max_img_id=None, output_folder=None, ): self.train = train self.data_folder = data_folder # root image folders self.data_corr_folder = data_corr_folder self.json_file = os.path.join(jsonfile) self.is_train = train # training set or test set self.img_res = img_res self.hm_res = hm_res self.sigma = sigma self.augmentation = augmentation self.evaluation = evaluation self.unlabeled = unlabeled self.num_classes = num_classes self.max_img_id = max_img_id self.cidread2cid = dict() self.output_folder = output_folder self.manual_path_list = [] self.session_id_train_list = session_id_train_list self.folder_train_list = folder_train_list if self.output_folder is None: raise ValueError( "Please provide an output_folder relative to images folder") assert (not self.evaluation or not self.augmentation) # self eval then not augmentation assert not self.unlabeled or evaluation # if unlabeled then evaluation self.annotation_dict = dict() if isfile(self.json_file) and not self.unlabeled: logger.debug("Searching for json file") read_json( self.annotation_dict, self.json_file, self.folder_train_list, self.cidread2cid, ) if not self.unlabeled and self.train and self.manual_path_list: logger.debug("Searching for manual corrections") read_manual_corrections( self.annotation_dict, self.output_folder, self.manual_path_list, self.cidread2cid, self.num_classes, ) if self.unlabeled: logger.debug("Searching unlabeled") read_unlabeled_folder( self.annotation_dict, self.unlabeled, self.output_folder, self.cidread2cid, self.max_img_id, ) # make sure data is in the folder for folder_name, image_name in self.annotation_dict.copy().keys(): image_file = os.path.join( self.data_folder, folder_name.replace("_network", ""), image_name + ".jpg", ) if not os.path.isfile(image_file): self.annotation_dict.pop((folder_name, image_name), None) print("FileNotFound: {}/{} ".format(folder_name, image_name)) normalize_annotations(self.annotation_dict, self.num_classes, self.cidread2cid) self.annotation_key = list(self.annotation_dict.keys()) if self.evaluation: # sort keys self.annotation_key.sort(key=lambda x: x[0] + "_" + x[1].split("_") [3] + "_" + x[1].split("_")[1]) self.mean, self.std = self._compute_mean() logger.debug("Folders inside {} data: {}".format( "train" if self.train else "validation", set([k[0] for k in self.annotation_key]), )) logger.debug( "Successfully imported {} Images in Drosophila Dataset".format( len(self)))
mse_tarsus=am["tarsus_tip"].avg, ) bar.next() bar.finish() return (predictions, heatmap, am["losses"].avg, am["acces"].avg, am["mse"].avg) if __name__ == "__main__": # prepare parser = create_parser() args = parser.parse_args() args.train_joints = np.arange(args.num_classes) logger.debug(f"Training joints: {args.train_joints}") logger.debug(f"Acc joints: {args.acc_joints}") args.checkpoint = (args.checkpoint.replace(" ", "_").replace("(", "_").replace( ")", "_")) args.checkpoint = os.path.join( args.checkpoint, get_time() + "_{}_{}_{}_{}_{}_{}".format( "predict" if args.unlabeled else "training", args.arch, args.stacks, args.img_res, args.blocks, args.name, ),
def prepare_bundle_adjust_param(self, camera_id_list=None, ignore_joint_list=None, max_num_images=1000): logger.debug( "Calibration ignore joint list {}".format(ignore_joint_list)) if ignore_joint_list is None: ignore_joint_list = config["skeleton"].ignore_joint_id if camera_id_list is None: camera_id_list = list(range(self.num_cameras)) camera_params = np.zeros(shape=(len(camera_id_list), 13), dtype=float) cam_list = [self.cam_list[c] for c in camera_id_list] for i, cid in enumerate(camera_id_list): camera_params[i, 0:3] = np.squeeze(cam_list[cid].rvec) camera_params[i, 3:6] = np.squeeze(cam_list[cid].tvec) camera_params[i, 6] = cam_list[cid].focal_length_x camera_params[i, 7] = cam_list[cid].focal_length_y camera_params[i, 8:13] = np.squeeze(cam_list[cid].distort) point_indices = [] camera_indices = [] points2d_ba = [] points3d_ba = [] points3d_ba_source = dict() points3d_ba_source_inv = dict() point_index_counter = 0 data_shape = self.points3d_m.shape if data_shape[0] > max_num_images: logger.debug( "There are too many ({}) images for calibration. Selecting {} randomly." .format(data_shape[0], max_num_images)) img_id_list = np.random.randint(0, high=data_shape[0] - 1, size=(max_num_images)) else: logger.debug("Using {} images for calibration".format( data_shape[0])) img_id_list = np.arange(data_shape[0] - 1) for img_id in img_id_list: for j_id in range(data_shape[1]): cam_list_iter = list() points2d_iter = list() for cam in cam_list: if j_id in ignore_joint_list: continue if np.any(self.points3d_m[img_id, j_id, :] == 0): continue if np.any(cam[img_id, j_id, :] == 0): continue if not config["skeleton"].camera_see_joint( cam.cam_id, j_id): continue if cam.cam_id == 3: continue cam_list_iter.append(cam) points2d_iter.append(cam[img_id, j_id, :]) # the point is seen by at least two cameras, add it to the bundle adjustment if len(cam_list_iter) >= 2: points3d_iter = self.points3d_m[img_id, j_id, :] points2d_ba.extend(points2d_iter) points3d_ba.append(points3d_iter) point_indices.extend([point_index_counter] * len(cam_list_iter)) points3d_ba_source[(img_id, j_id)] = point_index_counter points3d_ba_source_inv[point_index_counter] = (img_id, j_id) point_index_counter += 1 camera_indices.extend( [cam.cam_id for cam in cam_list_iter]) c = 0 # make sure stripes from both sides share the same point id's # TODO move this into config file if "fly" in config["name"]: for idx, point_idx in enumerate(point_indices): img_id, j_id = points3d_ba_source_inv[point_idx] if (config["skeleton"].is_tracked_point( j_id, config["skeleton"].Tracked.STRIPE) and j_id > config["skeleton"].num_joints // 2): if ( img_id, j_id - config["skeleton"].num_joints // 2, ) in points3d_ba_source: point_indices[idx] = points3d_ba_source[( img_id, j_id - config["skeleton"].num_joints // 2)] c += 1 logger.debug("Replaced {} points".format(c)) points3d_ba = np.squeeze(np.array(points3d_ba)) points2d_ba = np.squeeze(np.array(points2d_ba)) cid2cidx = { v: k for (k, v) in enumerate(np.sort(np.unique(camera_indices))) } camera_indices = [cid2cidx[cid] for cid in camera_indices] camera_indices = np.array(camera_indices) point_indices = np.array(point_indices) n_cameras = camera_params.shape[0] n_points = points3d_ba.shape[0] x0 = np.hstack((camera_params.ravel(), points3d_ba.ravel())) return ( x0.copy(), points2d_ba.copy(), n_cameras, n_points, camera_indices, point_indices, )
def __init__( self, image_folder, output_folder, cam_list=None, calibration=None, num_images=900, cam_id_list=range(config["num_cameras"]), cid2cidread=None, heatmap=None, pred=None, hm_path=None, pred_path=None, ): self.folder = image_folder self.folder_output = output_folder self.dict_name = image_folder self.points3d_m = None self.bone_param = None self.num_images = num_images self.num_joints = config["skeleton"].num_joints self.heatmap_shape = config["heatmap_shape"] self.image_shape = config["image_shape"] self.num_cameras = len(cam_id_list) if cid2cidread is not None: self.cid2cidread = cid2cidread else: _, cid2cidread = read_camera_order(self.folder_output) self.cid2cidread = cid2cidread if cam_list: logger.debug("Camera list is already given, skipping loading.") self.cam_list = cam_list else: self.cam_list = list() if pred_path is None: pred_path = find_pred_path(self.folder_output) if pred_path is None: logger.debug("no pred file under {}".format( self.folder_output)) if pred is None and pred_path is not None: logger.debug("loading pred path {}".format(pred_path)) if pred_path.endswith(".json"): pred = load_pred_from_json(pred_path, os.path.basename(image_folder), self.num_images) else: pred = np.load(file=pred_path, mmap_mode="r", allow_pickle=True) num_images_in_pred = pred.shape[ 1] if pred is not None else num_images if hm_path is None: hm_path = find_hm_path(self.folder_output) if hm_path is None: logger.debug("no heatmap file under {}".format( self.folder_output)) if heatmap is None and hm_path is not None: hm_shape = ( config["num_cameras"] + 1, num_images_in_pred, config["num_predict"], self.heatmap_shape[0], self.heatmap_shape[1], ) logger.debug("Heatmap shape: {}".format(hm_shape)) logger.debug("Reading hm from {}".format(hm_path)) heatmap = load_heatmap(hm_path, hm_shape) if self.num_images is not None and self.num_images < num_images_in_pred: if pred is not None: pred = pred[:, :self.num_images] if heatmap is not None: heatmap = heatmap[:, :self.num_images] for cam_id in cam_id_list: cam_id_read = cid2cidread[cam_id] pred_cam = pred2pred_cam(pred, cam_id, cam_id_read, self.image_shape, num_images_in_pred) self.cam_list.append( Camera( cid=cam_id, cid_read=cam_id_read, image_folder=image_folder, hm=heatmap, points2d=pred_cam, )) if calibration is None: logger.debug("Reading calibration from {}".format( self.folder_output)) calibration = read_calib(self.folder_output) if calibration is not None: _ = self.load_network(calibration)