-
Notifications
You must be signed in to change notification settings - Fork 17
Expand file tree
/
Copy pathrender.py
More file actions
313 lines (248 loc) · 11.8 KB
/
render.py
File metadata and controls
313 lines (248 loc) · 11.8 KB
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
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
import time
import os
import logging
import shutil
import glob
from pathlib import Path
from typing import Optional
import hydra
from omegaconf import DictConfig
from pixie.utils import (run_cmd, resolve_paths,
get_output_paths, create_directories, validate_config,
get_latest_nerf_run, save_contextual_config,
generate_material_segmentation,
generate_neural_segmentation,
get_material_vlm_segmentation_path,
get_material_neural_segmentation_path,
set_logger)
from pixie.viz_utils import compile_video
import cv2
def get_video_fps(video_path: str) -> float:
"""Extracts frames-per-second (FPS) from a video file."""
if not os.path.exists(video_path):
raise FileNotFoundError(f"Video file not found at {video_path}")
video = cv2.VideoCapture(str(video_path))
fps = video.get(cv2.CAP_PROP_FPS)
video.release()
return fps
def move_files_to_folder(src_dir: Path, dst_dir: Path) -> None:
"""Move all files from src_dir to dst_dir."""
dst_dir.mkdir(parents=True, exist_ok=True)
for file_path in src_dir.iterdir():
if file_path.is_file():
shutil.move(str(file_path), str(dst_dir / file_path.name))
# Remove empty source directory
if src_dir.exists() and not any(src_dir.iterdir()):
src_dir.rmdir()
def render_nerf_model(cfg: DictConfig, paths: dict) -> str:
"""Render RGB images from the trained NeRF model."""
# Get the latest NeRF training run
nerf_output_dir = paths['nerf_output']
latest_run = get_latest_nerf_run(nerf_output_dir)
if not latest_run:
raise RuntimeError(
f"No NeRF training run found in {nerf_output_dir}. Please run the pipeline first."
)
config_path = os.path.join(latest_run, "config.yml")
if not os.path.exists(config_path):
raise RuntimeError(f"NeRF config file not found at {config_path}")
# Create render output directory
render_output_dir = f"{paths['render_output']}/rgb_renders"
os.makedirs(render_output_dir, exist_ok=True)
# Check if we need to render (check for existing rendered images in subdirectories)
rendered_images_exist = any(glob.glob(f"{render_output_dir}/*/*/*.{ext}")
for ext in ['png', 'jpg', 'jpeg'])
# Render RGB images if needed
if cfg.overwrite or not rendered_images_exist:
render_cmd = [
"ns-render",
"dataset",
"--load-config",
config_path,
"--output-path",
render_output_dir,
"--rendered-output-names",
cfg.output_rendering.nerf_render.rendered_output_names,
"--split",
cfg.output_rendering.nerf_render.split,
]
run_cmd(render_cmd, step_name="RENDER_RGB")
logging.info(
f"RGB rendering completed. Output directory: {render_output_dir}")
else:
logging.info(
f"RGB renders already exist at {render_output_dir}. Skipping rendering."
)
return render_output_dir
def render_blender_gs(cfg: DictConfig, paths: dict) -> str:
"""Render Gaussian Splatting outputs using Blender."""
blender_output_dir = paths['blender_output']
render_cfg = cfg.output_rendering.blender_render_gs
# Get the ply_dir from physgaussian_output
ply_dir = os.path.join(paths['physgaussian_output'], f"sample_{cfg.physics.sample_id}", "ply_files")
assert os.path.exists(ply_dir), f"PLY directory does not exist: {ply_dir}"
# Create output directory for GS renders
gs_output_dir = os.path.join(blender_output_dir, "gs_renders")
os.makedirs(gs_output_dir, exist_ok=True)
# Check if video already exists
video_output = os.path.join(gs_output_dir, f"output.mp4")
if os.path.exists(video_output) and not cfg.overwrite:
logging.info(f"GS render video already exists at {video_output}. Skipping rendering.")
return gs_output_dir
# Build command to run render_gs.py
cmd = [
"blender", "-b", "-P", "pixie/blender/render_gs.py", "--",
"--obj_id", cfg.obj_id,
"--ply_dir", ply_dir,
"--output_dir", gs_output_dir,
"--data_dir", paths['data_dir'],
"--blend_file_path", paths['blend_file_path'],
"--blender_gs_addon_path", paths['blender_gs_addon_path']
]
# Add place_on_ground flag if enabled
if render_cfg.get('place_on_ground', False):
cmd.append("--place_on_ground")
# Add optional arguments based on render config
if render_cfg.get('rotate_around') is not None:
cmd.extend(["--rotate_around", str(render_cfg.rotate_around)])
if render_cfg.get('camera_id') is not None:
cmd.extend(["--camera_id", str(render_cfg.camera_id)])
if render_cfg.get('focal_length_ratio') is not None:
cmd.extend(["--focal_length_ratio", str(render_cfg.focal_length_ratio)])
if render_cfg.get('transparent', False):
cmd.append("--transparent")
if render_cfg.get('resolution_x') is not None:
cmd.extend(["--resolution_x", str(render_cfg.resolution_x)])
if render_cfg.get('resolution_y') is not None:
cmd.extend(["--resolution_y", str(render_cfg.resolution_y)])
if render_cfg.get('cycles_samples') is not None:
cmd.extend(["--cycles_samples", str(render_cfg.cycles_samples)])
if render_cfg.get('num_renders') is not None:
cmd.extend(["--num_renders", str(render_cfg.num_renders)])
if render_cfg.get('start_frame', 0) != 0:
cmd.extend(["--start_frame", str(render_cfg.start_frame)])
if render_cfg.get('is_dropping', False):
cmd.append("--is_dropping")
if render_cfg.get('save_blend', False):
cmd.append("--save_blend")
if render_cfg.get('bg_color') is not None:
cmd.extend(["--bg_color"] + [str(c) for c in render_cfg.bg_color])
# Add init_xyz and xyz_rotation
init_xyz = render_cfg.get('init_xyz', [0.0, 0.0, 0.0])
cmd.extend(["--init_xyz"] + [str(x) for x in init_xyz])
xyz_rotation = render_cfg.get('xyz_rotation', [0.0, 0.0, 0.0])
cmd.extend(["--xyz_rotation"] + [str(r) for r in xyz_rotation])
if cfg.overwrite:
cmd.append("--overwrite")
# Compile frames to video
frames_dir = Path(gs_output_dir)
# Run the rendering command
run_cmd(cmd, step_name="RENDER_BLENDER_GS")
save_contextual_config(cfg, blender_output_dir, context="blender_gs_render")
assert frames_dir.exists(), f"Frames directory {frames_dir} does not exist"
assert any(frames_dir.glob('*.png')), f"No frames found in {frames_dir}"
logging.info(f"Compiling frames to video: {video_output}")
# Get FPS from physics video or use default
physics_video_path = os.path.join(paths["physgaussian_output"],
f"sample_{cfg.physics.sample_id}",
"frames", "output.mp4")
fps = render_cfg.default_fps
if os.path.exists(physics_video_path):
fps = get_video_fps(physics_video_path)
logging.info(f"FPS: {fps}")
compile_video(frames_dir, video_output, fps)
logging.info(f"Gaussian Splatting rendering completed. Output directory: {gs_output_dir}")
return gs_output_dir
def render_blender_glb(cfg: DictConfig, paths: dict) -> str:
blender_output_dir = paths['blender_output']
render_cfg = cfg.output_rendering.blender_render_glb
## get the point_cloud_path which contains the material physics predictions
if cfg.material_mode == 'vlm':
point_cloud_path = get_material_vlm_segmentation_path(
cfg, paths['render_output'], paths)
elif cfg.material_mode == 'neural':
point_cloud_path = get_material_neural_segmentation_path(
cfg, paths['render_output'], paths)
else:
raise ValueError(f"Invalid material mode: {cfg.material_mode}")
for feature, cmap in zip(render_cfg.features, render_cfg.cmaps):
cmd = (
f"blender -b -P pixie/blender/apply_feature_colors.py -- "
f"--obj_ids {cfg.obj_id} "
f"--feature {feature} "
f"--pred_ply {point_cloud_path} "
f"--output_dir {blender_output_dir} "
f"--colormap {cmap} "
f"--data_dir {paths['data_dir']} "
f"--blend_file_path {paths['blend_file_path']} "
f"--overwrite "
)
if render_cfg.rotate_video:
cmd += f"--rotate_video "
if render_cfg.views:
cmd += f"--views {render_cfg.views} "
if render_cfg.focal_length:
cmd += f"--focal_length {render_cfg.focal_length} "
if render_cfg.camera_id:
cmd += f"--camera_id {render_cfg.camera_id} "
frames_dir = Path(blender_output_dir) / feature
video_output = frames_dir / f"{feature}.mp4"
if os.path.exists(video_output):
continue
run_cmd(cmd, step_name="RENDER_BLENDER")
save_contextual_config(cfg, blender_output_dir, context="blender_output_render")
# # Compile frames to video if rotate_video is enabled
if render_cfg.rotate_video:
# Frames are saved in blender_output_dir/feature subdirectories
assert frames_dir.exists(), f"Frames directory {frames_dir} does not exist"
assert any(frames_dir.glob('*.png')), f"No frames found in {frames_dir}"
logging.info(f"Compiling frames to video: {video_output}")
video_path = os.path.join(paths["physgaussian_output"],
f"sample_{cfg.physics.sample_id}",
"frames", "output.mp4")
fps = cfg.output_rendering.blender_render_glb.default_fps
if os.path.exists(video_path):
fps = get_video_fps(video_path)
logging.info(f"FPS: {fps}")
compile_video(frames_dir, video_output, fps)
return blender_output_dir
def render_gs_model(cfg: DictConfig, paths: dict) -> str:
render_gs_cmd = [
"python", "third_party/PhysGaussian/gaussian-splatting/render.py",
"--model_path", paths['gs_output'],
]
output_dir = os.path.join(paths['gs_output'], "test")
if not os.path.exists(output_dir):
run_cmd(render_gs_cmd, step_name="RENDER_GS")
return output_dir
@hydra.main(version_base=None, config_path="config", config_name="config")
def main(cfg: DictConfig) -> None:
"""Main function to render images from the trained NeRF model."""
# Set up logging first
set_logger()
# Validate configuration
validate_config(cfg)
# Resolve paths and system-specific settings
cfg = resolve_paths(cfg)
# Get all output paths
paths = get_output_paths(cfg, cfg.obj_id)
# Create necessary directories
create_directories(paths)
start_time = time.time()
if 'nerf' in cfg.output_rendering.render:
# Render RGB images from test set
output_dir = render_nerf_model(cfg, paths)
logging.info(f"[NERF] Output directory: {output_dir}")
if 'gs' in cfg.output_rendering.render:
output_dir = render_gs_model(cfg, paths)
logging.info(f"[GS] Output directory: {output_dir}")
if 'blender_glb' in cfg.output_rendering.render:
output_dir = render_blender_glb(cfg, paths)
logging.info(f"[BLENDER_GLB] Output directory: {output_dir}")
if 'blender_gs' in cfg.output_rendering.render:
output_dir = render_blender_gs(cfg, paths)
logging.info(f"[BLENDER_GS] Output directory: {output_dir}")
end_time = time.time()
logging.info(f"Rendering completed in {end_time - start_time:.2f} seconds")
if __name__ == "__main__":
main()