app.py

import os
import time
import glob
import json
import numpy as np
import trimesh
import argparse
from scipy.spatial.transform import Rotation
import PIL.Image
from PIL import Image
import math
import trimesh.transformations as tf
from trimesh.exchange.gltf import export_glb

os.environ['PYOPENGL_PLATFORM'] = 'egl'

import gradio as gr
import spaces

def parse_args():
    parser = argparse.ArgumentParser(description='Create animations for 3D models')
    
    parser.add_argument(
        '--input',
        type=str,
        default='./data/demo_glb/',
        help='Input file or directory path (default: ./data/demo_glb/)'
    )
    
    parser.add_argument(
        '--log_path',
        type=str,
        default='./results/demo',
        help='Output directory path (default: results/demo)'
    )
    
    parser.add_argument(
        '--animation_type',
        type=str,
        default='rotate',
        choices=['rotate', 'float', 'explode', 'assemble', 'pulse', 'swing'],
        help='Type of animation to apply'
    )
    
    parser.add_argument(
        '--animation_duration',
        type=float,
        default=3.0,
        help='Duration of animation in seconds'
    )
    
    parser.add_argument(
        '--fps',
        type=int,
        default=30,
        help='Frames per second for animation'
    )
    
    return parser.parse_args()

def get_input_files(input_path):
    if os.path.isfile(input_path):
        return [input_path]
    elif os.path.isdir(input_path):
        return glob.glob(os.path.join(input_path, '*'))
    else:
        raise ValueError(f"Input path {input_path} is neither a file nor a directory")

args = parse_args()

LOG_PATH = args.log_path
os.makedirs(LOG_PATH, exist_ok=True)

print(f"Output directory: {LOG_PATH}")

def normalize_mesh(mesh):
    """Normalize mesh to fit in a unit cube centered at origin"""
    if isinstance(mesh, trimesh.Scene):
        # Scene 객체 처리
        # 씬에서 모든 메시 추출
        meshes = []
        for geometry in mesh.geometry.values():
            if isinstance(geometry, trimesh.Trimesh):
                meshes.append(geometry)
        
        if not meshes:
            raise ValueError("No meshes found in scene")
        
        # 모든 메시의 정점을 결합하여 경계 상자 계산
        all_vertices = np.vstack([m.vertices for m in meshes])
        bounds = np.array([all_vertices.min(axis=0), all_vertices.max(axis=0)])
        center = (bounds[0] + bounds[1]) / 2
        scale = 1.0 / (bounds[1] - bounds[0]).max()
        
        # 각 메시를 정규화하여 새 씬 생성
        normalized_scene = trimesh.Scene()
        for mesh_obj in meshes:
            normalized_mesh = mesh_obj.copy()
            normalized_mesh.vertices = (normalized_mesh.vertices - center) * scale
            normalized_scene.add_geometry(normalized_mesh)
        
        return normalized_scene, center, scale
    else:
        # 일반 Trimesh 객체 처리
        vertices = mesh.vertices
        bounds = np.array([vertices.min(axis=0), vertices.max(axis=0)])
        center = (bounds[0] + bounds[1]) / 2
        scale = 1.0 / (bounds[1] - bounds[0]).max()
        
        # 복사본 생성하여 원본 변경 방지
        normalized_mesh = mesh.copy()
        normalized_mesh.vertices = (vertices - center) * scale
        
        return normalized_mesh, center, scale

def create_rotation_animation(mesh, duration=3.0, fps=30):
    """Create a rotation animation around the Y axis"""
    num_frames = int(duration * fps)
    frames = []
    
    # Normalize the mesh for consistent animation
    mesh, original_center, original_scale = normalize_mesh(mesh)
    
    for frame_idx in range(num_frames):
        t = frame_idx / (num_frames - 1)  # Normalized time [0, 1]
        angle = t * 2 * math.pi  # Full rotation
        
        if isinstance(mesh, trimesh.Scene):
            # Scene 객체인 경우 각 메시에 회전 적용
            frame_scene = trimesh.Scene()
            
            for node_name, transform, geometry_name in mesh.graph.nodes_geometry:
                # 메시를 복사하고 회전 적용
                mesh_copy = mesh.geometry[geometry_name].copy()
                
                # 메시의 중심점 계산
                center_point = mesh_copy.centroid if hasattr(mesh_copy, 'centroid') else np.zeros(3)
                
                # 회전 행렬 생성
                rotation_matrix = tf.rotation_matrix(angle, [0, 1, 0], center_point)
                
                # 변환 적용
                mesh_copy.apply_transform(rotation_matrix)
                
                # 씬에 추가
                frame_scene.add_geometry(mesh_copy, node_name=node_name)
            
            frames.append(frame_scene)
        else:
            # 일반 Trimesh 객체인 경우 직접 회전 적용
            animated_mesh = mesh.copy()
            rotation_matrix = tf.rotation_matrix(angle, [0, 1, 0])
            animated_mesh.apply_transform(rotation_matrix)
            frames.append(animated_mesh)
    
    return frames

