Update app.py

app.py

@@ -1,57 +1,135 @@
+# app.py
 import gradio as gr
-import torch
+import cv2
+import mediapipe as mp
 import numpy as np
-from PIL import Image
-import torchvision.transforms as transforms
-from decalib.deca import DECA
-from decalib.utils import util
-from decalib.utils.config import cfg as deca_cfg
-
-# Initialize DECA
-deca_cfg.model.use_tex = False
-deca = DECA(config=deca_cfg)
-
-def preprocess_image(image):
-    transform = transforms.Compose([
-        transforms.Resize(224),
-        transforms.ToTensor(),
-        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-    ])
-    return transform(image).unsqueeze(0)
-
-def reconstruct_face(image):
-    # Preprocess the image
-    input_image = preprocess_image(Image.fromarray(image))
-    
-    # Run DECA
-    with torch.no_grad():
-        codedict = deca.encode(input_image)
-        opdict, visdict = deca.decode(codedict)
-    
-    # Get the reconstructed face
-    reconstructed_face = util.tensor2image(visdict['shape_images'][0])
-    
-    # Get the 3D mesh
-    vertices = opdict['vertices'][0].cpu().numpy()
-    faces = deca.flame.faces_tensor.cpu().numpy()
-    
-    return reconstructed_face, vertices, faces
-
-def process_image(input_image):
-    reconstructed_face, vertices, faces = reconstruct_face(input_image)
-    return reconstructed_face, [vertices, faces]
-
-# Define the Gradio interface
-iface = gr.Interface(
-    fn=process_image,
-    inputs=gr.Image(),
-    outputs=[
-        gr.Image(label="Reconstructed Face"),
-        gr.Model3D(label="3D Face Model")
-    ],
-    title="3D Face Reconstruction from a Single Image",
-    description="Upload an image of a face to generate a 3D reconstruction."
-)
-
-# Launch the app
-iface.launch()
+import time
+from pathlib import Path
+import tempfile
+import os
+from face_reconstruction_main import FaceMeshDetector
+import plotly.graph_objects as go
+import json
+
+# Initialize the detector
+detector = FaceMeshDetector()
+
+def process_frame(frame):
+    """Process a single frame and return the image with face mesh overlay"""
+    if frame is None:
+        return None
+    
+    img, faces = detector.detect_faces(frame)
+    return img, faces
+
+def create_3d_plot(vertices):
+    """Create a 3D plotly figure from vertices"""
+    x, y, z = zip(*vertices)
+    
+    fig = go.Figure(data=[go.Scatter3d(
+        x=x, y=y, z=z,
+        mode='markers',
+        marker=dict(
+            size=2,
+            color=z,
+            colorscale='Viridis',
+        )
+    )])
+    
+    fig.update_layout(
+        scene=dict(
+            xaxis_title='X',
+            yaxis_title='Y',
+            zaxis_title='Z'
+        ),
+        margin=dict(l=0, r=0, b=0, t=0)
+    )
+    
+    return fig
+
+def save_obj_and_display(frame):
+    """Save frame as OBJ and return both file path and 3D visualization"""
+    if frame is None:
+        return None, None, "No frame provided"
+    
+    img, faces = process_frame(frame)
+    
+    if not faces:
+        return None, None, "No face detected in the frame"
+    
+    # Scale vertices
+    vertices = np.array(faces[0]) * 150
+    
+    # Create temporary directory if it doesn't exist
+    temp_dir = Path('temp')
+    temp_dir.mkdir(exist_ok=True)
+    
+    # Generate timestamp
+    timestamp = time.strftime("%Y%m%d-%H%M%S")
+    filename = f"face_mesh_{timestamp}"
+    
+    # Save OBJ file
+    obj_path = temp_dir / f"{filename}.obj"
+    success = detector.save_obj_file(vertices, filename)
+    
+    if not success:
+        return None, None, "Failed to save OBJ file"
+    
+    # Create 3D visualization
+    fig = create_3d_plot(vertices)
+    
+    return str(obj_path), fig, "Model generated successfully"
+
+def cleanup_old_files():
+    """Remove files older than 1 hour"""
+    temp_dir = Path('temp')
+    if temp_dir.exists():
+        current_time = time.time()
+        for file in temp_dir.glob('*.obj'):
+            if current_time - file.stat().st_mtime > 3600:  # 1 hour
+                file.unlink()
+
+def create_interface():
+    """Create the Gradio interface"""
+    with gr.Blocks() as interface:
+        gr.Markdown("""
+        # 3D Face Reconstruction
+        Upload an image or use your webcam to generate a 3D face model.
+        
+        Instructions:
+        1. Upload an image or use the webcam
+        2. Click 'Generate 3D Model'
+        3. View the 3D model in the interactive viewer
+        4. Download the OBJ file for use in 3D software
+        """)
+        
+        with gr.Row():
+            with gr.Column():
+                # Input methods
+                input_image = gr.Image(source="webcam", type="numpy")
+                generate_button = gr.Button("Generate 3D Model")
+            
+            with gr.Column():
+                # Output displays
+                obj_file = gr.File(label="Download OBJ File")
+                plot_3d = gr.Plot(label="3D Preview")
+                status_text = gr.Textbox(label="Status")
+        
+        # Set up event handler
+        generate_button.click(
+            fn=save_obj_and_display,
+            inputs=[input_image],
+            outputs=[obj_file, plot_3d, status_text]
+        )
+        
+        # Cleanup old files periodically
+        cleanup_old_files()
+    
+    return interface
+
+# Create and launch the interface
+interface = create_interface()
+
+# Launch locally for testing
+if __name__ == "__main__":
+    interface.launch()