Support mp4s

app.py

@@ -1,19 +1,58 @@
 import gradio as gr
 from PIL import Image, ImageDraw, ImageFont, ImageSequence
 import numpy as np
+import cv2
+import os
+import tempfile
 
-def create_grid(image_file, grid_x, grid_y, font_size, font_color, position, border_size, border_color):
-    image_file.file.seek(0)  # Seek to the beginning of the file
+global stored_frames
 
-    # Open the image from bytes
-    img = Image.open(image_file.name)
+def load_and_store_frames(image_file):
+    global stored_frames
 
-    total_frames = img.n_frames
-    selected_frames_count = grid_x * grid_y
+    # Make sure file exists
+    if image_file is None:
+        return "File not found", ""
+
+    print("Loading frames for {}".format(image_file.name))
+
+    if image_file.name.endswith('.mp4'):
+        frames = extract_frames_from_video(image_file.name)
+        video_path = image_file.name
+    else:  # it's a gif
+        img = Image.open(image_file.name)
+        frames = []
+        for i in range(0, img.n_frames):
+            img.seek(i)
+            frames.append(img.copy())
+        # Convert GIF to MP4 for preview
+        fourcc = cv2.VideoWriter_fourcc(*'mp4v') 
+        tmp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
+        video_path = tmp_file.name
+        convert_gif_to_video(image_file.name, tmp_file.name, 1 / (img.info.get('duration', 100) / 1000.0))
+
+    stored_frames = frames  # Store the frames for later use
+    return "Frames loaded successfully", video_path
+
+def generate_grid(grid_x, grid_y, font_size, font_color, position, border_size, border_color):
+    global stored_frames
+    print(f"Processing grid with {grid_x} x {grid_y} grid size, font size {font_size}, font color {font_color}, position {position}, border size {border_size}, border color {border_color}")
+
+    if stored_frames is None:
+        load_and_store_frames()
+    
 
+    grid_img, output_info = create_grid(stored_frames, grid_x, grid_y, font_size, font_color, position, border_size, border_color)
+    details = f"Total Frames: {len(stored_frames)}\n\n{output_info}"
+    return grid_img, details
+
+def create_grid(frames, grid_x, grid_y, font_size, font_color, position, border_size, border_color):
+    total_frames = len(frames)
+    selected_frames_count = grid_x * grid_y
+    
     # Select evenly spaced frames
     selected_frames_indices = np.linspace(0, total_frames - 1, selected_frames_count).astype(int)
-    selected_frames = [img.seek(i) or img.copy() for i in selected_frames_indices]
+    selected_frames = [frames[i] for i in selected_frames_indices]
 
     # Modify frames by adding border and number
     modified_frames = []
