Update app.py

app.py

@@ -8,59 +8,63 @@ import torch
 import cv2
 import tempfile
 import shutil
+import traceback
 from PIL import Image
 from typing import Iterable
 from gradio.themes import Soft
 from gradio.themes.utils import colors, fonts, sizes
 from transformers import Sam3Processor, Sam3Model
 
+# ---------------------------------------------------------
+# 0. HEADLESS RENDERING SETUP
+# ---------------------------------------------------------
+# Essential for running 3D visualizers (pyrender) in cloud environments
+os.environ["PYOPENGL_PLATFORM"] = "egl"
+
 # ---------------------------------------------------------
 # 1. SETUP & DYNAMIC IMPORTS
 # ---------------------------------------------------------
 
-# Define repo details
 REPO_URL = "https://github.com/facebookresearch/sam-3d-body.git"
 REPO_DIR = "sam-3d-body"
 
 def setup_environment():
-    """
-    Clones the repo if missing and sets up paths.
-    """
+    """Clones repo and sets up Python paths."""
     print("--- Checking Environment for SAM 3D Body ---")
     
-    # 1. Clone Repository if not present
+    # 1. Clone Repository
     if not os.path.exists(REPO_DIR):
         print(f"Cloning {REPO_URL}...")
         try:
             subprocess.run(["git", "clone", REPO_URL], check=True)
             print("Repository cloned.")
-            
-            # Optional: Install the package in editable mode to ensure internal imports work
-            print("Installing sam-3d-body package...")
+            # Install in editable mode to ensure package discovery works
             subprocess.run([sys.executable, "-m", "pip", "install", "-e", f"./{REPO_DIR}"], check=True)
         except Exception as e:
-            print(f"Error during setup: {e}")
+            print(f"Git clone/Install failed: {e}")
 
     # 2. Add paths to sys.path
-    # Add the root of the repo
     repo_abs = os.path.abspath(REPO_DIR)
-    if repo_abs not in sys.path:
-        sys.path.append(repo_abs)
-    
-    # Add the notebook folder (where utils.py lives)
     notebook_abs = os.path.abspath(os.path.join(REPO_DIR, "notebook"))
+    
+    if repo_abs not in sys.path:
+        sys.path.insert(0, repo_abs)
     if notebook_abs not in sys.path:
-        sys.path.append(notebook_abs)
-
-    print(f"Paths added: {repo_abs}, {notebook_abs}")
+        sys.path.insert(0, notebook_abs)
+        
+    print(f"Python Paths: {sys.path[:2]}...")
 
-# Run setup
 setup_environment()
 
-# 3. Attempt Imports
+# Global variables for models and error tracking
+sam3d_estimator = None
+sam3d_visualizer = None
+sam3d_load_error = None
+SAM3D_AVAILABLE = False
+
 try:
-    # utils.py is in sam-3d-body/notebook/utils.py
-    # We added 'sam-3d-body/notebook' to path, so we can import directly from utils
+    # Try importing from the utils file in the cloned repo
+    # This expects 'notebook/utils.py' to exist
     from utils import (
         setup_sam_3d_body, 
         setup_visualizer, 
@@ -68,11 +72,18 @@ try:
         visualize_3d_mesh, 
         save_mesh_results
     )
+    
+    print("Loading SAM 3D Body Estimator (this may take time)...")
+    # Initialize the model immediately to catch errors early
+    sam3d_estimator = setup_sam_3d_body(hf_repo_id="facebook/sam-3d-body-dinov3")
+    sam3d_visualizer = setup_visualizer()
     SAM3D_AVAILABLE = True
-    print("SAM 3D Body utilities imported successfully.")
-except ImportError as e:
-    print(f"CRITICAL WARNING: SAM 3D Body utils not found: {e}")
-    print("Ensure dependencies (mmhuman3d, detectron2, etc.) are installed in requirements.txt")
+    print("SAM 3D Body Model loaded successfully.")
+
+except Exception as e:
+    # Capture the exact error (e.g., missing mmhuman3d)
+    sam3d_load_error = f"{type(e).__name__}: {str(e)}\n{traceback.format_exc()}"
+    print(f"CRITICAL ERROR loading SAM 3D Body:\n{sam3d_load_error}")
     SAM3D_AVAILABLE = False
 
 # ---------------------------------------------------------
@@ -144,12 +155,10 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"Using device: {device}")
 
 # ---------------------------------------------------------
-# 3. MODEL LOADING
+# 3. MODEL LOADING (SAM3)
 # ---------------------------------------------------------
-
-# --- Load SAM3 (Segmentation) ---
 try:
-    print("Loading SAM3 Model and Processor...")
+    print("Loading SAM3 Model...")
     sam3_model = Sam3Model.from_pretrained("facebook/sam3").to(device)
     sam3_processor = Sam3Processor.from_pretrained("facebook/sam3")
     print("SAM3 Model loaded successfully.")
@@ -158,22 +167,6 @@ except Exception as e:
     sam3_model = None 
     sam3_processor = None
 
