Update app.py

app.py

@@ -1,6 +1,7 @@
 import os
 import sys
 import logging
+import tempfile
 import numpy as np
 import cv2
 import trimesh
@@ -9,64 +10,49 @@ import streamlit as st
 from pathlib import Path
 from scipy.spatial.transform import Rotation as R
 
-# Explicitly set up rendering backend
-# Add these lines before any pyrender imports
-os.environ['PYOPENGL_PLATFORM'] = 'egl'
-import OpenGL.GL as gl
+# Ensure temporary directories
+os.makedirs('temp', exist_ok=True)
+os.makedirs('output', exist_ok=True)
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 
-class AdvancedModelProcessor:
-    """Advanced processor for 3D models with enhanced animation capabilities"""
+class ModelProcessor:
+    """Comprehensive model processing utilities"""
     
     @staticmethod
-    def detect_model_type(file_path):
-        """Detect the type of 3D model"""
-        file_ext = Path(file_path).suffix.lower()
-        rigged_extensions = ['.fbx', '.dae', '.blend']
-        static_extensions = ['.obj', '.stl', '.ply', '.glb', '.gltf']
-        
-        if file_ext in rigged_extensions:
-            return 'rigged'
-        elif file_ext in static_extensions:
-            return 'static'
-        else:
-            return 'unknown'
-    
-    @staticmethod
-    def analyze_model(file_path):
-        """Analyze the 3D model's characteristics"""
+    def validate_model(file_path):
+        """Validate the uploaded 3D model file"""
         try:
-            # Load the mesh
-            mesh = trimesh.load(file_path)
+            # Check file size
+            file_size = os.path.getsize(file_path)
+            if file_size == 0:
+                raise ValueError("Empty file uploaded")
             
-            # Ensure we have a trimesh object
-            if isinstance(mesh, trimesh.Scene):
-                # If it's a scene, combine meshes
-                mesh = trimesh.util.concatenate([
-                    trimesh.Trimesh(vertices=geom.vertices, faces=geom.faces) 
-                    for geom in mesh.geometry.values()
-                ])
+            # Try loading the model
+            try:
+                mesh = trimesh.load(file_path)
+            except Exception as e:
+                raise ValueError(f"Invalid 3D model file: {str(e)}")
             
-            # Basic model information
-            info = {
-                'vertices': len(mesh.vertices),
-                'faces': len(mesh.faces),
-                'is_watertight': mesh.is_watertight,
-                'volume': mesh.volume,
-                'centroid': mesh.centroid.tolist() if hasattr(mesh, 'centroid') else None,
-                'bounding_box': mesh.bounding_box.extents.tolist() if hasattr(mesh, 'bounding_box') else None
-            }
+            # Check mesh validity
+            if isinstance(mesh, trimesh.Scene):
+                # For scene, ensure it has geometry
+                if not mesh.geometry:
+                    raise ValueError("Model scene contains no valid geometry")
+            else:
+                # For direct mesh, check vertices and faces
+                if len(mesh.vertices) == 0 or len(mesh.faces) == 0:
+                    raise ValueError("Model has no vertices or faces")
             
-            return info
+            return True
         except Exception as e:
-            logging.error(f"Model analysis error: {str(e)}")
-            return None
-    
+            st.error(f"Model validation error: {str(e)}")
+            return False
+
     @staticmethod
     def prepare_model(mesh):
-        """Prepare model for rendering by centering and scaling"""
+        """Prepare model for rendering"""
         # Ensure we have a trimesh object
         if isinstance(mesh, trimesh.Scene):
             # Combine meshes if it's a scene
@@ -84,150 +70,126 @@ class AdvancedModelProcessor:
         
         return mesh
 
-class AnimationPatterns:
-    """Provides different animation patterns"""
+class AnimationGenerator:
+    """Handles 3D model animation generation"""
     
     @staticmethod
