Update app.py

app.py

@@ -8,6 +8,7 @@ import torch
 import cv2
 import tempfile
 import shutil
+import glob
 from PIL import Image
 from typing import Iterable
 from gradio.themes import Soft
@@ -16,46 +17,53 @@ from gradio.themes.utils import colors, fonts, sizes
 # ---------------------------------------------------------
 # 1. ENVIRONMENT SETUP & REPO CLONING
 # ---------------------------------------------------------
-# Define the repository path
 REPO_URL = "https://github.com/facebookresearch/sam-3d-body.git"
 REPO_DIR = "sam-3d-body"
 
 def setup_sam_3d_env():
-    """Clones the repo and sets up paths."""
+    """
+    Clones the repo, installs dependencies, and fixes sys.path 
+    so that 'utils', 'tools', and 'sam_3d_body' can be imported.
+    """
     # 1. Clone if not exists
     if not os.path.exists(REPO_DIR):
         print(f"Cloning SAM 3D Body repository from {REPO_URL}...")
         try:
             subprocess.run(["git", "clone", REPO_URL], check=True)
-            # Install the package in editable mode to handle internal imports
-            print("Installing sam-3d-body package...")
+            print("Installing sam-3d-body package in editable mode...")
+            # We install using pip to resolve internal package dependencies
             subprocess.run([sys.executable, "-m", "pip", "install", "-e", REPO_DIR], check=True)
+            
+            # Install other requirements usually needed
+            subprocess.run([sys.executable, "-m", "pip", "install", "trimesh", "opencv-python", "matplotlib"], check=True)
         except subprocess.CalledProcessError as e:
             print(f"Error during setup: {e}")
             return False
 
-    # 2. Add paths to sys.path
+    # 2. Add Critical Paths to sys.path
     repo_abs_path = os.path.abspath(REPO_DIR)
     notebook_path = os.path.join(repo_abs_path, "notebook")
     
-    # Add repo root (for sam_3d_body package)
+    # CRITICAL: Add repo root first so 'import tools' and 'import sam_3d_body' work inside utils.py
     if repo_abs_path not in sys.path:
         sys.path.insert(0, repo_abs_path)
+        print(f"Added to sys.path: {repo_abs_path}")
     
-    # Add notebook folder (for utils.py)
+    # Add notebook folder so we can 'import utils'
     if notebook_path not in sys.path:
         sys.path.insert(0, notebook_path)
+        print(f"Added to sys.path: {notebook_path}")
         
     return True
 
-# Run setup
+# Run setup immediately
 env_ready = setup_sam_3d_env()
 
 # ---------------------------------------------------------
 # 2. IMPORTS
 # ---------------------------------------------------------
 
-# Import SAM3 (Transformers)
+# --- Import SAM3 (Segmentation) ---
 try:
     from transformers import Sam3Processor, Sam3Model
     SAM3_AVAILABLE = True
@@ -63,24 +71,24 @@ except ImportError:
     print("Warning: transformers library not found or outdated. SAM3 will be disabled.")
     SAM3_AVAILABLE = False
 
-# Import SAM 3D Body Utils
+# --- Import SAM 3D Body Utils ---
+# We use a specific alias to avoid confusion with standard python utils
+sam3d_utils = None
 SAM3D_AVAILABLE = False
+
 if env_ready:
     try:
-        # Import specific functions from the notebook/utils.py
-        # Note: We rely on the path insertion above to find 'utils'
-        from utils import (
-            setup_sam_3d_body, 
-            setup_visualizer, 
-            visualize_2d_results, 
-            visualize_3d_mesh, 
-            save_mesh_results
-        )
+        # Now that sys.path is fixed, this import should work
+        # and utils.py will successfully find 'tools' and 'sam_3d_body'
+        import utils as sam3d_utils_module
+        sam3d_utils = sam3d_utils_module
         SAM3D_AVAILABLE = True
         print("SAM 3D Body utils imported successfully.")
     except ImportError as e:
         print(f"Error importing SAM 3D Body utils: {e}")
-        print("Ensure requirements are installed (pytorch3d, opencv, etc.)")
+        print("This usually happens if 'tools' or 'sam_3d_body' cannot be found by utils.py")
+        import traceback
+        traceback.print_exc()
 
 # ---------------------------------------------------------
 # 3. THEME DEFINITION
@@ -154,7 +162,7 @@ print(f"Using device: {device}")
 # 4. LOAD MODELS
 # ---------------------------------------------------------
 
