import gradio as gr
import cv2
import os
import zipfile

def extract_frames(video_path):
    output_folder = "extracted_frames"
    os.makedirs(output_folder, exist_ok=True)
    
    cap = cv2.VideoCapture(video_path)
    fps = cap.get(cv2.CAP_PROP_FPS)
    
    if fps == 0 or fps is None:
        return "Error: Could not determine FPS of the video."
    
    extracted_images = []
    frame_count = 0
    saved_frame_count = 1  # Start naming frames from 1
    
    frame_interval = max(1, int(fps / 60))  # Avoid division by zero
    
    while True:
        success, frame = cap.read()
        if not success:
            break
        
        if frame_count % frame_interval == 0:  # Extract 60 frames per second
            image_path = os.path.join(output_folder, f"{saved_frame_count}.jpg")
            cv2.imwrite(image_path, frame)
            extracted_images.append(image_path)
            saved_frame_count += 1  # Increment file name sequentially
        
        frame_count += 1
    
    cap.release()
    
    if not extracted_images:
        return "Error: No frames extracted."
    
    zip_filename = "extracted_frames.zip"
    with zipfile.ZipFile(zip_filename, 'w') as zipf:
        for img in extracted_images:
            zipf.write(img, os.path.basename(img))
    
    return zip_filename

gui = gr.Interface(
    fn=extract_frames,
    inputs=gr.Video(label="Upload Video"),
    outputs=gr.File(label="Download Extracted Frames as ZIP"),
    title="Video Frame Extractor",
    description="Upload a video to extract frames at 60 FPS and download as a ZIP file."
)

gui.launch()