Upload app.py

app.py

@@ -6,6 +6,7 @@ Two-step interactive demo:
 """
 
 import csv
+import datetime
 import os
 import shutil
 import subprocess
@@ -13,6 +14,13 @@ import tempfile
 from pathlib import Path
 from typing import List, Optional, Tuple
 
+
+def log(msg: str):
+    """Print with Beijing time (UTC+8) prefix."""
+    beijing_time = datetime.datetime.now(datetime.timezone(datetime.timedelta(hours=8)))
+    timestamp = beijing_time.strftime("%Y-%m-%d %H:%M:%S")
+    print(f"[{timestamp}] {msg}")
+
 import cv2
 import gradio as gr
 import numpy as np
@@ -76,12 +84,12 @@ def load_model(checkpoint_path: str = "checkpoints"):
 
     if model_path is None:
         model_path = HUGGINGFACE_REPO
-        print(f"Local checkpoint not found at '{checkpoint_path}', loading from HuggingFace: {HUGGINGFACE_REPO}")
+        log(f"Local checkpoint not found at '{checkpoint_path}', loading from HuggingFace: {HUGGINGFACE_REPO}")
 
     # Skip redundant backbone weight download – from_pretrained will
     # overwrite all parameters from the safetensors file anyway.
     Sat3DGen._skip_backbone_weights = True
-    print(f"Loading model from {model_path} ...")
+    log(f"Loading model from {model_path} ...")
     MODEL = Sat3DGen.from_pretrained(model_path).to(DEVICE)
     Sat3DGen._skip_backbone_weights = False
     MODEL.eval()
@@ -91,7 +99,7 @@ def load_model(checkpoint_path: str = "checkpoints"):
         T.ToTensor(),
         T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
     ])
-    print("Model loaded successfully.")
+    log("Model loaded successfully.")
 
 
 # ---------------------------------------------------------------------------
@@ -281,38 +289,38 @@ def generate_mesh(sat_image_pil: Image.Image, mesh_resolution: int = 256, progre
     if sat_image_pil is None:
         raise gr.Error("Please upload a satellite image first.")
 
-    print("[generate_mesh] >>> Start")
+    log("[generate_mesh] >>> Start")
     load_model()
-    print("[generate_mesh] Model loaded")
+    log("[generate_mesh] Model loaded")
 
     progress(0.1, desc="Preprocessing satellite image...")
-    print("[generate_mesh] Preprocessing satellite image...")
+    log("[generate_mesh] Preprocessing satellite image...")
     sat_input = SAT_TRANSFORM(sat_image_pil.convert("RGB")).unsqueeze(0).to(DEVICE)
 
     progress(0.3, desc="Generating triplane features...")
-    print("[generate_mesh] Generating triplane features...")
+    log("[generate_mesh] Generating triplane features...")
     with torch.no_grad():
         triplane = MODEL.from_sat_to_triplane(sat_input)
-    print("[generate_mesh] Triplane generated successfully")
+    log("[generate_mesh] Triplane generated successfully")
 
     progress(0.5, desc="Extracting 3D mesh (this may take a moment)...")
-    print(f"[generate_mesh] Extracting 3D mesh (resolution={mesh_resolution})...")
+    log(f"[generate_mesh] Extracting 3D mesh (resolution={mesh_resolution})...")
     with torch.no_grad():
         vertices, faces, vertex_colors = MODEL.extract_mesh(triplane, mesh_resolution=mesh_resolution)
-    print(f"[generate_mesh] Mesh extracted: {vertices.shape[0]} vertices, {faces.shape[0]} faces")
+    log(f"[generate_mesh] Mesh extracted: {vertices.shape[0]} vertices, {faces.shape[0]} faces")
 
     vertices = vertices[:, [1, 2, 0]]
 
     # Save mesh
     mesh_path = str(RESULTS_DIR / "mesh.obj")
     save_obj(vertices, faces, vertex_colors, mesh_path)
-    print(f"[generate_mesh] OBJ saved to {mesh_path}")
+    log(f"[generate_mesh] OBJ saved to {mesh_path}")
 
     # Also save triplane to state for Step 2
     state = {"triplane": triplane, "sat_image": sat_image_pil}
 
     progress(0.9, desc="Preparing 3D visualization...")
-    print("[generate_mesh] Converting OBJ → GLB for 3D preview...")
+    log("[generate_mesh] Converting OBJ → GLB for 3D preview...")
 
     # Create a glb file for Gradio's Model3D component.
     # Use a tempfile so Gradio can reliably serve it via its file cache.
@@ -329,17 +337,17 @@ def generate_mesh(sat_image_pil: Image.Image, mesh_resolution: int = 256, progre
             raise gr.Error("Failed to load mesh geometry.")
     if not hasattr(mesh_trimesh, 'vertex_normals') or mesh_trimesh.vertex_normals is None or len(mesh_trimesh.vertex_normals) == 0:
         mesh_trimesh.vertex_normals  # triggers auto-computation
-    print(f"[generate_mesh] Mesh has {len(mesh_trimesh.vertices)} verts, {len(mesh_trimesh.faces)} faces, normals: {mesh_trimesh.vertex_normals.shape}")
+    log(f"[generate_mesh] Mesh has {len(mesh_trimesh.vertices)} verts, {len(mesh_trimesh.faces)} faces, normals: {mesh_trimesh.vertex_normals.shape}")
     mesh_trimesh.export(glb_path_local, file_type="glb")
-    print(f"[generate_mesh] GLB saved to {glb_path_local} ({os.path.getsize(glb_path_local)} bytes)")
+    log(f"[generate_mesh] GLB saved to {glb_path_local} ({os.path.getsize(glb_path_local)} bytes)")
 
     tmp_glb = tempfile.NamedTemporaryFile(suffix=".glb", delete=False)
     shutil.copy2(glb_path_local, tmp_glb.name)
     tmp_glb.close()
-    print(f"[generate_mesh] GLB copied to temp file: {tmp_glb.name}")
+    log(f"[generate_mesh] GLB copied to temp file: {tmp_glb.name}")
 
     progress(1.0, desc="Done!")
-    print("[generate_mesh] <<< 3D mesh generated successfully!")
+    log("[generate_mesh] <<< 3D mesh generated successfully!")
     return tmp_glb.name, mesh_path, state
 
 
@@ -366,14 +374,14 @@ def render_trajectory_video(
       Top row:    satellite image (with camera marker)  |  panorama RGB
       Bottom row: 4 perspective views in a horizontal row (left, front, right, back)
     """