-    def rotate_y(frame_num, total_frames):
-        """Rotate around Y-axis"""
-        angle = (frame_num / total_frames) * 2 * np.pi
-        rotation = R.from_rotvec(angle * np.array([0, 1, 0])).as_matrix()
-        transform = np.eye(4)
-        transform[:3, :3] = rotation
-        return transform
-
-def create_pyrender_scene(mesh):
-    """Create a pyrender scene with the model"""
-    try:
-        # Create scene with advanced lighting
-        scene = pyrender.Scene(
-            ambient_light=np.array([0.4, 0.4, 0.4, 1.0]),
-            bg_color=np.array([0.1, 0.1, 0.1, 1.0])
-        )
-        
-        # Convert trimesh to pyrender mesh
-        py_mesh = pyrender.Mesh.from_trimesh(mesh, smooth=True)
-        
-        # Add mesh to scene
-        scene.add(py_mesh)
-        
-        # Advanced camera setup
-        camera = pyrender.PerspectiveCamera(
-            yfov=np.pi / 3.0,
-            aspectRatio=1.0,
-            znear=0.1,
-            zfar=10.0
-        )
-        
-        # Camera positioning
-        camera_pose = np.array([
-            [1.0, 0.0, 0.0, 0.0],
-            [0.0, 1.0, 0.0, 0.0],
-            [0.0, 0.0, 1.0, 2.5],  # Adjusted camera distance
-            [0.0, 0.0, 0.0, 1.0]
-        ])
-        scene.add(camera, pose=camera_pose)
-        
-        # Enhanced lighting setup
-        directional_light = pyrender.DirectionalLight(
-            color=np.array([1.0, 1.0, 1.0]),
-            intensity=2.0
-        )
-        scene.add(directional_light, pose=camera_pose)
-        
-        return scene
-    except Exception as e:
-        logging.error(f"Scene creation error: {str(e)}")
-        raise
-
-def render_frame(scene, width=800, height=600):
-    """Render a high-quality frame"""
-    try:
-        # Create renderer with EGL backend
-        renderer = pyrender.OffscreenRenderer(
-            viewport_width=width,
-            viewport_height=height,
-            point_size=1.0
-        )
-        
-        # Render the scene
-        color, _ = renderer.render(scene, flags=pyrender.RenderFlags.RGBA)
-        renderer.delete()
-        return color
-    except Exception as e:
-        logging.error(f"Frame rendering error: {str(e)}")
-        # Print more detailed error information
-        logging.error(f"Detailed error: {sys.exc_info()}")
-        raise
+    def create_scene(mesh):
+        """Create a pyrender scene with the model"""
+        try:
+            # Create scene with advanced lighting
+            scene = pyrender.Scene(
+                ambient_light=np.array([0.4, 0.4, 0.4, 1.0]),
+                bg_color=np.array([0.1, 0.1, 0.1, 1.0])
+            )
+            
+            # Convert trimesh to pyrender mesh
+            py_mesh = pyrender.Mesh.from_trimesh(mesh, smooth=True)
+            
+            # Add mesh to scene
+            scene.add(py_mesh)
+            
+            # Camera setup
+            camera = pyrender.PerspectiveCamera(
+                yfov=np.pi / 3.0,
+                aspectRatio=1.0,
+                znear=0.1,
+                zfar=10.0
+            )
+            
+            # Camera positioning
+            camera_pose = np.array([
+                [1.0, 0.0, 0.0, 0.0],
+                [0.0, 1.0, 0.0, 0.0],
+                [0.0, 0.0, 1.0, 2.5],
+                [0.0, 0.0, 0.0, 1.0]
+            ])
+            scene.add(camera, pose=camera_pose)
+            
+            # Lighting
+            directional_light = pyrender.DirectionalLight(
+                color=np.array([1.0, 1.0, 1.0]),
+                intensity=2.0
+            )
+            scene.add(directional_light, pose=camera_pose)
+            
+            return scene
+        except Exception as e:
+            st.error(f"Scene creation error: {str(e)}")
+            return None
 
-def generate_animation(model_path, settings):
-    """Generate an animation video"""
-    try:
-        # Load and prepare the model
-        model_processor = AdvancedModelProcessor()
-        
-        # Load the mesh
-        mesh = trimesh.load(model_path)
-        
-        # Prepare the mesh
-        prepared_mesh = model_processor.prepare_model(mesh)
-        
-        # Ensure output directory exists
-        os.makedirs('output', exist_ok=True)
-        
-        # Prepare video writer
-        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-        out = cv2.VideoWriter('output/animation.mp4', 
-                               fourcc, 
-                               settings['fps'], 
-                               (settings['width'], settings['height']))
-        
-        # Animation generation
-        total_frames = settings['duration'] * settings['fps']
-        progress_bar = st.progress(0)
-        status_text = st.empty()
-        
-        for frame_num in range(total_frames):
-            # Get transformation
-            transform = AnimationPatterns.rotate_y(frame_num, total_frames)
+    @staticmethod
+    def generate_animation(model_path, settings):
+        """Generate animation from 3D model"""
+        try:
+            # Load the mesh
+            mesh = trimesh.load(model_path)
             
-            # Create scene
-            scene = create_pyrender_scene(prepared_mesh)
+            # Prepare the mesh
+            prepared_mesh = ModelProcessor.prepare_model(mesh)
             
-            # Render frame
-            frame = render_frame(scene, settings['width'], settings['height'])
+            # Prepare video writer
+            output_path = 'output/3d_model_animation.mp4'
+            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+            out = cv2.VideoWriter(output_path, 
+                                   fourcc, 
+                                   settings['fps'], 
+                                   (settings['width'], settings['height']))
             
-            # Convert and write frame
-            processed_frame = cv2.cvtColor(frame, cv2.COLOR_RGBA2BGR)
-            out.write(processed_frame)
+            # Animation generation
+            total_frames = settings['duration'] * settings['fps']
+            progress_bar = st.progress(0)
+            status_text = st.empty()
             