@@ -25,9 +64,9 @@ def create_grid(image_file, grid_x, grid_y, font_size, font_color, position, bor
 
     positions = {
         "Top Left": (20, 20),
-        "Top Right": (img.width - 20 - font_size, 20),
-        "Bottom Left": (20, img.height - 20 - font_size),
-        "Bottom Right": (img.width - 20 - font_size, img.height - 20 - font_size)
+        "Top Right": (frames[0].width - 20 - font_size, 20),
+        "Bottom Left": (20, frames[0].height - 20 - font_size),
+        "Bottom Right": (frames[0].width - 20 - font_size, frames[0].height - 20 - font_size)
     }
 
     for i, frame in enumerate(selected_frames):
@@ -55,26 +94,95 @@ def create_grid(image_file, grid_x, grid_y, font_size, font_color, position, bor
 
     output_info = f"Grid size: {grid_x} x {grid_y}\n\nSelected Frames: {selected_frames_count} / {total_frames} ({selected_frames_count / total_frames * 100:.2f}%)"
     return grid_img, output_info
-def gif_info(image_file, grid_x, grid_y, font_size, font_color, position, border_size, border_color):
-    image_file.file.seek(0)
-    img = Image.open(image_file.name)
-    total_frames = img.n_frames
-    frame_rate = img.info.get('duration', 100) / 1000.0  # Convert to seconds
 
-    # Generate grid using create_grid function
-    grid_img, output_info = create_grid(image_file, grid_x, grid_y, font_size, font_color, position, border_size, border_color)
+def extract_frames_from_video(video_file):
+    """Extract frames from an MP4 video."""
+    frames = []
+    cap = cv2.VideoCapture(video_file)
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+        # Convert BGR format (used by OpenCV) to RGB
+        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        frames.append(Image.fromarray(frame_rgb))
+    cap.release()
+    return frames
+
+
+def convert_gif_to_video(gif_path, output_video_path, frame_rate):
+    # Load the gif
+    gif = Image.open(gif_path)
+    # Define the codec and create VideoWriter object
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    out = cv2.VideoWriter(output_video_path, fourcc, frame_rate, (gif.width, gif.height))
+    
+    # Iterate over the frames of the gif
+    for frame_index in range(gif.n_frames):
+        gif.seek(frame_index)
+        # Convert the PIL Image to an array
+        frame_arr = np.array(gif.convert("RGB"))
+        # Convert RGB to BGR format
+        frame_bgr = cv2.cvtColor(frame_arr, cv2.COLOR_RGB2BGR)
+        # Write the frame to the video
+        out.write(frame_bgr)
     
-    gif_details = f"![image](/file={image_file.name})\n\n**Total Frames:** {total_frames}\n\n**Frame Rate:** {frame_rate} sec/frame\n\n{output_info}"
+    out.release()
+
+def gif_or_video_info(image_file, grid_x, grid_y, font_size, font_color, position, border_size, border_color):
+    image_file.file.seek(0)
+    video_path = ""
+    
+    if image_file.name.endswith('.mp4'):
+        video_path = image_file.name
+        cap = cv2.VideoCapture(image_file.name)
+        frame_rate = cap.get(cv2.CAP_PROP_FPS)  # Get the actual frame rate of the video
+        frames = extract_frames_from_video(image_file.name)
+        total_frames = len(frames)
+        cap.release()
+    else:  # it's a gif
+        img = Image.open(image_file.name)
+        frames = []
+        for i in range(0, img.n_frames):
+            img.seek(i)
+            frames.append(img.copy())
+
+        total_frames = img.n_frames
+        frame_rate = 1 / (img.info.get('duration', 100) / 1000.0)  # Convert to seconds
+        
+        # Convert GIF to MP4 and save it to a temp path
+        fourcc = cv2.VideoWriter_fourcc(*'mp4v') 
+        tmp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
+        video_path = tmp_file.name
+        convert_gif_to_video(image_file.name, tmp_file.name, frame_rate)
 
-    return grid_img, gif_details
+    grid_img, output_info = create_grid(frames, grid_x, grid_y, font_size, font_color, position, border_size, border_color)
+    details = f"**Total Frames:** {total_frames}\n\n**Frame Rate:** {frame_rate} frames/sec\n\n{output_info}"
+
+    return grid_img, details, video_path
+
+def gif_info(image_file, grid_x, grid_y, font_size, font_color, position, border_size, border_color):
+    return gif_or_video_info(image_file, grid_x, grid_y, font_size, font_color, position, border_size, border_color)
+
+def mirror(x):
+    return x
 
 with gr.Blocks() as app:
-    gr.Markdown('## GIF Grid Generator')
-    gr.Markdown('Upload a GIF to generate a grid from its frames. Use the sliders to adjust the grid size and text settings.')
+    gr.Markdown('## vid2grid Generator')
+    gr.Markdown('Upload a GIF or MP4 to generate a grid from its frames. Use the sliders to adjust the grid size and text settings.\n\nThis is particularly useful for use with multi modal models such as GPT-4V to retrieve descriptions of short videos or gifs, [more example here.](https://twitter.com/zachnagengast/status/1712896232170180651)\n\n **Note:** The grid will be generated only after clicking the "Generate Grid" button.')
     with gr.Row():
         with gr.Column():
-            control_image = gr.File(label="Upload a GIF", type="file", elem_id="file_upload", file_types=[".gif"])
-            gif_details = gr.Markdown("Upload a gif")
+            control_image = gr.File(label="Upload a short MP4 or GIF", type="file", elem_id="file_upload", file_types=[".gif", ".mp4"])
+            video_preview = gr.Video(interactive=False, label="Preview", format="mp4")
+            gif_details = gr.Markdown("No file found.")
+            # gr.Examples(
+            #     examples=[os.path.join(os.path.dirname(__file__), "demo.mp4")],
+            #     inputs=[control_image],
+            #     outputs=[gif_details, video_preview],
+            #     fn=load_and_store_frames,
+            #     cache_examples=True,
+            # )
+            process_button = gr.Button("Generate Grid")  # New button to trigger the heavy computation
             grid_x_slider = gr.Slider(minimum=1, maximum=10, step=1, value=3, label="Grid X Size")
             grid_y_slider = gr.Slider(minimum=1, maximum=10, step=1, value=3, label="Grid Y Size")
             font_color_dropdown = gr.Dropdown(choices=["Black", "White", "Red", "Green", "Blue"], value="White", label="Numbering Color")
@@ -86,14 +194,18 @@ with gr.Blocks() as app:
             result_image = gr.Image(label="Generated Grid")
 
     # Use .change() method to listen for changes in any of the controls
-    control_image.change(gif_info, inputs=[control_image, grid_x_slider, grid_y_slider, font_size_slider, font_color_dropdown, position_radio, border_size_slider, border_color_dropdown], outputs=[result_image, gif_details])
-    grid_x_slider.change(gif_info, inputs=[control_image, grid_x_slider, grid_y_slider, font_size_slider, font_color_dropdown, position_radio, border_size_slider, border_color_dropdown], outputs=[result_image, gif_details])
-    grid_y_slider.change(gif_info, inputs=[control_image, grid_x_slider, grid_y_slider, font_size_slider, font_color_dropdown, position_radio, border_size_slider, border_color_dropdown], outputs=[result_image, gif_details])
-    font_size_slider.change(gif_info, inputs=[control_image, grid_x_slider, grid_y_slider, font_size_slider, font_color_dropdown, position_radio, border_size_slider, border_color_dropdown], outputs=[result_image, gif_details])
-    font_color_dropdown.change(gif_info, inputs=[control_image, grid_x_slider, grid_y_slider, font_size_slider, font_color_dropdown, position_radio, border_size_slider, border_color_dropdown], outputs=[result_image, gif_details])
-    position_radio.change(gif_info, inputs=[control_image, grid_x_slider, grid_y_slider, font_size_slider, font_color_dropdown, position_radio, border_size_slider, border_color_dropdown], outputs=[result_image, gif_details])
-    border_size_slider.change(gif_info, inputs=[control_image, grid_x_slider, grid_y_slider, font_size_slider, font_color_dropdown, position_radio, border_size_slider, border_color_dropdown], outputs=[result_image, gif_details])
-    border_color_dropdown.change(gif_info, inputs=[control_image, grid_x_slider, grid_y_slider, font_size_slider, font_color_dropdown, position_radio, border_size_slider, border_color_dropdown], outputs=[result_image, gif_details])
+    control_image.upload(load_and_store_frames, inputs=[control_image], outputs=[gif_details, video_preview])
+
+    # grid_x_slider.change(generate_grid, inputs=[grid_x_slider, grid_y_slider, font_size_slider, font_color_dropdown, position_radio, border_size_slider, border_color_dropdown], outputs=[result_image, gif_details, video_preview])
+    # grid_y_slider.change(generate_grid, inputs=[grid_x_slider, grid_y_slider, font_size_slider, font_color_dropdown, position_radio, border_size_slider, border_color_dropdown], outputs=[result_image, gif_details])
+    # font_size_slider.change(generate_grid, inputs=[grid_x_slider, grid_y_slider, font_size_slider, font_color_dropdown, position_radio, border_size_slider, border_color_dropdown], outputs=[result_image, gif_details])
+    # font_color_dropdown.change(generate_grid, inputs=[grid_x_slider, grid_y_slider, font_size_slider, font_color_dropdown, position_radio, border_size_slider, border_color_dropdown], outputs=[result_image, gif_details])
+    # position_radio.change(generate_grid, inputs=[grid_x_slider, grid_y_slider, font_size_slider, font_color_dropdown, position_radio, border_size_slider, border_color_dropdown], outputs=[result_image, gif_details])
+    # border_size_slider.change(generate_grid, inputs=[grid_x_slider, grid_y_slider, font_size_slider, font_color_dropdown, position_radio, border_size_slider, border_color_dropdown], outputs=[result_image, gif_details])
+    # border_color_dropdown.change(generate_grid, inputs=[grid_x_slider, grid_y_slider, font_size_slider, font_color_dropdown, position_radio, border_size_slider, border_color_dropdown], outputs=[result_image, gif_details])
+    
+    process_button.click(generate_grid, inputs=[grid_x_slider, grid_y_slider, font_size_slider, font_color_dropdown, position_radio, border_size_slider, border_color_dropdown], outputs=[result_image, gif_details])
 
 if __name__ == "__main__":
+    stored_frames = None
     app.launch()