app.py

import os
import sys
import logging
import tempfile
import numpy as np
import cv2
import trimesh
import pyrender
import streamlit as st
from pathlib import Path
from scipy.spatial.transform import Rotation as R

# Ensure temporary directories
os.makedirs('temp', exist_ok=True)
os.makedirs('output', exist_ok=True)

# Configure logging
logging.basicConfig(level=logging.INFO)

class ModelProcessor:
    """Comprehensive model processing utilities"""
    
    @staticmethod
    def validate_model(file_path):
        """Validate the uploaded 3D model file"""
        try:
            # Check file size
            file_size = os.path.getsize(file_path)
            if file_size == 0:
                raise ValueError("Empty file uploaded")
            
            # Try loading the model
            try:
                mesh = trimesh.load(file_path)
            except Exception as e:
                raise ValueError(f"Invalid 3D model file: {str(e)}")
            
            # 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 True
        except Exception as e:
            st.error(f"Model validation error: {str(e)}")
            return False

    @staticmethod
    def prepare_model(mesh):
        """Prepare model for rendering"""
        # Ensure we have a trimesh object
        if isinstance(mesh, trimesh.Scene):
            # Combine meshes if it's a scene
            mesh = trimesh.util.concatenate([
                trimesh.Trimesh(vertices=geom.vertices, faces=geom.faces) 
                for geom in mesh.geometry.values()
            ])
        
        # Center the mesh
        mesh.vertices -= mesh.vertices.mean(axis=0)
        
        # Scale to fit in unit sphere
        scale = np.max(np.abs(mesh.vertices))
        mesh.vertices /= scale
        
        return mesh

class AnimationGenerator:
    """Handles 3D model animation generation"""
    
    @staticmethod
    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

    @staticmethod
    def generate_animation(model_path, settings):
        """Generate animation from 3D model"""
        try:
            # Load the mesh
            mesh = trimesh.load(model_path)
            
            # Prepare the mesh
            prepared_mesh = ModelProcessor.prepare_model(mesh)
            
            # 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']))
            
            # Animation generation
            total_frames = settings['duration'] * settings['fps']
            progress_bar = st.progress(0)
            status_text = st.empty()
            
            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
        
        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")
    
    # File upload section
    st.write("Upload a 3D Model (obj, stl, ply, fbx, glb, gltf, dae)")
    uploaded_file = st.file_uploader(
        "Choose a file", 
        type=['obj', 'stl', 'ply', 'fbx', 'glb', 'gltf', 'dae'],
        accept_multiple_files=False
    )
    
    # Animation settings
    st.sidebar.header("Animation Settings")
    settings = {
        'duration': st.sidebar.slider("Animation Duration (seconds)", 1, 30, 5),
        'fps': st.sidebar.slider("Frames per Second", 15, 60, 30),
        'width': st.sidebar.select_slider("Video Width", [480, 720, 1080], 720),
        'height': st.sidebar.select_slider("Video Height", [480, 720, 1080], 720),
    }
    
    # Generate animation button
    if st.sidebar.button("Generate Animation"):
        if uploaded_file is not None:
            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 = AnimationGenerator.generate_animation(temp_file_path, settings)
                    
                    # Clean up temporary file
                    os.unlink(temp_file_path)
                    
                    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"Error processing file: {str(e)}")
        else:
            st.warning("Please upload a 3D model file")

if __name__ == "__main__":
    main()