import gradio as gr
import torch
import numpy as np
import cv2
from PIL import Image
from transformers import AutoImageProcessor, AutoModelForVideoClassification, ViTForImageClassification

video_processor = AutoImageProcessor.from_pretrained("facebook/timesformer-base-finetuned-k400")
video_model = AutoModelForVideoClassification.from_pretrained("facebook/timesformer-base-finetuned-k400")

image_processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224")
image_model = ViTForImageClassification.from_pretrained("google/vit-base-patch16-224")

def extract_frames(video_path, num_frames=8):
    cap = cv2.VideoCapture(video_path)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    frames = []
    if total_frames == 0:
        return frames
    frame_indices = np.linspace(0, total_frames - 1, num_frames).astype(int)
    for idx in frame_indices:
        cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
        ret, frame = cap.read()
        if not ret:
            break
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        frame = cv2.resize(frame, (224, 224))
        frames.append(frame)
    cap.release()
    return frames

def format_top_predictions(logits, labels, top_k=3):
    probs = torch.nn.functional.softmax(logits, dim=-1)[0]
    top_probs, top_idxs = torch.topk(probs, k=top_k)
    results = []
    for prob, idx in zip(top_probs, top_idxs):
        results.append(f"{labels[idx.item()]} ({prob.item():.2f})")
    return results

def predict(file):
    if file is None:
        return "Please upload an image or video file.", None

    file_path = file.name
    video_exts = [".mp4", ".avi", ".mov", ".mkv"]

    try:
        if any(file_path.lower().endswith(ext) for ext in video_exts):
            frames = extract_frames(file_path)
            if len(frames) == 0:
                return "Failed to extract frames from video.", None
            inputs = video_processor(frames, return_tensors="pt")
            with torch.no_grad():
                outputs = video_model(**inputs)
            logits = outputs.logits
            labels = video_model.config.id2label
            top_preds = format_top_predictions(logits, labels, top_k=3)
            top_confidence = float(torch.nn.functional.softmax(logits, dim=-1)[0].max())
            if top_confidence < 0.5:
                pred_text = "Uncertain prediction: " + ", ".join(top_preds)
            else:
                pred_text = "Top Video Predictions: " + ", ".join(top_preds)
            return pred_text, file_path
        else:
            image = Image.open(file_path).convert("RGB")
            inputs = image_processor(images=image, return_tensors="pt")
            with torch.no_grad():
                outputs = image_model(**inputs)
            logits = outputs.logits
            labels = image_model.config.id2label
            top_preds = format_top_predictions(logits, labels, top_k=3)
            top_confidence = float(torch.nn.functional.softmax(logits, dim=-1)[0].max())
            if top_confidence < 0.5:
                pred_text = "Uncertain prediction: " + ", ".join(top_preds)
            else:
                pred_text = "Top Image Predictions: " + ", ".join(top_preds)
            return pred_text, None
    except Exception as e:
        return f"Error: {e}", None

iface = gr.Interface(
    fn=predict,
    inputs=gr.File(file_types=[".mp4", ".avi", ".mov", ".mkv", ".jpg", ".jpeg", ".png"]),
    outputs=[gr.Textbox(label="Prediction"), gr.Video(label="Uploaded Video")],
    title="Image and Video Classification",
    description="Upload an image or video. Shows top 3 predictions with confidence and a confidence threshold.",
    live=False
)

if __name__ == "__main__":
    iface.launch()