app.py

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

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

import gradio as gr
import spaces

# 결과 저장 경로
LOG_PATH = './results/demo'
os.makedirs(LOG_PATH, exist_ok=True)

def create_simple_rotation_animation(input_glb_path, output_glb_path, num_frames=30):
    """
    원본 GLB 파일에 회전 애니메이션을 적용한 새로운 GLB 파일 생성
    """
    try:
        # GLB 파일 로드
        scene = trimesh.load(input_glb_path)
        
        if isinstance(scene, trimesh.Scene):
            # 애니메이션 적용 (첫 번째 프레임만 사용)
            angle = math.pi / 4  # 45도 회전
            
            # 씬의 모든 메시에 회전 적용
            for node_name, transform, geometry_name in scene.graph.nodes_geometry:
                # 원본 위치 백업
                original_transform = scene.graph[node_name][0]
                
                # 회전 변환 계산
                rotation = tf.rotation_matrix(angle, [0, 1, 0])
                
                # 새 변환 = 원본 변환 * 회전 변환
                new_transform = np.dot(original_transform, rotation)
                
                # 변환 적용
                scene.graph[node_name] = new_transform
            
            # 회전된 GLB 저장
            scene.export(output_glb_path)
            print(f"Saved animated GLB to {output_glb_path}")
            return output_glb_path
            
        elif isinstance(scene, trimesh.Trimesh):
            # 단일 메시인 경우
            new_scene = trimesh.Scene()
            
            # 메시에 회전 적용
            angle = math.pi / 4  # 45도 회전
            rotation = tf.rotation_matrix(angle, [0, 1, 0])
            scene.apply_transform(rotation)
            
            # 회전된 메시를 씬에 추가
            new_scene.add_geometry(scene)
            
            # 씬을 GLB로 저장
            new_scene.export(output_glb_path)
            print(f"Saved animated GLB to {output_glb_path}")
            return output_glb_path
            
        else:
            print(f"Unsupported format: {type(scene)}")
            return None
            
    except Exception as e:
        print(f"Error creating animation: {str(e)}")
        return None

def create_textual_animation_gif(output_path, model_name, animation_type, duration=3.0, fps=30):
    """텍스트 기반의 간단한 애니메이션 GIF 생성 - 렌더링 실패 시 대체용"""
    try:
        # 간단한 프레임 시퀀스 생성
        frames = []
        num_frames = int(duration * fps)
        if num_frames > 60:  # 너무 많은 프레임은 효율적이지 않음
            num_frames = 60
            
        for i in range(num_frames):
            t = i / (num_frames - 1)  # 0~1 범위
            angle = t * 360  # 전체 회전
            
            # 새 이미지 생성
            img = Image.new('RGB', (640, 480), color=(240, 240, 240))
            draw = ImageDraw.Draw(img)
            
            # 정보 텍스트
            draw.text((50, 50), f"Model: {os.path.basename(model_name)}", fill=(0, 0, 0))
            draw.text((50, 100), f"Animation Type: {animation_type}", fill=(0, 0, 0))
            draw.text((50, 150), f"Frame: {i+1}/{num_frames}", fill=(0, 0, 0))
            
            # 애니메이션 유형에 따른 시각적 효과
            center_x, center_y = 320, 240
            if animation_type == 'rotate':
                # 회전하는 사각형
                radius = 100
                x = center_x + radius * math.cos(math.radians(angle))
                y = center_y + radius * math.sin(math.radians(angle))
                draw.rectangle((x-40, y-40, x+40, y+40), outline=(0, 0, 0), fill=(255, 0, 0))
                
            elif animation_type == 'float':
                # 위아래로 움직이는 원
                offset_y = 50 * math.sin(2 * math.pi * t)
                draw.ellipse((center_x-50, center_y-50+offset_y, center_x+50, center_y+50+offset_y), 
                            outline=(0, 0, 0), fill=(0, 0, 255))
                
            elif animation_type == 'explode' or animation_type == 'assemble':
                # 바깥쪽/안쪽으로 움직이는 여러 도형
                scale = t if animation_type == 'explode' else 1 - t
                for j in range(8):
                    angle_j = j * 45
                    dist = 120 * scale
                    x = center_x + dist * math.cos(math.radians(angle_j))
                    y = center_y + dist * math.sin(math.radians(angle_j))
                    
                    if j % 3 == 0:
                        draw.rectangle((x-20, y-20, x+20, y+20), outline=(0, 0, 0), fill=(255, 0, 0))
                    elif j % 3 == 1:
                        draw.ellipse((x-20, y-20, x+20, y+20), outline=(0, 0, 0), fill=(0, 255, 0))
                    else:
                        draw.polygon([(x, y-20), (x+20, y+20), (x-20, y+20)], outline=(0, 0, 0), fill=(0, 0, 255))
                        
            elif animation_type == 'pulse':
                # 크기가 변하는 원
                scale = 0.5 + 0.5 * math.sin(2 * math.pi * t)
                radius = 100 * scale
                draw.ellipse((center_x-radius, center_y-radius, center_x+radius, center_y+radius), 
                           outline=(0, 0, 0), fill=(0, 255, 0))
                
            elif animation_type == 'swing':
                # 좌우로 움직이는 삼각형
                angle_offset = 30 * math.sin(2 * math.pi * t)
                points = [
                    (center_x + 100 * math.cos(math.radians(angle_offset)), center_y - 80),
                    (center_x + 100 * math.cos(math.radians(120 + angle_offset)), center_y + 40),
                    (center_x + 100 * math.cos(math.radians(240 + angle_offset)), center_y + 40)
                ]
                draw.polygon(points, outline=(0, 0, 0), fill=(255, 165, 0))
            
            # 프레임 추가
            frames.append(img)
        
        # GIF로 저장
        frames[0].save(
            output_path,
            save_all=True,
            append_images=frames[1:],
            optimize=False,
            duration=int(1000 / fps),
            loop=0
        )
        print(f"Created textual animation GIF at {output_path}")
        return output_path
    except Exception as e:
        print(f"Error creating textual animation: {str(e)}")
        return None