def create_float_animation(mesh, duration=3.0, fps=30, amplitude=0.2):
    """Create a floating animation where the mesh moves up and down"""
    num_frames = int(duration * fps)
    frames = []
    
    # Normalize the mesh for consistent animation
    mesh, original_center, original_scale = normalize_mesh(mesh)
    
    for frame_idx in range(num_frames):
        t = frame_idx / (num_frames - 1)  # Normalized time [0, 1]
        y_offset = amplitude * math.sin(2 * math.pi * t)
        
        if isinstance(mesh, trimesh.Scene):
            # Scene 객체인 경우
            frame_scene = trimesh.Scene()
            
            for node_name, transform, geometry_name in mesh.graph.nodes_geometry:
                # 메시를 복사
                mesh_copy = mesh.geometry[geometry_name].copy()
                
                # 변환 적용
                translation_matrix = tf.translation_matrix([0, y_offset, 0])
                mesh_copy.apply_transform(translation_matrix)
                
                # 씬에 추가
                frame_scene.add_geometry(mesh_copy, node_name=node_name)
            
            frames.append(frame_scene)
        else:
            # 일반 Trimesh 객체인 경우
            animated_mesh = mesh.copy()
            translation_matrix = tf.translation_matrix([0, y_offset, 0])
            animated_mesh.apply_transform(translation_matrix)
            frames.append(animated_mesh)
    
    return frames

def create_explode_animation(mesh, duration=3.0, fps=30):
    """Create an explode animation where parts of the mesh move outward"""
    num_frames = int(duration * fps)
    frames = []
    
    # Normalize the mesh for consistent animation
    mesh, original_center, original_scale = normalize_mesh(mesh)
    
    # 씬인 경우 부분별 처리
    if isinstance(mesh, trimesh.Scene):
        for frame_idx in range(num_frames):
            t = frame_idx / (num_frames - 1)  # Normalized time [0, 1]
            
            # 새 씬 생성
            frame_scene = trimesh.Scene()
            
            # 각 노드별 폭발 애니메이션 적용
            for node_name, transform, geometry_name in mesh.graph.nodes_geometry:
                mesh_copy = mesh.geometry[geometry_name].copy()
                
                # 노드의 중심점 계산
                center_point = mesh_copy.centroid if hasattr(mesh_copy, 'centroid') else np.zeros(3)
                
                # 방향 벡터 (원점에서 객체 중심까지)
                direction = center_point
                if np.linalg.norm(direction) < 1e-10:
                    # 중심점이 원점과 너무 가까우면 랜덤 방향 선택
                    direction = np.random.rand(3) - 0.5
                
                direction = direction / np.linalg.norm(direction)
                
                # 폭발 이동 적용
                translation = direction * t * 0.5  # 폭발 강도 조절
                translation_matrix = tf.translation_matrix(translation)
                mesh_copy.apply_transform(translation_matrix)
                
                # 씬에 추가
                frame_scene.add_geometry(mesh_copy, node_name=f"{node_name}_{frame_idx}")
            
            frames.append(frame_scene)
    else:
        # 메시를 여러 부분으로 분할
        try:
            components = mesh.split(only_watertight=False)
            if len(components) <= 1:
                # 컴포넌트가 하나뿐이면 정점 기반 폭발 애니메이션 사용
                raise ValueError("Mesh cannot be split into components")
        except:
            # 분할이 실패하면 정점 기반 폭발 애니메이션 사용
            for frame_idx in range(num_frames):
                t = frame_idx / (num_frames - 1)  # Normalized time [0, 1]
                
                # 메시 복사
                animated_mesh = mesh.copy()
                vertices = animated_mesh.vertices.copy()
                
                # 각 정점을 중심에서 바깥쪽으로 이동
                directions = vertices.copy()
                norms = np.linalg.norm(directions, axis=1, keepdims=True)
                mask = norms > 1e-10
                directions[mask] = directions[mask] / norms[mask]
                directions[~mask] = np.random.rand(np.sum(~mask), 3) - 0.5
                
                # 폭발 계수 적용
                vertices += directions * t * 0.3
                animated_mesh.vertices = vertices
                
                frames.append(animated_mesh)
        else:
            # 컴포넌트 기반 폭발 애니메이션
            for frame_idx in range(num_frames):
                t = frame_idx / (num_frames - 1)  # Normalized time [0, 1]
                
                # 새 씬 생성
                scene = trimesh.Scene()
                
                # 각 컴포넌트를 중심에서 바깥쪽으로 이동
                for i, component in enumerate(components):
                    # 컴포넌트 복사
                    animated_component = component.copy()
                    
                    # 컴포넌트 중심점에서 방향 계산
                    direction = animated_component.centroid
                    if np.linalg.norm(direction) < 1e-10:
                        # 중심점이 원점과 너무 가까우면 랜덤 방향 선택
                        direction = np.random.rand(3) - 0.5
                    
                    direction = direction / np.linalg.norm(direction)
                    
                    # 폭발 이동 적용
                    translation = direction * t * 0.5
                    translation_matrix = tf.translation_matrix(translation)
                    animated_component.apply_transform(translation_matrix)
                    
                    # 씬에 추가
                    scene.add_geometry(animated_component, node_name=f"component_{i}")
                
                # 씬을 단일 메시로 변환 (근사치)
                animated_mesh = trimesh.util.concatenate(scene.dump())
                frames.append(animated_mesh)
    
    return frames

