app.py

import gradio as gr
import cv2
import numpy as np
import trimesh
import tempfile
import os

# -------------------------
# GLOBAL (checkerboard persistence)
# -------------------------
_checkerboard_colors = None

# -------------------------
# VIDEO LOADING (BGR → RGB FIXED ✅)
# -------------------------
def read_video_frames(video_path, start=0, end=None, frame_step=1):
    cap = cv2.VideoCapture(video_path)
    frames = []

    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    if end is None:
        end = total_frames

    count = 0

    while True:
        ret, frame = cap.read()
        if not ret or count >= end:
            break

        if count >= start and (count - start) % frame_step == 0:
            # FIX COLOR ORDER HERE
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frames.append(frame)

        count += 1

    cap.release()
    return np.array(frames)


# -------------------------
# DOWNSAMPLING
# -------------------------
def downsample_frames(frames, block_size=1, method='stride'):
    if block_size == 1:
        return frames

    z, h, w, c = frames.shape

    if method == 'stride':
        return frames[:, ::block_size, ::block_size]

    elif method == 'mean':
        new_h = h // block_size
        new_w = w // block_size
        out = np.zeros((z, new_h, new_w, c), dtype=np.uint8)

        for zi in range(z):
            for i in range(new_h):
                for j in range(new_w):
                    block = frames[
                        zi,
                        i*block_size:(i+1)*block_size,
                        j*block_size:(j+1)*block_size
                    ]
                    out[zi, i, j] = block.mean(axis=(0,1))
        return out


# -------------------------
# VOXEL MASK
# -------------------------
def frames_to_voxels(frames, threshold=10):
    return (np.sum(frames, axis=3) > threshold)


# -------------------------
# VOXEL → MESH (FIXED COLORS ✅)
# -------------------------
def voxels_to_mesh(frames, voxels, voxel_size=1.0):
    meshes = []
    z_len, h, w = voxels.shape

    for z in range(z_len):
        for y in range(h):
            for x in range(w):
                if voxels[z, y, x]:
                    color = frames[z, frames.shape[1] - 1 - y, x].astype(np.uint8)

                    cube = trimesh.creation.box(extents=[voxel_size]*3)
                    cube.apply_translation([x, y, z])

                    # Apply colors correctly (RGBA uint8)
                    rgba = np.append(color, 255)
                    cube.visual.face_colors = np.tile(rgba, (12,1))

                    meshes.append(cube)

    if meshes:
        return trimesh.util.concatenate(meshes)
    return trimesh.Scene()


# -------------------------
# RANDOM CHECKERBOARD (ONE-TIME COLORS ✅)
# -------------------------
def default_checkerboard():
    global _checkerboard_colors

    h, w, z_len = 10, 10, 2
    frames = np.zeros((z_len, h, w, 3), dtype=np.uint8)

    if _checkerboard_colors is None:
        _checkerboard_colors = np.random.randint(
            0, 256, size=(z_len, h, w, 3), dtype=np.uint8
        )

    for z in range(z_len):
        for y in range(h):
            for x in range(w):
                if (x + y + z) % 2 == 0:
                    frames[z, y, x] = [0, 0, 0]
                else:
                    frames[z, y, x] = _checkerboard_colors[z, y, x]

    voxels = frames_to_voxels(frames, threshold=1)
    mesh = voxels_to_mesh(frames, voxels, voxel_size=2)

    tmp = tempfile.gettempdir()
    obj = os.path.join(tmp, "checkerboard.obj")
    glb = os.path.join(tmp, "checkerboard.glb")

    mesh.export(obj)
    mesh.export(glb)

    return obj, glb, glb


# -------------------------
# MAIN GENERATOR
# -------------------------
def generate_voxel_files(
    video_file,
    start_frame,
    end_frame,
    frame_step,
    block_size,
    downsample_method
):
    if video_file is None:
        return default_checkerboard()

    frames = read_video_frames(
        video_file.name,
        start=start_frame,
        end=end_frame,
        frame_step=frame_step
    )

    frames = downsample_frames(
        frames,
        block_size=block_size,
        method=downsample_method
    )

    voxels = frames_to_voxels(frames)
    mesh = voxels_to_mesh(frames, voxels)

    tmp = tempfile.gettempdir()
    obj = os.path.join(tmp, "output.obj")
    glb = os.path.join(tmp, "output.glb")

    mesh.export(obj)
    mesh.export(glb)

    return obj, glb, glb


# -------------------------
# GRADIO UI
# -------------------------
iface = gr.Interface(
    fn=generate_voxel_files,
    inputs=[
        gr.File(label="Upload MP4 (or leave empty for checkerboard)"),
        gr.Slider(0, 500, value=0, step=1, label="Start Frame"),
        gr.Slider(0, 500, value=50, step=1, label="End Frame"),
        gr.Slider(1, 10, value=1, step=1, label="Frame Step"),
        gr.Slider(1, 32, value=1, step=1, label="Pixel Block Size"),
        gr.Radio(["stride", "mean"], value="stride", label="Downsample Method"),
    ],
    outputs=[
        gr.File(label="OBJ"),
        gr.File(label="GLB"),
        gr.Model3D(label="3D Preview"),
    ],
    title="MP4 → Voxels → 3D",
    description="If no file is uploaded, a random-color checkerboard appears."
)

if __name__ == "__main__":
    iface.launch()