-# --- Load SAM3 ---
+# --- 1. Load SAM3 ---
 sam3_model = None
 sam3_processor = None
 if SAM3_AVAILABLE:
@@ -166,21 +174,27 @@ if SAM3_AVAILABLE:
     except Exception as e:
         print(f"Error loading SAM3: {e}")
 
-# --- Load SAM 3D Body ---
+# --- 2. Load SAM 3D Body ---
 sam3d_estimator = None
 sam3d_visualizer = None
 
 if SAM3D_AVAILABLE:
     try:
-        print("Loading SAM 3D Body Estimator...")
-        # Note: This might require huggingface_hub login if the repo is gated, 
-        # but facebook/sam-3d-body-dinov3 is usually public.
-        sam3d_estimator = setup_sam_3d_body(hf_repo_id="facebook/sam-3d-body-dinov3")
-        sam3d_visualizer = setup_visualizer()
-        print("SAM 3D Body Loaded.")
+        print("Loading SAM 3D Body Estimator (this may take a moment)...")
+        # Initialize estimator using the utility function from the repo
+        # Note: detector_name="vitdet" is default, requiring 'tools' import to work
+        sam3d_estimator = sam3d_utils.setup_sam_3d_body(
+            hf_repo_id="facebook/sam-3d-body-dinov3", 
+            device=device
+        )
+        sam3d_visualizer = sam3d_utils.setup_visualizer()
+        print("SAM 3D Body Loaded Successfully.")
     except Exception as e:
         print(f"Error loading SAM 3D Body model: {e}")
+        # If it fails, we set the flag to False so the UI handles it gracefully
         SAM3D_AVAILABLE = False
+        import traceback
+        traceback.print_exc()
 
 # ---------------------------------------------------------
 # 5. INFERENCE FUNCTIONS
@@ -188,6 +202,7 @@ if SAM3D_AVAILABLE:
 
 @spaces.GPU
 def segment_image(input_image, text_prompt, threshold=0.5):
+    """Handler for Tab 1: Segmentation"""
     if input_image is None:
         raise gr.Error("Please upload an image.")
     if not text_prompt:
@@ -221,66 +236,80 @@ def segment_image(input_image, text_prompt, threshold=0.5):
 
 @spaces.GPU
 def process_3d_body(input_image):
+    """Handler for Tab 2: 3D Body Reconstruction"""
     if input_image is None:
         raise gr.Error("Please upload an image.")
+    
     if not SAM3D_AVAILABLE or sam3d_estimator is None:
-        raise gr.Error("SAM 3D Body libraries or model not available (Check logs for import errors).")
+        raise gr.Error("SAM 3D Body libraries or model failed to load. Check console logs.")
 
-    # Prepare Image
+    # Convert PIL to CV2 BGR for the estimator
     img_np = np.array(input_image.convert("RGB"))
     img_cv2 = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
     
-    # The utils/estimator usually requires a file path
+    # The estimator.process_one_image expects a 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 for {tmp_path}...")
+        
+        # 1. Run Inference
+        # process_one_image is a method of the estimator class inside sam-3d-body
         outputs = sam3d_estimator.process_one_image(tmp_path)
         
         if not outputs:
             return None, None, None, "No people detected."
 
-        # 1. 2D Visuals
-        vis_results_2d = visualize_2d_results(img_cv2, outputs, sam3d_visualizer)
-        # Handle case if visualize_2d_results returns list of images (one per person)
-        if isinstance(vis_results_2d, list) and len(vis_results_2d) > 0:
-            # Just take the first one or combine them?
-            # Usually it returns cropped visuals. Let's assume list of images.
+        # 2. 2D Keypoints Visualization
+        vis_results_2d = sam3d_utils.visualize_2d_results(img_cv2, outputs, sam3d_visualizer)
+        # Combine if multiple, or just take first for display simplicity. 
+        # Usually vis_results_2d is a list of full images with drawings.
+        if vis_results_2d:
+            # For simplicity, if multiple people, the last one overrides or we assume 1 main person
+            # Ideally we'd grid them, but for Gradio output, let's take the first result's image
             res_2d_rgb = cv2.cvtColor(vis_results_2d[0], cv2.COLOR_BGR2RGB)
         else:
             res_2d_rgb = img_np
 
