Update app.py

app.py

@@ -1,81 +1,40 @@
 import gradio as gr
-import os
-import zipfile
-import cv2
-import numpy as np
 from gradio_client import Client, handle_file
 
-# Function to handle video processing
-def process_video(style_image, zip_file):
-    # Extract images from the uploaded zip file
-    zip_path = zip_file.name
-    output_dir = "/tmp/extracted_images"  # Use a temporary directory for Hugging Face Spaces
-    os.makedirs(output_dir, exist_ok=True)
-    
-    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
-        zip_ref.extractall(output_dir)
-
-    # Get all image files from the extracted zip
-    image_files = [os.path.join(output_dir, f) for f in os.listdir(output_dir) if f.endswith(('.png', '.jpg', '.jpeg'))]
-
-    # Initialize the client for image-to-image color transfer
-    client = Client("Abu1998/Image_Color_Transfer_Video")
-
-    # Prepare the style image (color reference)
-    style_image_path = style_image.name
-    style_image_url = handle_file(style_image_path)
-
-    transformed_images = []
-
-    # Process each image and apply color transfer
-    for image_file in image_files:
-        image_url = handle_file(image_file)
-        result = client.predict(img1=style_image_url, img2=image_url, api_name="/_transfer_style")
+# Function to handle video-to-image extraction and return the extracted frames as a ZIP
+def extract_frames_from_video(video_file):
+    # Initialize the client for Video-To-Image API
+    client = Client("Abu1998/Video-To-Image")
 
-        # Save the transformed image in the temporary directory
-        transformed_image_path = os.path.join("/tmp/transformed_images", os.path.basename(image_file))
-        os.makedirs("/tmp/transformed_images", exist_ok=True)
-        with open(transformed_image_path, 'wb') as f:
-            f.write(result[0]['path'].read())
-
-        transformed_images.append(transformed_image_path)
-
-    # Convert transformed images back to video
-    frame_array = []
-    for image_path in transformed_images:
-        img = cv2.imread(image_path)
-        height, width, layers = img.shape
-        size = (width, height)
-        frame_array.append(img)
-
-    video_output_path = "/tmp/output_video.mp4"
-    out = cv2.VideoWriter(video_output_path, cv2.VideoWriter_fourcc(*'mp4v'), 20.0, size)
-
-    for frame in frame_array:
-        out.write(frame)
-
-    out.release()
+    # Upload the video file to the API
+    video_url = handle_file(video_file.name)
+    
+    # Send the video file to the API for processing
+    result = client.predict(
+        video_path={"video": video_url},
+        api_name="/predict"
+    )
 
-    return video_output_path
+    # The result should contain the link to the ZIP file with extracted frames
+    return result['filepath']
 
 
 # Gradio Interface setup
 def create_gradio_interface():
     with gr.Blocks() as app:
-        gr.Markdown("# Image Color Transfer Video Generator")
+        gr.Markdown("# Video to Image Frame Extractor")
 
-        # Input components
-        style_image = gr.File(type="file", label="Upload Style Image (Reference Image for Color Transfer)")  # Updated to gr.File
-        zip_file = gr.File(label="Upload Zip File of Images")
+        # Input component for video upload
+        video_file = gr.File(type="file", label="Upload Video")
 
-        # Output component
-        video_output = gr.Video(label="Generated Video")
+        # Output component for the extracted frames (ZIP file)
+        zip_output = gr.File(label="Download Extracted Frames as ZIP")
 
-        # Button to trigger processing
-        process_button = gr.Button("Generate Video")
+        # Button to trigger video-to-image extraction
+        process_button = gr.Button("Extract Frames")
 
         # Action for button click
-        process_button.click(fn=process_video, inputs=[style_image, zip_file], outputs=video_output)
+        process_button.click(fn=extract_frames_from_video, inputs=video_file, outputs=zip_output)
 
     app.launch(share=True)