import cv2
import numpy as np

import gradio as gr
import os
import shutil
import spaces

class PanoramaGenerator:
    def __init__(self, video_path):
        self.video_path = video_path
        self.frames = []

    def extract_frames(self, frame_interval=10):
        cap = cv2.VideoCapture(self.video_path)

        if not cap.isOpened():
            raise ValueError("Error opening video file")

        frame_count = 0
        while True:
            ret, frame = cap.read()

            if not ret:
                break

            # Extract frame
            if frame_count % frame_interval == 0:
                # Resize frame
                resized_frame = cv2.resize(frame, (0, 0), fx=0.5, fy=0.5)
                self.frames.append(resized_frame)

            frame_count += 1

        cap.release()
        return self.frames

    def create_panorama(self):
        if len(self.frames) < 2:
            raise ValueError("Not enough frames to create panorama")

        stitcher = cv2.Stitcher_create(mode=cv2.Stitcher_PANORAMA)

        # Attempt to stitch frames
        status, panorama = stitcher.stitch(self.frames)

        # Check if stitching was successful
        if status != cv2.Stitcher_OK:
            raise RuntimeError(f"Stitching failed with error code {status}")

        return panorama

    def save_panorama(self, output_path="panorama.jpg"):
        panorama = self.create_panorama()
        cv2.imwrite(output_path, panorama)
        return output_path




def process_video(video_file):
    # Save the uploaded video to a temporary file
    temp_video_path = "temp_video.mp4"
    shutil.copy(video_file.name, temp_video_path)

    # Generate the panorama
    panorama_gen = PanoramaGenerator(temp_video_path)
    panorama_gen.extract_frames()
    panorama_path = panorama_gen.save_panorama()

    # Clean up the temporary video file
    os.remove(temp_video_path)

    return panorama_path

@spaces.GPU
def generate_and_display(video_file):
    panorama_path = process_video(video_file)
    return panorama_path, panorama_path

with gr.Blocks() as demo:
    gr.Markdown("# Video to Panorama Generator")

    with gr.Row():
        with gr.Column():
            video_input = gr.File(label="Upload Video", file_types=["video"])
            generate_button = gr.Button("Generate Panorama")
        with gr.Column():
            panorama_output = gr.Image(label="Generated Panorama")
            download_button = gr.DownloadButton("Download Panorama", file_name="panorama.jpg")

    generate_button.click(
        fn=generate_and_display,
        inputs=video_input,
        outputs=[panorama_output, download_button]
    )

demo.launch()