import cv2
import gradio as gr
import numpy as np
from keras.models import load_model

# Load your trained model
model = load_model('model.h5')  # Ensure this matches the model filename in your Space

def process_frame(frame):
    # Preprocess the frame for the model
    img = cv2.resize(frame, (64, 64))  # Adjust size based on your model input
    img = img.astype('float32') / 255.0  # Normalize
    img = np.expand_dims(img, axis=0)  # Add batch dimension

    # Predict using the model
    prediction = model.predict(img)
    return prediction[0][1]  # Assuming category 1 is jumpscare

def convert_video_to_dumpscare(video_path, sensitivity):
    cap = cv2.VideoCapture(video_path)
    fps = cap.get(cv2.CAP_PROP_FPS)
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    
    output_path = 'output_dumpscare.mp4'
    out = cv2.VideoWriter(output_path, fourcc, fps, (int(cap.get(3)), int(cap.get(4))))

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break

        # Check if the frame is a jumpscare
        if process_frame(frame) > sensitivity:  # Use sensitivity slider
            out.write(frame)  # Write frame if it's a jumpscare

    cap.release()
    out.release()
    return output_path

# Gradio interface
def gradio_interface(video, sensitivity):
    output_video = convert_video_to_dumpscare(video.name, sensitivity)
    return output_video

# Set up Gradio app
iface = gr.Interface(
    fn=gradio_interface,
    inputs=[
        gr.Video(label="Import Video"),
        gr.Slider(minimum=0.0, maximum=1.0, default=0.5, label="Sensitivity"),
    ],
    outputs=gr.Video(label="Output Dumpscare Video"),
    title="Dumpscare Video Converter",
    description="Upload a video, set sensitivity, and click 'Cut' to process the video.",
    live=True
)

iface.launch()