-# --- Load SAM 3D Body ---
-sam3d_estimator = None
-sam3d_visualizer = None
-
-if SAM3D_AVAILABLE:
-    try:
-        print("Loading SAM 3D Body Estimator...")
-        # Note: The utils function usually expects the HF repo ID
-        sam3d_estimator = setup_sam_3d_body(hf_repo_id="facebook/sam-3d-body-dinov3")
-        sam3d_visualizer = setup_visualizer()
-        print("SAM 3D Body Model loaded successfully.")
-    except Exception as e:
-        print(f"Error loading SAM 3D Body model: {e}")
-        # Mark as unavailable if loading fails despite imports working
-        SAM3D_AVAILABLE = False
-
 # ---------------------------------------------------------
 # 4. INFERENCE FUNCTIONS
 # ---------------------------------------------------------
@@ -217,59 +210,53 @@ def process_3d_body(input_image):
     if input_image is None:
         raise gr.Error("Please upload an image.")
     
+    # Check if initialization failed
     if not SAM3D_AVAILABLE or sam3d_estimator is None:
-        raise gr.Error("SAM 3D Body model not available. Check server logs for import errors.")
+        # Raise the specific error captured during startup
+        error_msg = sam3d_load_error if sam3d_load_error else "Unknown initialization error."
+        raise gr.Error(f"Model Setup Failed. Logs:\n{error_msg}")
 
     # Convert PIL to CV2 BGR
     img_np = np.array(input_image.convert("RGB"))
     img_cv2 = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
     
-    # Helper requires a physical file path for inference
+    # Helper requires a physical file path
     with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp_file:
         tmp_path = tmp_file.name
         cv2.imwrite(tmp_path, img_cv2)
     
     try:
-        # Run Inference
-        print(f"Processing 3D Body on {tmp_path}...")
+        print(f"Running inference on {tmp_path}...")
         outputs = sam3d_estimator.process_one_image(tmp_path)
         
         if not outputs:
             return None, None, None, "No people detected."
 
-        # 1. Generate 2D Visualization
+        # 1. 2D Vis
         vis_results_2d = visualize_2d_results(img_cv2, outputs, sam3d_visualizer)
-        # Handle list of results
         res_2d_rgb = cv2.cvtColor(vis_results_2d[0], cv2.COLOR_BGR2RGB) if vis_results_2d else img_np
         
-        # 2. Generate 3D Visualization (Overlay Image)
+        # 2. 3D Overlay
         mesh_results_img = visualize_3d_mesh(img_cv2, outputs, sam3d_estimator.faces)
         res_3d_overlay_rgb = cv2.cvtColor(mesh_results_img[0], cv2.COLOR_BGR2RGB) if mesh_results_img else img_np
 
-        # 3. Save PLY Mesh to temp directory for Gradio Model3D
+        # 3. Save PLY
         output_dir = tempfile.mkdtemp()
         image_name = "gradio_mesh_result"
-        
-        # utils.save_mesh_results returns a list of paths to .ply files
         ply_files = save_mesh_results(img_cv2, outputs, sam3d_estimator.faces, output_dir, image_name)
         
-        ply_path = None
-        if ply_files and len(ply_files) > 0:
-            ply_path = ply_files[0] # Return the first person's mesh
+        ply_path = ply_files[0] if ply_files else None
             
         status = f"Success! Detected {len(outputs)} person(s)."
         
         return res_2d_rgb, res_3d_overlay_rgb, ply_path, status
 
     except Exception as e:
-        raise gr.Error(f"Inference failed: {str(e)}")
+        traceback.print_exc()
+        raise gr.Error(f"Inference Runtime Error: {str(e)}")
     finally:
-        # Cleanup input temp file
         if os.path.exists(tmp_path):
-            try:
-                os.remove(tmp_path)
-            except:
-                pass
+            os.remove(tmp_path)
 
 # ---------------------------------------------------------
 # 5. GRADIO UI LAYOUT
@@ -317,6 +304,15 @@ with gr.Blocks(css=css, theme=steel_blue_theme) as demo:
                         t2_input_image = gr.Image(label="Input Image", type="pil", height=350)
                         t2_run_btn = gr.Button("Generate 3D Body", variant="primary")
                         t2_status = gr.Textbox(label="Status", interactive=False)
+                        
+                        # Warning box if initialization failed
+                        if not SAM3D_AVAILABLE:
+                            gr.Markdown(
+                                "⚠️ **Warning: SAM 3D Body failed to load.**\n"
+                                f"Error: {sam3d_load_error}\n"
+                                "Please check `mmhuman3d` and `mmcv` dependencies.",
+                                elem_classes=["error-box"]
+                            )
 
                     with gr.Column(scale=2):
                         with gr.Row():
@@ -335,7 +331,7 @@ with gr.Blocks(css=css, theme=steel_blue_theme) as demo:
                     outputs=[t2_output_2d, t2_output_overlay, t2_output_3d, t2_status]
                 )
                 
-                # Check if example files exist before creating component to avoid errors
+                # Dynamic examples
                 ex_files = []
                 if os.path.exists("examples/player.jpg"): ex_files.append(["examples/player.jpg"])
                 if os.path.exists("examples/dancing.jpg"): ex_files.append(["examples/dancing.jpg"])