Update app.py

app.py

@@ -4,7 +4,7 @@ import math
 import gradio as gr
 from PIL import Image, ImageDraw
 import os
-from utils.keyframe_utils import generate_keyframe_prompt
+from utils.keyframe_utils import generate_keyframe_prompt, generate_all_keyframe_images
 
 # Load segments JSON
 def load_segments():
@@ -18,44 +18,27 @@ def get_keyframe_images(segment_id):
     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))
+            images.append(Image.open(img_path).resize((128, 72)))
         else:
-            images.append(Image.new("RGB", (256, 144), color=(200, 200, 200)))
+            img = Image.new("RGB", (128, 72), color=(200, 200, 200))
+            draw = ImageDraw.Draw(img)
+            draw.text((10, 30), "No Image", fill=(0, 0, 0))
+            images.append(img)
     return images
 
-
 def segment_display(segment):
-    """
-    Format one row for Gradio display:
-    - Truncate and wrap description
-    - Load and resize keyframe images to thumbnail size
-    """
     seg_id = segment.get("segment_id")
     description = segment.get("description", "")
     short_desc = description if len(description) <= 120 else description[:117] + "..."
 
-    # Wrap long text with newlines every ~40 chars
     def wrap_text(text, width=40):
         return "\n".join([text[i:i+width] for i in range(0, len(text), width)])
 
     display_desc = wrap_text(short_desc)
-
     row = [f"Segment {seg_id}", display_desc]
-
-    for i in range(1, 4):
-        img_path = f"keyframes/segment_{seg_id}_v{i}.png"
-        if os.path.exists(img_path):
-            img = Image.open(img_path).resize((128, 72))
-        else:
-            # Placeholder if not found
-            img = Image.new("RGB", (128, 72), color=(200, 200, 200))
-            draw = ImageDraw.Draw(img)
-            draw.text((10, 30), "No Image", fill=(0, 0, 0))
-        row.append(img)
-
+    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)
@@ -68,14 +51,12 @@ def paginate_segments(page=1, page_size=15):
     data = [segment_display(seg) for seg in subset]
     return headers, data, max_page
 
-# Function to return downloadable JSON file
 def download_prompts():
     if os.path.exists("all_prompts_output.json"):
         return "all_prompts_output.json"
     else:
         return None
 
-# Gradio interface
 def build_interface():
     with gr.Blocks() as demo:
         gr.Markdown("## 🎬 Keyframe Candidate Viewer")
@@ -91,15 +72,25 @@ def build_interface():
             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 ➡")
+        def generate_and_refresh():
+            segments = load_segments()
+            generate_all_keyframe_images(segments)
+            return update(1)
+
+        with gr.Row():
+            with gr.Row(elem_id="page-controls"):
+                prev_btn = gr.Button("⬅", size="sm")
+                next_btn = gr.Button("➡", size="sm")
 
-        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])
 
+        gr.Markdown("### 🖼️ 生成所有分镜的关键帧图像")
+        generate_btn = gr.Button("🛠️ 一键生成图像")
+        generate_btn.click(fn=generate_and_refresh, outputs=[table, total_pages_text, page_state])
+
         gr.Markdown("### 📥 下载所有生成的图像提示词")
         with gr.Row():
             download_btn = gr.Button("📥 导出 prompts JSON")
@@ -110,4 +101,4 @@ def build_interface():
 
 if __name__ == "__main__":
     demo = build_interface()
-    demo.launch()
+    demo.launch()