-            # Update progress
-            progress = (frame_num + 1) / total_frames
-            progress_bar.progress(progress)
-            status_text.text(f"Generating frame {frame_num + 1}/{total_frames}")
+            for frame_num in range(total_frames):
+                # Rotation transformation
+                angle = (frame_num / total_frames) * 2 * np.pi
+                rotation = R.from_rotvec(angle * np.array([0, 1, 0])).as_matrix()
+                transform = np.eye(4)
+                transform[:3, :3] = rotation
+                
+                # Create scene
+                scene = AnimationGenerator.create_scene(prepared_mesh)
+                
+                if scene is None:
+                    raise ValueError("Failed to create rendering scene")
+                
+                # Render frame
+                renderer = pyrender.OffscreenRenderer(
+                    viewport_width=settings['width'],
+                    viewport_height=settings['height']
+                )
+                color, _ = renderer.render(scene)
+                renderer.delete()
+                
+                # Convert and write frame
+                processed_frame = cv2.cvtColor(color, cv2.COLOR_RGBA2BGR)
+                out.write(processed_frame)
+                
+                # Update progress
+                progress = (frame_num + 1) / total_frames
+                progress_bar.progress(progress)
+                status_text.text(f"Generating frame {frame_num + 1}/{total_frames}")
+            
+            out.release()
+            status_text.text("Animation complete!")
+            return output_path
         
-        out.release()
-        status_text.text("Animation complete!")
-        return 'output/animation.mp4'
-    
-    except Exception as e:
-        logging.error(f"Animation generation error: {str(e)}")
-        # Print more detailed error information
-        logging.error(f"Detailed error: {sys.exc_info()}")
-        raise
+        except Exception as e:
+            st.error(f"Animation generation error: {str(e)}")
+            logging.error(f"Detailed error: {traceback.format_exc()}")
+            return None
 
 def main():
     st.set_page_config(page_title="3D Model Animation Generator", layout="wide")
     st.title("3D Model Animation Generator")
     
-    # Create necessary directories
-    for dir_name in ['temp', 'output']:
-        os.makedirs(dir_name, exist_ok=True)
-    
     # File upload section
+    st.write("Upload a 3D Model (obj, stl, ply, fbx, glb, gltf, dae)")
     uploaded_file = st.file_uploader(
-        "Upload 3D Model", 
-        type=['obj', 'stl', 'ply', 'fbx', 'glb', 'gltf', 'dae']
+        "Choose a file", 
+        type=['obj', 'stl', 'ply', 'fbx', 'glb', 'gltf', 'dae'],
+        accept_multiple_files=False
     )
     
     # Animation settings
@@ -242,42 +204,42 @@ def main():
     # Generate animation button
     if st.sidebar.button("Generate Animation"):
         if uploaded_file is not None:
-            # Save uploaded file
-            file_extension = Path(uploaded_file.name).suffix
-            model_path = f"temp/model{file_extension}"
-            
-            with open(model_path, "wb") as f:
-                f.write(uploaded_file.getbuffer())
-            
-            # Generate animation
             try:
+                # Create a temporary file
+                with tempfile.NamedTemporaryFile(delete=False, suffix=Path(uploaded_file.name).suffix) as temp_file:
+                    temp_file.write(uploaded_file.getvalue())
+                    temp_file_path = temp_file.name
+                
+                # Validate the model
+                if not ModelProcessor.validate_model(temp_file_path):
+                    st.error("Invalid model file. Please check the file and try again.")
+                    os.unlink(temp_file_path)
+                    return
+                
+                # Generate animation
                 with st.spinner("Generating Animation..."):
-                    video_path = generate_animation(model_path, settings)
-                    
-                    # Display and download video
-                    with open(video_path, 'rb') as f:
-                        video_bytes = f.read()
-                    
-                    st.video(video_bytes)
-                    st.download_button(
-                        label="Download Animation",
-                        data=video_bytes,
-                        file_name="3d_model_animation.mp4",
-                        mime="video/mp4"
-                    )
+                    video_path = AnimationGenerator.generate_animation(temp_file_path, settings)
                     
-                    # Show model analysis
-                    model_processor = AdvancedModelProcessor()
-                    model_info = model_processor.analyze_model(model_path)
+                    # Clean up temporary file
+                    os.unlink(temp_file_path)
                     
-                    st.subheader("Model Analysis")
-                    st.json(model_info)
+                    if video_path and os.path.exists(video_path):
+                        # Display and download video
+                        with open(video_path, 'rb') as f:
+                            video_bytes = f.read()
+                        
+                        st.video(video_bytes)
+                        st.download_button(
+                            label="Download Animation",
+                            data=video_bytes,
+                            file_name="3d_model_animation.mp4",
+                            mime="video/mp4"
+                        )
             
             except Exception as e:
-                st.error(f"Animation generation error: {str(e)}")
-                logging.error(f"Animation generation error: {str(e)}")
+                st.error(f"Error processing file: {str(e)}")
         else:
-            st.error("Please upload a 3D model")
+            st.warning("Please upload a 3D model file")
 
 if __name__ == "__main__":
     main()