def create_glb_with_animation(input_glb_path, animation_type, duration=3.0, fps=30):
    """
    업로드된 GLB 파일에 애니메이션을 적용한 새로운 GLB 파일 생성
    """
    try:
        base_filename = os.path.basename(input_glb_path).rsplit('.', 1)[0]
        output_glb_path = os.path.join(LOG_PATH, f"animated_{base_filename}.glb")
        
        # 원본 모델 로드
        scene = trimesh.load(input_glb_path)
        print(f"Loaded GLB: {type(scene)}")
        
        # 애니메이션 유형에 따라 처리
        if animation_type == 'rotate':
            # 시계방향 45도 회전
            angle = math.pi / 4
            axis = [0, 1, 0]  # Y축 (수직축)
        elif animation_type == 'float':
            # 수직으로 움직임
            angle = 0
            axis = [0, 1, 0]
            # Y축 방향으로 이동
            if isinstance(scene, trimesh.Scene):
                for node_name, transform, geometry_name in scene.graph.nodes_geometry:
                    # 가볍게 위로 이동
                    translation = tf.translation_matrix([0, 0.2, 0])
                    original_transform = scene.graph[node_name][0]
                    new_transform = np.dot(original_transform, translation)
                    scene.graph[node_name] = new_transform
            elif isinstance(scene, trimesh.Trimesh):
                translation = tf.translation_matrix([0, 0.2, 0])
                scene.apply_transform(translation)
        elif animation_type == 'explode':
            # 각 부분의 중심에서 약간 멀어지게
            angle = 0
            axis = [0, 1, 0]
            # 오브젝트가 여러 개일 경우 중심에서 바깥쪽으로 이동
            if isinstance(scene, trimesh.Scene):
                # 씬의 중심점 계산
                all_vertices = []
                for geometry_name, geometry in scene.geometry.items():
                    if hasattr(geometry, 'vertices') and len(geometry.vertices) > 0:
                        all_vertices.append(geometry.vertices)
                
                if all_vertices:
                    all_points = np.vstack(all_vertices)
                    center = np.mean(all_points, axis=0)
                    
                    for node_name, transform, geometry_name in scene.graph.nodes_geometry:
                        # 각 부분의 중심점 계산
                        geometry = scene.geometry[geometry_name]
                        if hasattr(geometry, 'centroid'):
                            part_center = geometry.centroid
                            # 중심에서 객체 방향 계산
                            direction = part_center - center
                            if np.linalg.norm(direction) > 0.001:
                                direction = direction / np.linalg.norm(direction)
                                # 방향으로 이동
                                translation = tf.translation_matrix(direction * 0.2)
                                original_transform = scene.graph[node_name][0]
                                new_transform = np.dot(original_transform, translation)
                                scene.graph[node_name] = new_transform
        elif animation_type == 'pulse':
            # 모델을 약간 키움
            scale_factor = 1.2
            
            if isinstance(scene, trimesh.Scene):
                # 씬의 중심점 계산
                all_vertices = []
                for geometry_name, geometry in scene.geometry.items():
                    if hasattr(geometry, 'vertices') and len(geometry.vertices) > 0:
                        all_vertices.append(geometry.vertices)
                
                if all_vertices:
                    all_points = np.vstack(all_vertices)
                    center = np.mean(all_points, axis=0)
                    
                    for node_name, transform, geometry_name in scene.graph.nodes_geometry:
                        # 중심 기준 스케일링
                        translate_to_center = tf.translation_matrix(-center)
                        scale = np.eye(4)
                        scale[:3, :3] *= scale_factor
                        translate_back = tf.translation_matrix(center)
                        
                        # 변환 적용
                        original_transform = scene.graph[node_name][0]
                        new_transform = np.dot(translate_back, np.dot(scale, np.dot(translate_to_center, original_transform)))
                        scene.graph[node_name] = new_transform
            elif isinstance(scene, trimesh.Trimesh):
                # 메시 중심 기준 스케일링
                center = scene.centroid
                translate_to_center = tf.translation_matrix(-center)
                scale = tf.scale_matrix(scale_factor)
                translate_back = tf.translation_matrix(center)
                scene.apply_transform(np.dot(translate_back, np.dot(scale, translate_to_center)))
        else:
            # 기본적으로 회전 적용
            angle = math.pi / 4
            axis = [0, 1, 0]
        
        # 회전 적용 (float, explode, pulse 아닌 경우)
        if angle != 0:
            if isinstance(scene, trimesh.Scene):
                # 씬의 모든 부분에 회전 적용
                for node_name, transform, geometry_name in scene.graph.nodes_geometry:
                    rotation = tf.rotation_matrix(angle, axis)
                    original_transform = scene.graph[node_name][0]
                    new_transform = np.dot(original_transform, rotation)
                    scene.graph[node_name] = new_transform
            elif isinstance(scene, trimesh.Trimesh):
                # 단일 메시에 회전 적용
                rotation = tf.rotation_matrix(angle, axis)
                scene.apply_transform(rotation)
        
        # 결과 저장
        scene.export(output_glb_path)
        print(f"Saved animated GLB to {output_glb_path}")
        return output_glb_path
    except Exception as e:
        print(f"Error creating animated GLB: {str(e)}")
        # 오류 발생 시 원본 파일 복사
        try:
            import shutil
            output_path = os.path.join(LOG_PATH, f"copy_{os.path.basename(input_glb_path)}")
            shutil.copy(input_glb_path, output_path)
            print(f"Copied original GLB to {output_path}")
            return output_path
        except Exception as copy_error:
            print(f"Error copying GLB: {copy_error}")
            return input_glb_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:
        # 1. 업로드된 GLB 파일에 실제 애니메이션 적용
        animated_glb_path = create_glb_with_animation(
            input_3d,
            animation_type,
            animation_duration,
            fps
        )
        
        # 2. 텍스트 기반 애니메이션 GIF 생성 (백업용)
        base_filename = os.path.basename(input_3d).rsplit('.', 1)[0]
        text_gif_path = os.path.join(LOG_PATH, f'text_animated_{base_filename}.gif')
        animated_gif_path = create_textual_animation_gif(
            text_gif_path,
            os.path.basename(input_3d),
            animation_type,
            animation_duration,
            fps
        )
        
        # 3. 메타데이터 생성
        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_{base_filename}.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