def create_assemble_animation(mesh, duration=3.0, fps=30):
    """Create an assembly animation (reverse of explode)"""
    # Get explode animation and reverse it
    explode_frames = create_explode_animation(mesh, duration, fps)
    return list(reversed(explode_frames))

def create_pulse_animation(mesh, duration=3.0, fps=30, min_scale=0.8, max_scale=1.2):
    """Create a pulsing animation where the mesh scales up and down"""
    num_frames = int(duration * fps)
    frames = []
    
    # Normalize the mesh for consistent animation
    mesh, original_center, original_scale = normalize_mesh(mesh)
    
    for frame_idx in range(num_frames):
        t = frame_idx / (num_frames - 1)  # Normalized time [0, 1]
        
        # Create a copy of the mesh to animate
        animated_mesh = mesh.copy()
        
        # Apply pulsing motion (sinusoidal scale)
        scale_factor = min_scale + (max_scale - min_scale) * (0.5 + 0.5 * math.sin(2 * math.pi * t))
        scale_matrix = tf.scale_matrix(scale_factor)
        animated_mesh.apply_transform(scale_matrix)
        
        # Add to frames
        frames.append(animated_mesh)
    
    return frames

def create_swing_animation(mesh, duration=3.0, fps=30, max_angle=math.pi/6):
    """Create a swinging animation where the mesh rotates back and forth"""
    num_frames = int(duration * fps)
    frames = []
    
    # Normalize the mesh for consistent animation
    mesh, original_center, original_scale = normalize_mesh(mesh)
    
    for frame_idx in range(num_frames):
        t = frame_idx / (num_frames - 1)  # Normalized time [0, 1]
        
        # Create a copy of the mesh to animate
        animated_mesh = mesh.copy()
        
        # Apply swinging motion (sinusoidal rotation)
        angle = max_angle * math.sin(2 * math.pi * t)
        rotation_matrix = tf.rotation_matrix(angle, [0, 1, 0])
        animated_mesh.apply_transform(rotation_matrix)
        
        # Add to frames
        frames.append(animated_mesh)
    
    return frames

