Update app.py

app.py

@@ -1,63 +1,87 @@
-import zipfile
+import gradio as gr
 import os
+import zipfile
 import cv2
+import numpy as np
 from gradio_client import Client, handle_file
-import shutil
 
-# Function to process the video and apply the color transformation to each frame
-def process_video(input_zip_path, style_image_path, output_video_path):
-    # Initialize the Hugging Face Client for the color transformation API
+# Downgrade numpy to version <2 if not already set in requirements
+os.system("pip install numpy<2")
+
+# 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 = "extracted_images"
+    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 zip file containing images
-    with zipfile.ZipFile(input_zip_path, 'r') as zip_ref:
-        zip_ref.extractall('input_images')  # Extract images to 'input_images' folder
+    # Prepare the style image (color reference)
+    style_image_path = style_image.name
+    style_image_url = handle_file(style_image_path)
 
-    # Process each image
-    image_files = [f for f in os.listdir('input_images') if f.endswith('.png') or f.endswith('.jpg')]
-    
-    # Prepare output folder
-    output_folder = 'output_images'
-    os.makedirs(output_folder, exist_ok=True)
+    transformed_images = []
 
-    # Apply color transformation to each image
+    # Process each image and apply color transfer
     for image_file in image_files:
-        input_image_path = os.path.join('input_images', image_file)
-        
-        # Use the style image and the current image for transformation
-        result = client.predict(
-            img1=handle_file(style_image_path),
-            img2=handle_file(input_image_path),
-            api_name="/_transfer_style"
-        )
-        
+        image_url = handle_file(image_file)
+        result = client.predict(img1=style_image_url, img2=image_url, api_name="/_transfer_style")
+
         # Save the transformed image
-        transformed_image_path = os.path.join(output_folder, f"transformed_{image_file}")
+        transformed_image_path = os.path.join("transformed_images", os.path.basename(image_file))
         with open(transformed_image_path, 'wb') as f:
-            f.write(result[0]['data'])
-
-    # Create a video from the transformed images
-    images = [os.path.join(output_folder, f) for f in sorted(os.listdir(output_folder))]
-    frame = cv2.imread(images[0])
-    height, width, layers = frame.shape
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    video_writer = cv2.VideoWriter(output_video_path, fourcc, 30.0, (width, height))
-
-    for image in images:
-        frame = cv2.imread(image)
-        video_writer.write(frame)
-    
-    video_writer.release()
-    
-    # Clean up temporary files
-    shutil.rmtree('input_images')
-    shutil.rmtree('output_images')
+            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 = "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()
+
+    return video_output_path
+
+
+# Gradio Interface setup
+def create_gradio_interface():
+    with gr.Blocks() as app:
+        gr.Markdown("# Image Color Transfer Video Generator")
+
+        # Input components
+        style_image = gr.Image(type="file", label="Upload Style Image (Reference Image for Color Transfer)")
+        zip_file = gr.File(label="Upload Zip File of Images")
+
+        # Output component
+        video_output = gr.Video(label="Generated Video")
+
+        # Button to trigger processing
+        process_button = gr.Button("Generate Video")
+
+        # Action for button click
+        process_button.click(fn=process_video, inputs=[style_image, zip_file], outputs=video_output)
 
-    print(f"Video saved to {output_video_path}")
+    app.launch(share=True)
 
-# Example usage
-input_zip_path = 'path_to_zip_of_images.zip'  # Path to the zip file of images from the video
-style_image_path = 'path_to_style_image.jpg'  # Path to the style image for color transformation
-output_video_path = 'output_video.mp4'  # Path where the final video will be saved
 
-process_video(input_zip_path, style_image_path, output_video_path)
+# Run the application
+if __name__ == "__main__":
+    create_gradio_interface()