-    print("[render_trajectory_video] >>> Start")
+    log("[render_trajectory_video] >>> Start")
     load_model()
 
     sat_size = sat_image_pil.size[0]
     positions, pixel_coords = read_trajectory_from_csv(trajectory_csv_path, sat_size)
     if len(positions) == 0:
         raise gr.Error(f"Trajectory file is empty: {trajectory_csv_path}")
-    print(f"[render_trajectory_video] Loaded {len(positions)} positions from {trajectory_csv_path}")
+    log(f"[render_trajectory_video] Loaded {len(positions)} positions from {trajectory_csv_path}")
 
     progress(0.1, desc="Extracting triplane features...")
     sat_tensor = SAT_TRANSFORM(sat_image_pil.convert("RGB")).unsqueeze(0).to(DEVICE)
@@ -399,7 +407,7 @@ def render_trajectory_video(
     for idx, position in enumerate(positions):
         progress(0.25 + 0.6 * idx / total_positions, desc=f"Rendering frame {idx + 1}/{total_positions}...")
         if idx % 10 == 0 or idx == total_positions - 1:
-            print(f"[render_trajectory_video] Rendering frame {idx + 1}/{total_positions}...")
+            log(f"[render_trajectory_video] Rendering frame {idx + 1}/{total_positions}...")
 
         c2w = position_to_c2w(position)
         c2w[:, :3, 3] = c2w[:, :3, 3] * MODEL.position_scale_factor
@@ -455,7 +463,7 @@ def render_trajectory_video(
         cv2.imwrite(str(frame_path), cv2.cvtColor(composed, cv2.COLOR_RGB2BGR))
 
     progress(0.9, desc="Encoding video...")
-    print("[render_trajectory_video] All frames rendered, encoding video with ffmpeg...")
+    log("[render_trajectory_video] All frames rendered, encoding video with ffmpeg...")
     video_path = str(RESULTS_DIR / "trajectory_video.mp4")
     ffmpeg_path = shutil.which("ffmpeg")
     if ffmpeg_path is None:
@@ -469,7 +477,7 @@ def render_trajectory_video(
         video_path,
     ], check=True, capture_output=True)
 
-    print(f"[render_trajectory_video] Video saved to {video_path}")
+    log(f"[render_trajectory_video] Video saved to {video_path}")
     progress(1.0, desc="Done!")
     return video_path
 
@@ -533,6 +541,9 @@ def build_demo():
 
                 with gr.Column(scale=2):
                     mesh_viewer = gr.Model3D(label="3D Mesh Preview", height=500)
+                    gr.Markdown(
+                        "⏳ *After generation completes, the 3D preview may take ~10 seconds to load. Please wait.*"
+                    )
                     download_button = gr.DownloadButton("💾 Download Mesh (.obj)", variant="secondary")
 
             if sample_sat_images:
@@ -670,6 +681,9 @@ def build_demo():
                         )
                     sky_status = gr.Markdown(value=default_sky_message)
                     render_button = gr.Button("🎬 Render Video", variant="primary", size="lg")
+                    gr.Markdown(
+                        "⏳ *Running on CPU — video rendering is slow (~5 min for 80 frames). Please be patient.*"
+                    )
 
                 # Middle column: trajectory preview
                 with gr.Column(scale=1):