def generate_gif_from_frames(frames, output_path, fps=30, resolution=(640, 480), background_color=(255, 255, 255, 255)):
    """Generate a GIF from animation frames"""
    gif_frames = []
    
    for frame in frames:
        # Create a scene with the frame
        if not isinstance(frame, trimesh.Scene):
            scene = trimesh.Scene(frame)
        else:
            scene = frame
        
        # 자동 카메라 설정 (전체 장면이 보이도록)
        try:
            scene.camera
        except:
            # 기본 카메라가 없으면 생성
            scene.set_camera()
        
        # 모든 객체가 보이도록 카메라 위치 조정
        scene.camera_transform = scene.camera_transform

        # Render the frame
        try:
            rendered_img = scene.save_image(resolution=resolution, background=background_color)
            pil_img = Image.open(rendered_img)
            gif_frames.append(pil_img)
        except Exception as e:
            print(f"Error rendering frame: {str(e)}")
            # 렌더링 실패 시 빈 이미지 생성
            gif_frames.append(Image.new('RGB', resolution, (255, 255, 255)))
    
    # Save as GIF
    if gif_frames:
        gif_frames[0].save(
            output_path,
            save_all=True,
            append_images=gif_frames[1:],
            optimize=False,
            duration=int(1000 / fps),
            loop=0
        )
        return output_path
    else:
        return None

def create_animation_mesh(input_mesh_path, animation_type='rotate', duration=3.0, fps=30):
    """Create animation from input mesh based on animation type"""
    # Load the mesh
    try:
        # 먼저 Scene으로 로드 시도
        loaded_obj = trimesh.load(input_mesh_path)
        
        # Scene인 경우 단일 메시로 변환 시도
        if isinstance(loaded_obj, trimesh.Scene):
            print("Loaded a scene, extracting meshes...")
            # 씬에서 모든 메시 추출
            meshes = []
            for geometry in loaded_obj.geometry.values():
                if isinstance(geometry, trimesh.Trimesh):
                    meshes.append(geometry)
            
            if not meshes:
                print("No meshes found in scene")
                return None, None
                
            # 모든 메시를 하나로 결합
            mesh = trimesh.util.concatenate(meshes)
        elif isinstance(loaded_obj, trimesh.Trimesh):
            mesh = loaded_obj
        else:
            print(f"Unsupported object type: {type(loaded_obj)}")
            return None, None
            
    except Exception as e:
        print(f"Error loading mesh: {str(e)}")
        return None, None
    
    # Generate animation frames based on animation type
    if animation_type == 'rotate':
        frames = create_rotation_animation(mesh, duration, fps)
    elif animation_type == 'float':
        frames = create_float_animation(mesh, duration, fps)
    elif animation_type == 'explode':
        frames = create_explode_animation(mesh, duration, fps)
    elif animation_type == 'assemble':
        frames = create_assemble_animation(mesh, duration, fps)
    elif animation_type == 'pulse':
        frames = create_pulse_animation(mesh, duration, fps)
    elif animation_type == 'swing':
        frames = create_swing_animation(mesh, duration, fps)
    else:
        print(f"Unknown animation type: {animation_type}")
        return None, None
    
    base_filename = os.path.basename(input_mesh_path).rsplit('.', 1)[0]
    
    # Save animated mesh as GLB
    try:
        animated_glb_path = os.path.join(LOG_PATH, f'animated_{base_filename}.glb')
        
        # For GLB output, we'll use the first frame for now
        # In a production environment, you'd want to use proper animation keyframes
        if frames and len(frames) > 0:
            # First frame for static GLB
            first_frame = frames[0]
            # Export as GLB
            scene = trimesh.Scene(first_frame)
            scene.export(animated_glb_path)
        else:
            return None, None
    except Exception as e:
        print(f"Error exporting GLB: {str(e)}")
        animated_glb_path = None
    
    # Create GIF for preview
    try:
        animated_gif_path = os.path.join(LOG_PATH, f'animated_{base_filename}.gif')
        generate_gif_from_frames(frames, animated_gif_path, fps)
    except Exception as e:
        print(f"Error creating GIF: {str(e)}")
        animated_gif_path = None
    
    return animated_glb_path, animated_gif_path

@spaces.GPU
def process_3d_model(input_3d, animation_type, animation_duration, fps):
    """Process a 3D model and apply animation"""
    print(f"Processing: {input_3d} with animation type: {animation_type}")
    
    try:
        # Create animation
        animated_glb_path, animated_gif_path = create_animation_mesh(
            input_3d,
            animation_type=animation_type,
            duration=animation_duration,
            fps=fps
        )
        
        if not animated_glb_path or not animated_gif_path:
            return "Error creating animation", None
        
        # Create a simple JSON metadata file
        metadata = {
            "animation_type": animation_type,
            "duration": animation_duration,
            "fps": fps,
            "original_model": os.path.basename(input_3d),
            "created_at": time.strftime("%Y-%m-%d %H:%M:%S")
        }
        
        json_path = os.path.join(LOG_PATH, f'metadata_{os.path.basename(input_3d).rsplit(".", 1)[0]}.json')
        with open(json_path, 'w') as f:
            json.dump(metadata, f, indent=4)
        
        return animated_glb_path, animated_gif_path, json_path
    except Exception as e:
        error_msg = f"Error processing file: {str(e)}"
        print(error_msg)
        return error_msg, None, None