-        # 2. 3D Overlay Visuals
-        mesh_results_img = visualize_3d_mesh(img_cv2, outputs, sam3d_estimator.faces)
-        if isinstance(mesh_results_img, list) and len(mesh_results_img) > 0:
-            res_3d_overlay_rgb = cv2.cvtColor(mesh_results_img[0], cv2.COLOR_BGR2RGB)
+        # 3. 3D Overlay Visualization
+        # visualize_3d_mesh returns a wide image (Original | Overlay | White | Side)
+        mesh_results_wide = sam3d_utils.visualize_3d_mesh(img_cv2, outputs, sam3d_estimator.faces)
+        if mesh_results_wide:
+            res_3d_overlay_rgb = cv2.cvtColor(mesh_results_wide[0], cv2.COLOR_BGR2RGB)
         else:
             res_3d_overlay_rgb = img_np
 
-        # 3. Save PLY for Model3D
+        # 4. Save PLY for Model3D
+        # Create a unique directory for this run
         output_dir = tempfile.mkdtemp()
         image_name = "gradio_mesh"
         
         # save_mesh_results returns list of paths to .ply files
-        ply_files = save_mesh_results(img_cv2, outputs, sam3d_estimator.faces, output_dir, image_name)
+        ply_files = sam3d_utils.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 mesh found
             
-        status = f"Detected {len(outputs)} person(s). Showing result for Person 0."
+        status_msg = f"Detected {len(outputs)} person(s). Displaying Person 0."
         
-        return res_2d_rgb, res_3d_overlay_rgb, ply_path, status
+        return res_2d_rgb, res_3d_overlay_rgb, ply_path, status_msg
 
     except Exception as e:
         import traceback
         traceback.print_exc()
-        raise gr.Error(f"Inference failed: {e}")
+        raise gr.Error(f"Inference failed: {str(e)}")
 
     finally:
+        # Cleanup input temp file
         if os.path.exists(tmp_path):
             os.remove(tmp_path)
 
@@ -294,6 +323,7 @@ css = """
     max-width: 1200px;
 }
 #main-title h1 {font-size: 2.1em !important; text-align: center;}
+.gradio-container {min-height: 0px !important;}
 """
 
 with gr.Blocks(css=css, theme=steel_blue_theme) as demo:
@@ -301,7 +331,7 @@ with gr.Blocks(css=css, theme=steel_blue_theme) as demo:
         gr.Markdown("# **SAM Integrated Vision Suite**", elem_id="main-title")
         
         with gr.Tabs():
-            # TAB 1: SEGMENTATION
+            # ================= TAB 1: SEGMENTATION =================
             with gr.Tab("SAM3 Segmentation"):
                 gr.Markdown("Segment objects using **SAM3** with text prompts.")
                 with gr.Row():
@@ -315,7 +345,10 @@ with gr.Blocks(css=css, theme=steel_blue_theme) as demo:
 
                 t1_btn.click(segment_image, [t1_input, t1_prompt, t1_thresh], [t1_output])
 
-            # TAB 2: 3D BODY
+                # Optional examples if files exist
+                # gr.Examples(...)
+
+            # ================= TAB 2: 3D BODY =================
             with gr.Tab("SAM 3D Body"):
                 gr.Markdown("Detect human bodies and reconstruct **3D Meshes**.")
                 
@@ -328,12 +361,12 @@ with gr.Blocks(css=css, theme=steel_blue_theme) as demo:
                     with gr.Column(scale=2):
                         with gr.Row():
                             t2_vis_2d = gr.Image(label="2D Detection", type="numpy")
-                            t2_vis_overlay = gr.Image(label="Mesh Overlay", type="numpy")
+                            t2_vis_overlay = gr.Image(label="3D Visualization (Original | Overlay | White | Side)", type="numpy")
                         
                         t2_model_3d = gr.Model3D(
                             label="Interactive 3D Mesh",
                             clear_color=[0.0, 0.0, 0.0, 0.0],
-                            camera_position=[0, 0, 2.5]
+                            camera_position=[0, 0, 4.0]
                         )
 
                 t2_btn.click(
@@ -341,12 +374,6 @@ with gr.Blocks(css=css, theme=steel_blue_theme) as demo:
                     inputs=[t2_input],
                     outputs=[t2_vis_2d, t2_vis_overlay, t2_model_3d, t2_status]
                 )
-                
-                gr.Examples(
-                    examples=[["examples/player.jpg"], ["examples/dancing.jpg"]],
-                    inputs=[t2_input],
-                    label="3D Body Examples"
-                )
 
 if __name__ == "__main__":
     demo.launch(mcp_server=True, ssr_mode=False, show_error=True)