app.py

# --- app.py ---
import json
import math
import gradio as gr
from PIL import Image
import os
from utils.keyframe_utils import generate_keyframe_prompt

# Load segments JSON
def load_segments():
    with open("segments_full.json", "r", encoding="utf-8") as f:
        segments = json.load(f)
    return segments

# Load pre-generated keyframe images (3 per segment)
def get_keyframe_images(segment_id):
    images = []
    for i in range(1, 4):
        img_path = f"keyframes/segment_{segment_id}_v{i}.png"
        if os.path.exists(img_path):
            images.append(Image.open(img_path))
        else:
            images.append(Image.new("RGB", (256, 144), color=(200, 200, 200)))
    return images

# UI component for a single segment row
def segment_display(segment):
    description = segment["description"]
    seg_id = segment["segment_id"]
    row = [f"Segment {seg_id}", description]
    row.extend(get_keyframe_images(seg_id))
    return row

# Pagination logic
def paginate_segments(page=1, page_size=15):
    segments = load_segments()
    total = len(segments)
    max_page = math.ceil(total / page_size)
    start = (page - 1) * page_size
    end = start + page_size
    subset = segments[start:end]

    headers = ["Segment ID", "Description", "Candidate 1", "Candidate 2", "Candidate 3"]
    data = [segment_display(seg) for seg in subset]
    return headers, data, max_page

# Gradio interface
def build_interface():
    with gr.Blocks() as demo:
        gr.Markdown("## 🎬 Keyframe Candidate Viewer")
        page_state = gr.State(1)
        table = gr.Dataframe(headers=[], datatype=["str", "str", "image", "image", "image"], row_count=15)
        total_pages_text = gr.Textbox(label="Page Info", interactive=False)

        def update(page):
            headers, rows, max_page = paginate_segments(page)
            return gr.update(headers=headers, value=rows), f"Page {page} of {max_page}", page

        prev_btn = gr.Button("⬅ Prev")
        next_btn = gr.Button("Next ➡")

        table.change(fn=lambda: None, inputs=[], outputs=[])
        prev_btn.click(fn=lambda p: max(1, p - 1), inputs=page_state, outputs=page_state).then(update, inputs=page_state, outputs=[table, total_pages_text, page_state])
        next_btn.click(fn=lambda p: p + 1, inputs=page_state, outputs=page_state).then(update, inputs=page_state, outputs=[table, total_pages_text, page_state])

        demo.load(fn=update, inputs=page_state, outputs=[table, total_pages_text, page_state])

    return demo

if __name__ == "__main__":
    demo = build_interface()
    demo.launch()