app.py

"""
GLB 애니메이션 생성기
────────────────────────────────────────────────────────
업로드한 GLB 3D 모델에 회전·부유·폭발 등 애니메이션을 적용해
① 변형된 GLB ② 실제 모델을 렌더링한 GIF ③ 메타데이터 JSON 을 돌려줍니다.
• headless 서버 대응: EGL + pyglet.headless + trimesh → 실패 시 pyrender 폴백
• 최대 60 fps·60 프레임로 제한 (데모 목적)
"""

# ──────────────────── 1. 공통 모듈 ────────────────────
import os, io, time, glob, json, math, shutil
import numpy as np
from PIL import Image

import pyglet
pyglet.options["headless"] = True            # 반드시 trimesh 전에!
os.environ["PYOPENGL_PLATFORM"] = "egl"      # Off-screen GL

import trimesh
import trimesh.transformations as tf

import gradio as gr
import spaces

LOG_PATH = "./results/demo"
os.makedirs(LOG_PATH, exist_ok=True)

# ──────────────────── 2. 렌더 유틸 ────────────────────
def _render_with_trimesh(scene: trimesh.Scene, res):
    png = scene.save_image(resolution=res, visible=True)
    if png is None:
        raise RuntimeError("trimesh.save_image returned None")
    return Image.open(io.BytesIO(png)).convert("RGB")

def _render_with_pyrender(mesh_or_scene, res):
    import pyrender  # 지연 임포트
    if isinstance(mesh_or_scene, trimesh.Scene):
        mesh = trimesh.util.concatenate(mesh_or_scene.dump())
    else:
        mesh = mesh_or_scene
    mesh = pyrender.Mesh.from_trimesh(mesh, smooth=False)
    scn = pyrender.Scene()
    scn.add(mesh)
    cam = pyrender.PerspectiveCamera(yfov=np.pi / 3)
    scn.add(cam, pose=tf.translation_matrix([0, 0, 3]))
    light = pyrender.DirectionalLight(intensity=3.0)
    scn.add(light, pose=tf.translation_matrix([0, 5, 5]))
    r = pyrender.OffscreenRenderer(*res)
    color, _ = r.render(scn, flags=pyrender.RenderFlags.RGBA)
    r.delete()
    return Image.fromarray(color[..., :3])

# ──────────────────── 3. 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 생성"""
    base = trimesh.load(input_glb_path)
    if isinstance(base, trimesh.Trimesh):
        base = trimesh.Scene(base)

    num_frames = min(int(duration * fps), 60)
    frames = []
    for i in range(num_frames):
        t = i / (num_frames - 1)
        scene = base.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.apply_transform(M)

        # ① trimesh → ② pyrender 폴백
        try:
            frame = _render_with_trimesh(scene, resolution)
        except Exception as e:
            print("trimesh 렌더 실패, pyrender 폴백:", e)
            frame = _render_with_pyrender(scene, resolution)

        frames.append(frame)

    frames[0].save(
        output_path,
        save_all=True,
        append_images=frames[1:],
        duration=int(1000 / fps),
        loop=0,
    )
    print("GIF saved:", output_path)
    return output_path

# ──────────────────── 4. GLB 변형 ────────────────────
def modify_glb_file(input_glb_path, output_glb_path, animation_type="rotate"):
    """단일 변환을 적용한 새 GLB 저장 (실제 애니메이션은 아님)"""
    try:
        scn = trimesh.load(input_glb_path)
        if not isinstance(scn, trimesh.Scene):
            scn = trimesh.Scene(scn)

        if animation_type == "rotate":
            T = tf.rotation_matrix(math.pi / 4, [0, 1, 0])
        elif animation_type == "float":
            T = tf.translation_matrix([0, 0.5, 0])
        elif animation_type == "pulse":
            T = tf.scale_matrix(1.2)
        elif animation_type == "explode":
            T = tf.translation_matrix([0.5, 0, 0])
        elif animation_type == "assemble":
            T = tf.translation_matrix([-0.5, 0, 0])
        elif animation_type == "swing":
            T = tf.rotation_matrix(math.pi / 8, [0, 0, 1])
        else:
            T = np.eye(4)

        scn.apply_transform(T)
        scn.export(output_glb_path)
        return output_glb_path
    except Exception as e:
        print("GLB 변형 실패, 원본 복사:", e)
        shutil.copy(input_glb_path, output_glb_path)
        return output_glb_path

# ──────────────────── 5. 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]
        glb_out = os.path.join(LOG_PATH, f"animated_{base}.glb")
        gif_out = os.path.join(LOG_PATH, f"preview_{base}.gif")
        json_out = os.path.join(LOG_PATH, f"metadata_{base}.json")

        modify_glb_file(input_3d, glb_out, animation_type)
        create_model_animation_gif(
            gif_out, input_3d, animation_type, animation_duration, fps
        )

        meta = 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_out, "w") as f:
            json.dump(meta, f, indent=4)

        return glb_out, gif_out, json_out

    except Exception as e:
        print("process_3d_model 실패:", e)
        err_gif = os.path.join(LOG_PATH, "error.gif")
        Image.new("RGB", (640, 480), (255, 0, 0)).save(err_gif)
        return input_3d, err_gif, None

# ──────────────────── 6. Gradio UI ────────────────────
with gr.Blocks(title="GLB 애니메이션 생성기") as demo:
    gr.Markdown(
        """
<h2><b>GLB 애니메이션 생성기 - 3D 모델 움직임 효과</b></h2>
업로드한 GLB 모델에 애니메이션을 적용해 변형된 GLB·GIF·JSON을 제공합니다.
"""
    )
    with gr.Row():
        with gr.Column():
            inp = gr.Model3D(label="3D 모델 업로드 (GLB)")
            typ = gr.Dropdown(
                label="애니메이션 유형",
                choices=["rotate", "float", "explode", "assemble", "pulse", "swing"],
                value="rotate",
            )
            dur = gr.Slider("애니메이션 길이 (초)", 1.0, 10.0, 3.0, 0.5)
            fps = gr.Slider("FPS", 15, 60, 30, 1)
            btn = gr.Button("애니메이션 생성")
        with gr.Column():
            out_glb = gr.Model3D(label="애니메이션된 GLB")
            out_gif = gr.Image(label="미리보기 GIF")
            out_json = gr.File(label="메타데이터 JSON")

    btn.click(
        fn=process_3d_model,
        inputs=[inp, typ, dur, fps],
        outputs=[out_glb, out_gif, out_json],
    )

    ex = [[f] for f in glob.glob("./data/demo_glb/*.glb")]
    if ex:
        gr.Examples(examples=ex, inputs=[inp])

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