# Gradio 인터페이스 설정
with gr.Blocks(title="GLB 애니메이션 생성기") as demo:
    # 제목 섹션
    gr.Markdown("""
    <h2><b>GLB 애니메이션 생성기 - 3D 모델 움직임 효과</b></h2>

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

    ❗️❗️❗️**중요사항:**
    - 이 데모는 업로드된 GLB 파일에 애니메이션을 적용합니다.
    - 다양한 애니메이션 스타일 중에서 선택하세요: 회전, 부유, 폭발, 조립, 펄스, 스윙.
    - 결과는 애니메이션된 GLB 파일과 미리보기용 GIF 파일로 제공됩니다.
    """)
    
    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("""
    ### 애니메이션 유형 설명
    - **회전(rotate)**: 모델이 Y축을 중심으로 회전합니다.
    - **부유(float)**: 모델이 위아래로 부드럽게 떠다닙니다.
    - **폭발(explode)**: 모델의 각 부분이 중심에서 바깥쪽으로 퍼져나갑니다.
    - **조립(assemble)**: 폭발 애니메이션의 반대 - 부품들이 함께 모입니다.
    - **펄스(pulse)**: 모델이 크기가 커졌다 작아졌다를 반복합니다.
    - **스윙(swing)**: 모델이 좌우로 부드럽게 흔들립니다.

    ### 팁
    - 애니메이션 길이와 FPS를 조절하여 움직임의 속도와 부드러움을 조절할 수 있습니다.
    - 복잡한 모델은 처리 시간이 더 오래 걸릴 수 있습니다.
    - GIF 미리보기는 빠른 참조용이며, 고품질 결과를 위해서는 애니메이션된 GLB 파일을 다운로드하세요.
    """)
    
    # 버튼 동작 설정
    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:
        gr.Examples(
            examples=example_files,
            inputs=[input_3d],
            examples_per_page=10,
        )

# 앱 실행
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)