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
if not (config['load_ply'] is True and os.path.exists(mesh_fi)):
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'])
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 ...")
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()
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'],
def load_edge_model(device: str, depth_edge_model_ckpt: str) -> Inpaint_Edge_Net:
print(f(_("Loading edge model at {datetime.now():%Y-%m-%d %H:%M:%S.%f}")))
depth_edge_model = Inpaint_Edge_Net(init_weights=True)
depth_edge_weight = torch.load(depth_edge_model_ckpt, map_location=torch.device(device))
depth_edge_model.load_state_dict(depth_edge_weight)
return depth_edge_model.to(device)
if not (config['load_ply'] is True and os.path.exists(mesh_fi)):
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'])
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 ...")
depth_edge_model = Inpaint_Edge_Net(init_weights=True)
depth_edge_weight = torch.load(config['depth_edge_model_ckpt'])
depth_edge_model.load_state_dict(depth_edge_weight)
depth_edge_model = depth_edge_model.cuda(config['gpu_ids'])
depth_edge_model.eval()
depth_feat_model = Inpaint_Depth_Net()
depth_feat_weight = torch.load(config['depth_feat_model_ckpt'])
depth_feat_model.load_state_dict(depth_feat_weight, strict=True)
depth_feat_model = depth_feat_model.cuda(config['gpu_ids'])
depth_feat_model.eval()
depth_feat_model = depth_feat_model.cuda(config['gpu_ids'])
rgb_model = Inpaint_Color_Net()
rgb_feat_weight = torch.load(config['rgb_feat_model_ckpt'])
rgb_model.load_state_dict(rgb_feat_weight)
rgb_model.eval()
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
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) # do bilateral filtering
depth = vis_depths[-1]
model = None
torch.cuda.empty_cache()
## MODEL INITS
print("Start Running 3D_Photo ...")
print(f"Loading edge model at {time.time()}")
depth_edge_model = Inpaint_Edge_Net(
init_weights=True) # init edge inpainting model
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() # in eval mode
print(f"Loading depth model at {time.time()}")
depth_feat_model = Inpaint_Depth_Net() # init depth inpainting model
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)
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")