Update app.py

app.py

@@ -11,6 +11,9 @@ import subprocess
 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 = []
@@ -24,14 +27,18 @@ def move_uploaded_files(files, temp_dir):
 
 def process_images(files):
     n_view = len(files)
+    
     files_str = ','.join([f.name for f in files])
-    hashdir = hashlib.md5(files_str.encode('utf8')).hexdigest()
-    hashdir_path = os.path.join(TMP_DIR, hashdir)
+    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",
@@ -41,11 +48,69 @@ def process_images(files):
     ]
     
     result = subprocess.run(command, capture_output=True, text=True)
-    result_tree = subprocess.run(["tree", hashdir_path], capture_output=True, text=True)
+    stdout.append(result.stdout)
+    stderr.append(result.stderr)
+
+    # ----- (2) Train: jointly optimize pose -----
+    # python -W ignore ./train_joint.py \
+    #         -s ${SOURCE_PATH} \
+    #         -m ${MODEL_PATH}  \
+    #         --n_views ${N_VIEW}  \
+    #         --scene ${SCENE} \
+    #         --iter ${gs_train_iter} \
+    #         --optim_pose \
+    
+    scene = hashdir_name
+    model_path = os.path.join(hashdir_path, 'infer', scene, f'{n_view}_views_{gs_train_iter}Iter_{pose_lr}PoseLR')
+    
+    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",
+    ]
 
+    result = subprocess.run(command, capture_output=True, text=True)
+    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",
+    ]
+    
+    result = subprocess.run(command, capture_output=True, text=True)
+    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(hashdir_path, 'sparse/0/', 'points3D.ply')
     
-    return ply_file, result_tree.stdout, result.stdout, result.stderr
+    return result_tree.stdout, stdout, stderr
 
 def run():
     with gr.Blocks() as demo:
@@ -56,7 +121,7 @@ def 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.Model3D(label="3D Model"),
+            # gr.Model3D(label="3D Model"),
             gr.Textbox(label="Output dirs"), 
             gr.Textbox(label="Log (stdout)"), 
             gr.Textbox(label="Log (stderr)"),