forked from vt-vl-lab/3d-photo-inpainting
-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathmain.py
176 lines (155 loc) · 8.04 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
import numpy as np
import argparse
import glob
import os
from functools import partial
import vispy
import scipy.misc as misc
from tqdm import tqdm
import yaml
import time
import sys
from mesh import write_ply, read_ply, output_3d_photo
from utils import get_MiDaS_samples, read_MiDaS_depth
import torch
import cv2
import PIL
from skimage.transform import resize
import imageio
import copy
from networks import Inpaint_Color_Net, Inpaint_Depth_Net, Inpaint_Edge_Net
from MiDaS.run import run_depth
from boostmonodepth_utils import run_boostmonodepth
from MiDaS.monodepth_net import MonoDepthNet
import MiDaS.MiDaS_utils as MiDaS_utils
from bilateral_filtering import sparse_bilateral_filtering
dir_path = os.path.dirname(os.path.realpath(__file__))
parser = argparse.ArgumentParser()
parser.add_argument('--config', type=str, default=os.path.join(dir_path, 'argument.yml'), help='Configure of post processing')
parser.add_argument('--size', type=int, default=960, help='size of longest dimension')
parser.add_argument('--num_frames', type=int, default=2)
parser.add_argument('--x_shift', type=float, default=0)
parser.add_argument('--y_shift', type=float, default=0)
parser.add_argument('--z_shift', type=float, default=0)
parser.add_argument('--input', required=True, type=str, help="Input file")
parser.add_argument('--output', type=str, required=True, help="Output file")
args = parser.parse_args()
config = yaml.load(open(os.path.join(dir_path, args.config), 'r'))
if config['offscreen_rendering'] is True:
vispy.use(app='egl')
filename, ext = os.path.basename(args.input).split('.')
config['src_folder'] = os.path.dirname(args.input)
config['specific'] = filename
config['img_format'] = '.%s'%ext
config['mesh_folder'] = os.path.join(dir_path, config['mesh_folder'])
config['video_folder'] = os.path.join(dir_path, config['video_folder'])
config['depth_folder'] = os.path.join(dir_path, config['depth_folder'])
config['depth_edge_model_ckpt'] = os.path.join(dir_path, config['depth_edge_model_ckpt'])
config['depth_feat_model_ckpt'] = os.path.join(dir_path, config['depth_feat_model_ckpt'])
config['rgb_feat_model_ckpt'] = os.path.join(dir_path, config['rgb_feat_model_ckpt'])
config['traj_types'] = ['double-straight-line']
config['video_postfix'] = ['zoom-in']
config['longer_side_len'] = args.size
config['num_frames'] = 2
config['save_ply'] = False
config['x_shift_range'] = [args.x_shift]
config['y_shift_range'] = [args.y_shift]
config['z_shift_range'] = [args.z_shift]
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['use_boostmonodepth'] is True:
run_boostmonodepth(sample['ref_img_fi'], config['src_folder'], config['depth_folder'])
elif 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, results = 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, save_video=False)
result = results[0]
PIL.Image.fromarray(result).save(args.output)