_HEADER_ = '''
<h2><b>GLB 애니메이션 생성기 - 3D 모델 움직임 효과</b></h2>

이 데모를 통해 정적인 3D 모델(GLB 파일)에 다양한 애니메이션 효과를 적용할 수 있습니다.

❗️❗️❗️**중요사항:**
- 이 데모는 업로드된 GLB 파일에 애니메이션을 적용합니다.
- 다양한 애니메이션 스타일 중에서 선택하세요: 회전, 부유, 폭발, 조립, 펄스, 스윙.
- 결과는 애니메이션된 GLB 파일과 미리보기용 GIF 파일로 제공됩니다.
'''

_INFO_ = r"""
### 애니메이션 유형 설명
- **회전(rotate)**: 모델이 Y축을 중심으로 회전합니다.
- **부유(float)**: 모델이 위아래로 부드럽게 떠다닙니다.
- **폭발(explode)**: 모델의 각 부분이 중심에서 바깥쪽으로 퍼져나갑니다.
- **조립(assemble)**: 폭발 애니메이션의 반대 - 부품들이 함께 모입니다.
- **펄스(pulse)**: 모델이 크기가 커졌다 작아졌다를 반복합니다.
- **스윙(swing)**: 모델이 좌우로 부드럽게 흔들립니다.

### 팁
- 애니메이션 길이와 FPS를 조절하여 움직임의 속도와 부드러움을 조절할 수 있습니다.
- 복잡한 모델은 처리 시간이 더 오래 걸릴 수 있습니다.
- GIF 미리보기는 빠른 참조용이며, 고품질 결과를 위해서는 애니메이션된 GLB 파일을 다운로드하세요.
"""

# Gradio 인터페이스 설정
def create_gradio_interface():
    with gr.Blocks(title="GLB 애니메이션 생성기") as demo:
        # 제목 섹션
        gr.Markdown(_HEADER_)
        
        with gr.Row():
            with gr.Column():
                # 입력 컴포넌트
                input_3d = gr.Model3D(label="3D 모델 파일 업로드 (GLB 포맷)")
                
                with gr.Row():
                    animation_type = gr.Dropdown(
                        label="애니메이션 유형", 
                        choices=["rotate", "float", "explode", "assemble", "pulse", "swing"],
                        value="rotate"
                    )
                
                with gr.Row():
                    animation_duration = gr.Slider(
                        label="애니메이션 길이 (초)",
                        minimum=1.0,
                        maximum=10.0,
                        value=3.0,
                        step=0.5
                    )
                    fps = gr.Slider(
                        label="초당 프레임 수",
                        minimum=15,
                        maximum=60,
                        value=30,
                        step=1
                    )
                
                submit_btn = gr.Button("모델 처리 및 애니메이션 생성")
                
            with gr.Column():
                # 출력 컴포넌트
                output_3d = gr.Model3D(label="애니메이션 적용된 3D 모델")
                output_gif = gr.Image(label="애니메이션 미리보기 (GIF)")
                output_json = gr.File(label="메타데이터 파일 다운로드")
        
        # 애니메이션 유형 설명
        gr.Markdown(_INFO_)
        
        # 버튼 동작 설정
        submit_btn.click(
            fn=process_3d_model,
            inputs=[input_3d, animation_type, animation_duration, fps],
            outputs=[output_3d, output_gif, output_json]
        )
        
        # 예제 준비
        example_files = [ [f] for f in glob.glob('./data/demo_glb/*.glb') ]
        
        if example_files:
            example = gr.Examples(
                examples=example_files,
                inputs=[input_3d],
                examples_per_page=10,
            )
    
    return demo

# 메인 실행 부분
if __name__ == "__main__":
    demo = create_gradio_interface()
    demo.launch(share=True, server_name="0.0.0.0", server_port=7860)