def load_rgb_model(device: str, rgb_feat_model_ckpt: str) -> Inpaint_Color_Net: print(f(_("Loading rgb model at {datetime.now():%Y-%m-%d %H:%M:%S.%f}"))) rgb_model = Inpaint_Color_Net() rgb_feat_weight = torch.load(rgb_feat_model_ckpt, map_location=torch.device(device)) rgb_model.load_state_dict(rgb_feat_weight) rgb_model.eval() return rgb_model.to(device)
map_location=torch.device(device)) depth_edge_model.load_state_dict(depth_edge_weight) depth_edge_model = depth_edge_model.to(device) depth_edge_model.eval() depth_feat_model = Inpaint_Depth_Net() depth_feat_weight = torch.load(config['depth_feat_model_ckpt'], map_location=torch.device(device)) depth_feat_model.load_state_dict(depth_feat_weight, strict=True) depth_feat_model = depth_feat_model.to(device) depth_feat_model.eval() depth_feat_model = depth_feat_model.to(device) rgb_model = Inpaint_Color_Net() rgb_feat_weight = torch.load(config['rgb_feat_model_ckpt'], map_location=torch.device(device)) rgb_model.load_state_dict(rgb_feat_weight) rgb_model.eval() rgb_model = rgb_model.to(device) graph = None rt_info = write_ply(image, depth, sample['int_mtx'], mesh_fi, config, rgb_model, depth_edge_model, depth_edge_model, depth_feat_model) if rt_info is False: continue rgb_model = None color_feat_model = None depth_edge_model = None depth_feat_model = None torch.cuda.empty_cache() if config['save_ply'] is True or config['load_ply'] is True: verts, colors, faces, Height, Width, hFov, vFov = read_ply(mesh_fi)
def make_video(params_file): print("Making Video...") with open(params_file) as f: params = json.load(f) os.makedirs(params['mesh_folder'], exist_ok=True) os.makedirs(params['video_folder'], exist_ok=True) os.makedirs(params['depth_folder'], exist_ok=True) sample_list = get_MiDaS_samples(params, params['specific']) normal_canvas, all_canvas = None, None device = "cuda" depth = None try: sample = sample_list[0] except: sample = sample_list # print("Current Source ==> ", sample['src_pair_name']) mesh_fi = os.path.join(params['mesh_folder'], sample['src_pair_name'] + '.ply') image = imageio.imread(sample['ref_img_fi']) # print(f"Running depth extraction at {time.time()}") if params['require_midas'] is True: run_depth([sample['ref_img_fi']], params['src_dir'], params['depth_folder'], params['MiDaS_model_ckpt'], MonoDepthNet, MiDaS_utils, target_w=640) if 'npy' in params['depth_format']: params['output_h'], params['output_w'] = np.load( sample['depth_fi']).shape[:2] else: params['output_h'], params['output_w'] = imageio.imread( sample['depth_fi']).shape[:2] frac = params['longer_side_len'] / max(params['output_h'], params['output_w']) params['output_h'], params['output_w'] = int( params['output_h'] * frac), int(params['output_w'] * frac) params['original_h'], params['original_w'] = params['output_h'], params[ 'output_w'] if image.ndim == 2: image = image[..., None].repeat(3, -1) if np.sum(np.abs(image[..., 0] - image[..., 1])) == 0 and np.sum( np.abs(image[..., 1] - image[..., 2])) == 0: params['gray_image'] = True else: params['gray_image'] = False image = cv2.resize(image, (params['output_w'], params['output_h']), interpolation=cv2.INTER_AREA) depth = read_MiDaS_depth(sample['depth_fi'], 3.0, params['output_h'], params['output_w']) mean_loc_depth = depth[depth.shape[0] // 2, depth.shape[1] // 2] if not (params['load_ply'] is True and os.path.exists(mesh_fi)): vis_photos, vis_depths = sparse_bilateral_filtering( depth.copy(), image.copy(), params, num_iter=params['sparse_iter'], spdb=False) depth = vis_depths[-1] model = None torch.cuda.empty_cache() depth_edge_model = Inpaint_Edge_Net(init_weights=True) depth_edge_weight = torch.load(params['depth_edge_model_ckpt'], map_location=torch.device(device)) depth_edge_model.load_state_dict(depth_edge_weight) depth_edge_model = depth_edge_model.to(device) depth_edge_model.eval() depth_feat_model = Inpaint_Depth_Net() depth_feat_weight = torch.load(params['depth_feat_model_ckpt'], map_location=torch.device(device)) depth_feat_model.load_state_dict(depth_feat_weight, strict=True) depth_feat_model = depth_feat_model.to(device) depth_feat_model.eval() depth_feat_model = depth_feat_model.to(device) rgb_model = Inpaint_Color_Net() rgb_feat_weight = torch.load(params['rgb_feat_model_ckpt'], map_location=torch.device(device)) rgb_model.load_state_dict(rgb_feat_weight) rgb_model.eval() rgb_model = rgb_model.to(device) graph = None rt_info = write_ply("", image, depth, sample['int_mtx'], mesh_fi, params, rgb_model, depth_edge_model, depth_edge_model, depth_feat_model) if rt_info is False: return rgb_model = None color_feat_model = None depth_edge_model = None depth_feat_model = None torch.cuda.empty_cache() if params['save_ply'] is True or params['load_ply'] is True: verts, colors, faces, Height, Width, hFov, vFov = read_ply(mesh_fi) else: verts, colors, faces, Height, Width, hFov, vFov = rt_info videos_poses, video_basename = copy.deepcopy( sample['tgts_poses']), sample['tgt_name'] top = (params.get('original_h') // 2 - sample['int_mtx'][1, 2] * params['output_h']) left = (params.get('original_w') // 2 - sample['int_mtx'][0, 2] * params['output_w']) down, right = top + params['output_h'], left + params['output_w'] border = [int(xx) for xx in [top, down, left, right]] normal_canvas, all_canvas = output_3d_photo( verts.copy(), colors.copy(), faces.copy(), copy.deepcopy(Height), copy.deepcopy(Width), copy.deepcopy(hFov), copy.deepcopy(vFov), copy.deepcopy(sample['tgt_pose']), sample['video_postfix'], copy.deepcopy(sample['ref_pose']), copy.deepcopy(params['video_folder']), image.copy(), copy.deepcopy(sample['int_mtx']), params, image, videos_poses, video_basename, params.get('original_h'), params.get('original_w'), border=border, depth=depth, normal_canvas=normal_canvas, all_canvas=all_canvas, mean_loc_depth=mean_loc_depth) print("Done!")
def Main(config_dict): config = config_dict print(config) if config['offscreen_rendering'] is True: vispy.use(app='egl') os.makedirs(config['mesh_folder'], exist_ok=True) os.makedirs(config['video_folder'], exist_ok=True) os.makedirs(config['depth_folder'], exist_ok=True) sample_list = get_MiDaS_samples(config['src_folder'], config['depth_folder'], config, config['specific']) normal_canvas, all_canvas = None, None if isinstance(config["gpu_ids"], int) and (config["gpu_ids"] >= 0): device = config["gpu_ids"] else: device = "cpu" print(f"running on device {device}") for idx in tqdm(range(len(sample_list))): depth = None sample = sample_list[idx] print("Current Source ==> ", sample['src_pair_name']) mesh_fi = os.path.join(config['mesh_folder'], sample['src_pair_name'] +'.ply') image = imageio.imread(sample['ref_img_fi']) print(f"Running depth extraction at {time.time()}") if config['require_midas'] is True: run_depth([sample['ref_img_fi']], config['src_folder'], config['depth_folder'], config['MiDaS_model_ckpt'], MonoDepthNet, MiDaS_utils, target_w=640) if 'npy' in config['depth_format']: config['output_h'], config['output_w'] = np.load(sample['depth_fi']).shape[:2] else: config['output_h'], config['output_w'] = imageio.imread(sample['depth_fi']).shape[:2] frac = config['longer_side_len'] / max(config['output_h'], config['output_w']) config['output_h'], config['output_w'] = int(config['output_h'] * frac), int(config['output_w'] * frac) config['original_h'], config['original_w'] = config['output_h'], config['output_w'] if image.ndim == 2: image = image[..., None].repeat(3, -1) if np.sum(np.abs(image[..., 0] - image[..., 1])) == 0 and np.sum(np.abs(image[..., 1] - image[..., 2])) == 0: config['gray_image'] = True else: config['gray_image'] = False image = cv2.resize(image, (config['output_w'], config['output_h']), interpolation=cv2.INTER_AREA) depth = read_MiDaS_depth(sample['depth_fi'], 3.0, config['output_h'], config['output_w']) mean_loc_depth = depth[depth.shape[0]//2, depth.shape[1]//2] if not(config['load_ply'] is True and os.path.exists(mesh_fi)): vis_photos, vis_depths = sparse_bilateral_filtering(depth.copy(), image.copy(), config, num_iter=config['sparse_iter'], spdb=False) depth = vis_depths[-1] model = None torch.cuda.empty_cache() print("Start Running 3D_Photo ...") print(f"Loading edge model at {time.time()}") depth_edge_model = Inpaint_Edge_Net(init_weights=True) depth_edge_weight = torch.load(config['depth_edge_model_ckpt'], map_location=torch.device(device)) depth_edge_model.load_state_dict(depth_edge_weight) depth_edge_model = depth_edge_model.to(device) depth_edge_model.eval() print(f"Loading depth model at {time.time()}") depth_feat_model = Inpaint_Depth_Net() depth_feat_weight = torch.load(config['depth_feat_model_ckpt'], map_location=torch.device(device)) depth_feat_model.load_state_dict(depth_feat_weight, strict=True) depth_feat_model = depth_feat_model.to(device) depth_feat_model.eval() depth_feat_model = depth_feat_model.to(device) print(f"Loading rgb model at {time.time()}") rgb_model = Inpaint_Color_Net() rgb_feat_weight = torch.load(config['rgb_feat_model_ckpt'], map_location=torch.device(device)) rgb_model.load_state_dict(rgb_feat_weight) rgb_model.eval() rgb_model = rgb_model.to(device) graph = None print(f"Writing depth ply (and basically doing everything) at {time.time()}") rt_info = write_ply(image, depth, sample['int_mtx'], mesh_fi, config, rgb_model, depth_edge_model, depth_edge_model, depth_feat_model) if rt_info is False: continue rgb_model = None color_feat_model = None depth_edge_model = None depth_feat_model = None torch.cuda.empty_cache() if config['save_ply'] is True or config['load_ply'] is True: verts, colors, faces, Height, Width, hFov, vFov = read_ply(mesh_fi) else: verts, colors, faces, Height, Width, hFov, vFov = rt_info print(f"Making video at {time.time()}") videos_poses, video_basename = copy.deepcopy(sample['tgts_poses']), sample['tgt_name'] top = (config.get('original_h') // 2 - sample['int_mtx'][1, 2] * config['output_h']) left = (config.get('original_w') // 2 - sample['int_mtx'][0, 2] * config['output_w']) down, right = top + config['output_h'], left + config['output_w'] border = [int(xx) for xx in [top, down, left, right]] normal_canvas, all_canvas = output_3d_photo(verts.copy(), colors.copy(), faces.copy(), copy.deepcopy(Height), copy.deepcopy(Width), copy.deepcopy(hFov), copy.deepcopy(vFov), copy.deepcopy(sample['tgt_pose']), sample['video_postfix'], copy.deepcopy(sample['ref_pose']), copy.deepcopy(config['video_folder']), image.copy(), copy.deepcopy(sample['int_mtx']), config, image, videos_poses, video_basename, config.get('original_h'), config.get('original_w'), border=border, depth=depth, normal_canvas=normal_canvas, all_canvas=all_canvas, mean_loc_depth=mean_loc_depth) move_to_dir("video", "static") return 0
class Image3D: def __init__(self, config) -> None: self.config = yaml.load(open(config, 'r')) if os.name == 'nt': vispy.use(app='pyqt5') else: vispy.use(app='egl') if isinstance(self.config['gpu_ids'], int) and (self.config['gpu_ids'] >= 0): self.device = self.config['gpu_ids'] else: self.device = "cpu" print(f"Loading edge model at {time.time()}") self.depth_edge_model = Inpaint_Edge_Net(init_weights=True) self.depth_edge_weight = torch.load(self.config['depth_edge_model_ckpt'], map_location=torch.device(self.device)) self.depth_edge_model.load_state_dict(self.depth_edge_weight) self.depth_edge_model = self.depth_edge_model.to(self.device) self.depth_edge_model.eval() print(f"Loading depth model at {time.time()}") self.depth_feat_model = Inpaint_Depth_Net() self.depth_feat_weight = torch.load(self.config['depth_feat_model_ckpt'], map_location=torch.device(self.device)) self.depth_feat_model.load_state_dict(self.depth_feat_weight, strict=True) self.depth_feat_model = self.depth_feat_model.to(self.device) self.depth_feat_model.eval() self.depth_feat_model = self.depth_feat_model.to(self.device) print(f"Loading rgb model at {time.time()}") self.rgb_model = Inpaint_Color_Net() self.rgb_feat_weight = torch.load(self.config['rgb_feat_model_ckpt'], map_location=torch.device(self.device)) self.rgb_model.load_state_dict(self.rgb_feat_weight) self.rgb_model.eval() self.rgb_model = self.rgb_model.to(self.device) self.graph = None def run_3dimage(self, src) -> str: src_folder = src + '/image' self.config['mesh_folder'] = src self.config['video_folder'] = src self.config['depth_folder'] = src + '/depth' sample_list = get_MiDaS_samples(src_folder, self.config['depth_folder'], self.config, self.config['specific']) normal_canvas, all_canvas = None, None for idx in range(len(sample_list)): depth = None sample = sample_list[idx] print("Current Source ==> ", sample['src_pair_name']) mesh_fi = os.path.join(self.config['mesh_folder'], sample['src_pair_name'] +'.ply') image = imageio.imread(sample['ref_img_fi']) self.config['output_h'], self.config['output_w'] = imageio.imread(sample['depth_fi'], as_gray=True).shape[:2] frac = self.config['longer_side_len'] / max(self.config['output_h'], self.config['output_w']) self.config['output_h'], self.config['output_w'] = int(self.config['output_h'] * frac), int(self.config['output_w'] * frac) self.config['original_h'], self.config['original_w'] = self.config['output_h'], self.config['output_w'] if image.ndim == 2: image = image[..., None].repeat(3, -1) if np.sum(np.abs(image[..., 0] - image[..., 1])) == 0 and np.sum(np.abs(image[..., 1] - image[..., 2])) == 0: self.config['gray_image'] = True else: self.config['gray_image'] = False image = cv2.resize(image, (self.config['output_w'], self.config['output_h']), interpolation=cv2.INTER_AREA) depth = read_MiDaS_depth(sample['depth_fi'], 3.0, self.config['output_h'], self.config['output_w']) mean_loc_depth = depth[depth.shape[0]//2, depth.shape[1]//2] if not(self.config['load_ply'] is True and os.path.exists(mesh_fi)): _, vis_depths = sparse_bilateral_filtering(depth.copy(), image.copy(), self.config, num_iter=self.config['sparse_iter'], spdb=False) depth = vis_depths[-1] torch.cuda.empty_cache() print("Start Running 3D_Photo ...") print(f"Writing depth ply (and basically doing everything) at {time.time()}") rt_info = write_ply(image, depth, sample['int_mtx'], mesh_fi, self.config, self.rgb_model, self.depth_edge_model, self.depth_edge_model, self.depth_feat_model) if rt_info is False: continue torch.cuda.empty_cache() if self.config['save_ply'] is True or self.config['load_ply'] is True: verts, colors, faces, Height, Width, hFov, vFov = read_ply(mesh_fi) else: verts, colors, faces, Height, Width, hFov, vFov = rt_info print(f"Making video at {time.time()}") videos_poses, video_basename = copy.deepcopy(sample['tgts_poses']), sample['tgt_name'] top = (self.config.get('original_h') // 2 - sample['int_mtx'][1, 2] * self.config['output_h']) left = (self.config.get('original_w') // 2 - sample['int_mtx'][0, 2] * self.config['output_w']) down, right = top + self.config['output_h'], left + self.config['output_w'] border = [int(xx) for xx in [top, down, left, right]] normal_canvas, all_canvas = output_3d_photo(verts.copy(), colors.copy(), faces.copy(), copy.deepcopy(Height), copy.deepcopy(Width), copy.deepcopy(hFov), copy.deepcopy(vFov), copy.deepcopy(sample['tgt_pose']), sample['video_postfix'], copy.deepcopy(sample['ref_pose']), copy.deepcopy(self.config['video_folder']), image.copy(), copy.deepcopy(sample['int_mtx']), self.config, image, videos_poses, video_basename, self.config.get('original_h'), self.config.get('original_w'), border=border, depth=depth, normal_canvas=normal_canvas, all_canvas=all_canvas, mean_loc_depth=mean_loc_depth) torch.cuda.empty_cache()
def render_mp4(input_dir, output_dir, filer_type): print("invoked render start") parser = argparse.ArgumentParser() parser.add_argument("--config", type=str, default="argument.yml", help="Configure of post processing") args = parser.parse_args() config = yaml.load(open(args.config, "r")) if config["offscreen_rendering"] is True: vispy.use(app="egl") os.makedirs(config["mesh_folder"], exist_ok=True) os.makedirs(config["video_folder"], exist_ok=True) os.makedirs(config["depth_folder"], exist_ok=True) sample_list = get_MiDaS_samples(input_dir, config["depth_folder"], config, filer_type, config["specific"]) normal_canvas, all_canvas = None, None if isinstance(config["gpu_ids"], int) and (config["gpu_ids"] >= 0): device = config["gpu_ids"] else: device = "cpu" print(f"running on device {device}") for idx in tqdm(range(len(sample_list))): depth = None sample = sample_list[idx] print("Current Source ==> ", sample["src_pair_name"]) mesh_fi = os.path.join(config["mesh_folder"], sample["src_pair_name"] + ".ply") image = imageio.imread(sample["ref_img_fi"]) print(f"Running depth extraction at {time.time()}") if config["require_midas"] is True: run_depth( [sample["ref_img_fi"]], config["src_folder"], config["depth_folder"], config["MiDaS_model_ckpt"], MonoDepthNet, MiDaS_utils, target_w=640, ) if "npy" in config["depth_format"]: config["output_h"], config["output_w"] = np.load( sample["depth_fi"]).shape[:2] else: config["output_h"], config["output_w"] = imageio.imread( sample["depth_fi"]).shape[:2] frac = config["longer_side_len"] / max(config["output_h"], config["output_w"]) config["output_h"], config["output_w"] = int( config["output_h"] * frac), int(config["output_w"] * frac) config["original_h"], config["original_w"] = config[ "output_h"], config["output_w"] if image.ndim == 2: image = image[..., None].repeat(3, -1) if np.sum(np.abs(image[..., 0] - image[..., 1])) == 0 and np.sum( np.abs(image[..., 1] - image[..., 2])) == 0: config["gray_image"] = True else: config["gray_image"] = False image = cv2.resize(image, (config["output_w"], config["output_h"]), interpolation=cv2.INTER_AREA) depth = read_MiDaS_depth(sample["depth_fi"], 3.0, config["output_h"], config["output_w"]) mean_loc_depth = depth[depth.shape[0] // 2, depth.shape[1] // 2] if not (config["load_ply"] is True and os.path.exists(mesh_fi)): vis_photos, vis_depths = sparse_bilateral_filtering( depth.copy(), image.copy(), config, num_iter=config["sparse_iter"], spdb=False) depth = vis_depths[-1] model = None torch.cuda.empty_cache() print("Start Running 3D_Photo ...") print(f"Loading edge model at {time.time()}") depth_edge_model = Inpaint_Edge_Net(init_weights=True) depth_edge_weight = torch.load(config["depth_edge_model_ckpt"], map_location=torch.device(device)) depth_edge_model.load_state_dict(depth_edge_weight) depth_edge_model = depth_edge_model.to(device) depth_edge_model.eval() print(f"Loading depth model at {time.time()}") depth_feat_model = Inpaint_Depth_Net() depth_feat_weight = torch.load(config["depth_feat_model_ckpt"], map_location=torch.device(device)) depth_feat_model.load_state_dict(depth_feat_weight, strict=True) depth_feat_model = depth_feat_model.to(device) depth_feat_model.eval() depth_feat_model = depth_feat_model.to(device) print(f"Loading rgb model at {time.time()}") rgb_model = Inpaint_Color_Net() rgb_feat_weight = torch.load(config["rgb_feat_model_ckpt"], map_location=torch.device(device)) rgb_model.load_state_dict(rgb_feat_weight) rgb_model.eval() rgb_model = rgb_model.to(device) graph = None print( f"Writing depth ply (and basically doing everything) at {time.time()}" ) rt_info = write_ply( image, depth, sample["int_mtx"], mesh_fi, config, rgb_model, depth_edge_model, depth_edge_model, depth_feat_model, ) if rt_info is False: continue rgb_model = None color_feat_model = None depth_edge_model = None depth_feat_model = None torch.cuda.empty_cache() if config["save_ply"] is True or config["load_ply"] is True: verts, colors, faces, Height, Width, hFov, vFov = read_ply(mesh_fi) else: verts, colors, faces, Height, Width, hFov, vFov = rt_info print(f"Making video at {time.time()}") videos_poses, video_basename = copy.deepcopy( sample["tgts_poses"]), sample["tgt_name"] top = config.get("original_h") // 2 - sample["int_mtx"][ 1, 2] * config["output_h"] left = config.get("original_w") // 2 - sample["int_mtx"][ 0, 2] * config["output_w"] down, right = top + config["output_h"], left + config["output_w"] border = [int(xx) for xx in [top, down, left, right]] normal_canvas, all_canvas = output_3d_photo( verts.copy(), colors.copy(), faces.copy(), copy.deepcopy(Height), copy.deepcopy(Width), copy.deepcopy(hFov), copy.deepcopy(vFov), copy.deepcopy(sample["tgt_pose"]), sample["video_postfix"], copy.deepcopy(sample["ref_pose"]), copy.deepcopy(output_dir), image.copy(), copy.deepcopy(sample["int_mtx"]), config, image, videos_poses, video_basename, config.get("original_h"), config.get("original_w"), border=border, depth=depth, normal_canvas=normal_canvas, all_canvas=all_canvas, mean_loc_depth=mean_loc_depth, ) print("invoked render end")