Spaces:
Paused
Paused
| import numpy as np | |
| import os | |
| import gradio as gr | |
| import numpy as np | |
| import tempfile | |
| import shutil | |
| import hashlib | |
| import subprocess | |
| import time | |
| TMP_DIR = os.path.join(tempfile.gettempdir(), 'gradio') | |
| os.makedirs(TMP_DIR, exist_ok=True) | |
| gs_train_iter = 1000 | |
| pose_lr = '1x' | |
| def move_uploaded_files(files, temp_dir): | |
| saved_files = [] | |
| for file in files: | |
| target_file_path = os.path.join(temp_dir, os.path.basename(file.name)) | |
| saved_files.append(target_file_path) | |
| if os.path.exists(file.name): | |
| shutil.move(file.name, target_file_path) | |
| return saved_files | |
| def process_images(files): | |
| n_view = len(files) | |
| files_str = ','.join([f.name for f in files]) | |
| hashdir_name = hashlib.md5(files_str.encode('utf8')).hexdigest() | |
| hashdir_path = os.path.join(TMP_DIR, hashdir_name) | |
| hashdir_path_img = os.path.join(hashdir_path, 'images') | |
| os.makedirs(hashdir_path_img, exist_ok=True) | |
| saved_files = move_uploaded_files(files, hashdir_path_img) | |
| saved_files_str = ' '.join(saved_files) | |
| stdout, stderr, = [], [] | |
| # ----- (1) Dust3r_coarse_geometric_initialization ----- | |
| command = [ | |
| "python", "-W", "ignore", | |
| "./coarse_init_infer.py", | |
| "--img_base_path", hashdir_path, | |
| "--n_views", str(n_view), | |
| "--focal_avg", | |
| ] | |
| start_time = time.time() | |
| result = subprocess.run(command, capture_output=True, text=True) | |
| end_time = time.time() | |
| stdout.append(f"Time taken for (1) Dust3r_coarse_geometric_initialization: {end_time-start_time} sec.") | |
| stdout.append(result.stdout) | |
| stderr.append(result.stderr) | |
| # ----- (2) Train: jointly optimize pose ----- | |
| scene = 'scene' | |
| subdir = f'{n_view}_views_{gs_train_iter}Iter_{pose_lr}PoseLR' | |
| model_path = os.path.join(hashdir_path, 'infer', scene, subdir) | |
| # python -W ignore ./train_joint.py \ | |
| # -s ${SOURCE_PATH} \ | |
| # -m ${MODEL_PATH} \ | |
| # --n_views ${N_VIEW} \ | |
| # --scene ${SCENE} \ | |
| # --iter ${gs_train_iter} \ | |
| # --optim_pose \ | |
| command = [ | |
| "python", "-W", "ignore", | |
| "./train_joint.py", | |
| "-s", hashdir_path, | |
| "-m", model_path, | |
| "--n_views", str(n_view), | |
| "--scene", scene, | |
| "--iter", str(gs_train_iter), | |
| "--optim_pose", | |
| ] | |
| start_time = time.time() | |
| result = subprocess.run(command, capture_output=True, text=True) | |
| end_time = time.time() | |
| stdout.append(f"Time taken for (2) Train: jointly optimize pose: {end_time-start_time} sec.") | |
| stdout.append(result.stdout) | |
| stderr.append(result.stderr) | |
| # ----- (3) Render interpolated pose & output video ----- | |
| # python -W ignore ./render_by_interp.py \ | |
| # -s ${SOURCE_PATH} \ | |
| # -m ${MODEL_PATH} \ | |
| # --n_views ${N_VIEW} \ | |
| # --scene ${SCENE} \ | |
| # --iter ${gs_train_iter} \ | |
| # --eval \ | |
| # --get_video | |
| command = [ | |
| "python", "-W", "ignore", | |
| "./render_by_interp.py", | |
| "-s", hashdir_path, | |
| "-m", model_path, | |
| "--n_views", str(n_view), | |
| "--scene", scene, | |
| "--iter", str(gs_train_iter), | |
| "--eval", | |
| "--get_video", | |
| ] | |
| start_time = time.time() | |
| result = subprocess.run(command, capture_output=True, text=True) | |
| end_time = time.time() | |
| stdout.append(f"Time taken for (3) Render interpolated pose & output video: {end_time-start_time} sec.") | |
| stdout.append(result.stdout) | |
| stderr.append(result.stderr) | |
| stdout = '\n'.join(stdout) | |
| stderr = '\n'.join(stderr) | |
| result_tree = subprocess.run(["tree", hashdir_path], capture_output=True, text=True) | |
| ply_file = os.path.join(model_path, 'point_cloud', f'iteration_{gs_train_iter}', 'point_cloud.ply') | |
| mp4_file = os.path.join(model_path, f'{scene}_{n_view}_view.mp4') | |
| return mp4_file, ply_file, result_tree.stdout, stdout, stderr | |
| def run(): | |
| with gr.Blocks() as demo: | |
| gr.Markdown("Image and Video Processor") | |
| # Gradio interface | |
| inputs = gr.File(file_count="multiple", file_types=["image"], label="Upload Images") | |
| btn = gr.Button(value="Run") | |
| # outputs = [gr.Image(type="numpy", label="Processed Images", interactive=False), gr.Video(label="Output Video")] | |
| # outputs = [gr.Textbox(label="Log"), gr.Image(type="filepath", label="Processed Images", interactive=False)] | |
| outputs = [ | |
| gr.Video(label="Output Video"), | |
| gr.Model3D(label="3D Model"), | |
| gr.Textbox(label="Output dirs"), | |
| gr.Textbox(label="Log (stdout)"), | |
| gr.Textbox(label="Log (stderr)"), | |
| ] | |
| btn.click(fn=process_images, inputs=inputs, outputs=outputs) | |
| demo.launch(server_name="0.0.0.0", server_port=7860) | |
| if __name__ == "__main__": | |
| run() |