app.py

"""
GLB 애니메이션 생성기 – 업로드한 3D 모델을 직접 렌더링해 GIF·GLB로 내보냅니다.
(Gradio + Trimesh 오프스크린 렌더링)

※ 주요 변경
1. create_model_animation_gif : 모델을 실제 렌더링해 프레임 생성
2. process_3d_model 내부에서 새 GIF 함수 사용
3. 예외 처리·주석 일부 보강
"""

import os, io, time, glob, json, math
import numpy as np
import trimesh
from PIL import Image
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)


# ------------------------------------------------------------------
# 1) 실제 GLB 모델을 렌더링해 GIF 생성
# ------------------------------------------------------------------
def create_model_animation_gif(
    output_path: str,
    input_glb_path: str,
    animation_type: str,
    duration: float = 3.0,
    fps: int = 30,
    resolution=(640, 480),
):
    """
    업로드된 GLB 모델을 오프스크린으로 렌더링해 애니메이션 GIF 생성
    """
    # (1) 모델 로드
    scene = trimesh.load(input_glb_path)
    if isinstance(scene, trimesh.Trimesh):
        scene = trimesh.Scene(scene)

    frames, num_frames = [], min(int(duration * fps), 60)
    for i in range(num_frames):
        t = i / (num_frames - 1)

        # (2) 프레임마다 변환 적용
        scene_i = scene.copy()
        if animation_type == "rotate":
            M = tf.rotation_matrix(2 * math.pi * t, [0, 1, 0])
        elif animation_type == "float":
            M = tf.translation_matrix([0, 0.5 * math.sin(2 * math.pi * t), 0])
        elif animation_type == "pulse":
            M = tf.scale_matrix(0.8 + 0.4 * math.sin(2 * math.pi * t))
        elif animation_type == "explode":
            M = tf.translation_matrix([0.5 * t, 0, 0])
        elif animation_type == "assemble":
            M = tf.translation_matrix([0.5 * (1 - t), 0, 0])
        elif animation_type == "swing":
            M = tf.rotation_matrix(math.pi / 6 * math.sin(2 * math.pi * t), [0, 0, 1])
        else:
            M = np.eye(4)

        scene_i.apply_transform(M)

        # (3) 오프스크린 렌더 → PNG 바이트
        png_bytes = scene_i.save_image(resolution=resolution, visible=True)
        frame = Image.open(io.BytesIO(png_bytes)).convert("RGB")
        frames.append(frame)

    # (4) GIF 저장
    frames[0].save(
        output_path,
        save_all=True,
        append_images=frames[1:],
        optimize=False,
        duration=int(1000 / fps),
        loop=0,
    )
    print(f"Created model animation GIF at {output_path}")
    return output_path


# ------------------------------------------------------------------
# 2) GLB 파일에 단순 변환 적용해 저장
# ------------------------------------------------------------------
def modify_glb_file(input_glb_path, output_glb_path, animation_type="rotate"):
    """
    업로드된 GLB 파일에 단일 변환(회전·이동·스케일)을 적용해 새 GLB로 저장
    """
    try:
        scene = trimesh.load(input_glb_path)
        if not isinstance(scene, trimesh.Scene):
            scene = trimesh.Scene(scene)

        # 변환 매트릭스 결정
        if animation_type == "rotate":
            transform = tf.rotation_matrix(math.pi / 4, [0, 1, 0])
        elif animation_type == "float":
            transform = tf.translation_matrix([0, 0.5, 0])
        elif animation_type == "pulse":
            transform = tf.scale_matrix(1.2)
        elif animation_type == "explode":
            transform = tf.translation_matrix([0.5, 0, 0])
        elif animation_type == "assemble":
            transform = tf.translation_matrix([-0.5, 0, 0])
        elif animation_type == "swing":
            transform = tf.rotation_matrix(math.pi / 8, [0, 0, 1])
        else:
            transform = np.eye(4)

        scene.apply_transform(transform)
        scene.export(output_glb_path)
        print(f"Exported modified GLB to {output_glb_path}")
        return output_glb_path

    except Exception as e:
        print("Error modifying GLB:", e)
        # 실패 시 원본 복사
        import shutil

        shutil.copy(input_glb_path, output_glb_path)
        print(f"Copied original GLB to {output_glb_path}")
        return output_glb_path


# ------------------------------------------------------------------
# 3) Gradio에서 호출할 메인 파이프라인
# ------------------------------------------------------------------
@spaces.GPU
def process_3d_model(input_3d, animation_type, animation_duration, fps):
    """
    ① GLB 변환 → ② GIF 렌더 → ③ 메타데이터 JSON 저장
    """
    try:
        base = os.path.splitext(os.path.basename(input_3d))[0]
        animated_glb = os.path.join(LOG_PATH, f"animated_{base}.glb")
        animated_gif = os.path.join(LOG_PATH, f"preview_{base}.gif")
        json_path = os.path.join(LOG_PATH, f"metadata_{base}.json")

        # 1) GLB 변형
        modify_glb_file(input_3d, animated_glb, animation_type)

        # 2) 실제 모델 기반 GIF
        create_model_animation_gif(
            animated_gif,
            input_3d,
            animation_type,
            animation_duration,
            fps,
        )

        # 3) 메타데이터 기록
        metadata = dict(
            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"),
        )
        with open(json_path, "w") as f:
            json.dump(metadata, f, indent=4)

        return animated_glb, animated_gif, json_path

    except Exception as e:
        # 치명적 오류 시: 원본 모델 + 빈 GIF
        print("Error in process_3d_model:", e)
        fallback_gif = os.path.join(LOG_PATH, "error.gif")
        Image.new("RGB", (640, 480), (255, 0, 0)).save(fallback_gif)
        return input_3d, fallback_gif, None


# ------------------------------------------------------------------
# 4) Gradio UI
# ------------------------------------------------------------------
with gr.Blocks(title="GLB 애니메이션 생성기") as demo:
    gr.Markdown(
        """
<h2><b>GLB 애니메이션 생성기 - 3D 모델 움직임 효과</b></h2>
정적인 3D 모델(GLB)에 회전·부유·폭발 등 애니메이션을 적용합니다.
"""
    )

    with gr.Row():
        with gr.Column():
            input_3d = gr.Model3D(label="3D 모델 업로드 (GLB)")
            animation_type = gr.Dropdown(
                label="애니메이션 유형",
                choices=["rotate", "float", "explode", "assemble", "pulse", "swing"],
                value="rotate",
            )
            animation_duration = gr.Slider(
                label="애니메이션 길이 (초)", minimum=1.0, maximum=10.0, value=3.0, step=0.5
            )
            fps = gr.Slider(label="FPS", minimum=15, maximum=60, value=30, step=1)
            submit_btn = gr.Button("애니메이션 생성")

        with gr.Column():
            output_3d = gr.Model3D(label="애니메이션된 GLB")
            output_gif = gr.Image(label="미리보기 GIF")
            output_json = gr.File(label="메타데이터 JSON")

    submit_btn.click(
        fn=process_3d_model,
        inputs=[input_3d, animation_type, animation_duration, fps],
        outputs=[output_3d, output_gif, output_json],
    )

    # 예제 GLB 폴더가 있다면 자동 등록
    example_glbs = [[f] for f in glob.glob("./data/demo_glb/*.glb")]
    if example_glbs:
        gr.Examples(examples=example_glbs, inputs=[input_3d])

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