import torch
import torch.nn as nn
import cv2
import os
import gradio as gr
from PIL import Image
from transformers import ViTImageProcessor, ViTModel

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
MODEL_PATH = "Vit_LSTM.pth" 
NUM_CLASSES = 5 
MAX_FRAMES = 16  

class ViT_LSTM(nn.Module):
    def __init__(self, feature_dim=768, hidden_dim=512, num_classes=NUM_CLASSES):
        super(ViT_LSTM, self).__init__()
        self.lstm = nn.LSTM(feature_dim, hidden_dim, batch_first=True, num_layers=2, bidirectional=True)
        self.fc = nn.Linear(hidden_dim * 2, num_classes)
        self.dropout = nn.Dropout(0.3)

    def forward(self, x):
        lstm_out, _ = self.lstm(x)
        lstm_out = lstm_out[:, -1, :]  
        out = self.dropout(lstm_out)
        out = self.fc(out)
        return out

vit_processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224")
vit_model = ViTModel.from_pretrained("google/vit-base-patch16-224").to(DEVICE)

model = ViT_LSTM()
model.load_state_dict(torch.load(MODEL_PATH, map_location=DEVICE))
model.to(DEVICE)
model.eval()
LABELS = ["BaseballPitch", "Basketball", "BenchPress", "Biking", "Billiards"]

def extract_vit_features(video_path, max_frames=MAX_FRAMES):
    cap = cv2.VideoCapture(video_path)
    frames = []
    frame_count = 0

    while cap.isOpened() and frame_count < max_frames:
        ret, frame = cap.read()
        if not ret:
            break
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)  
        frames.append(Image.fromarray(frame))
        frame_count += 1

    cap.release()

    if not frames:
        return None

    print(f"Extracted {len(frames)} frames from video.")

    inputs = vit_processor(images=frames, return_tensors="pt")["pixel_values"].to(DEVICE)

    with torch.no_grad():
        features = vit_model(inputs).last_hidden_state.mean(dim=1)  

    return features 

def predict_action(video_file):
    video_path = video_file.name  
    print(f"Received video path: {video_path}") 

    features = extract_vit_features(video_path)

    if features is None:
        return "No frames extracted, please upload a valid video."

    features = features.unsqueeze(0) 

    with torch.no_grad():
        output = model(features)
        predicted_class = torch.argmax(output, dim=1).item()

    return f"Predicted Action: {LABELS[predicted_class]}"

#  Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# Action Recognition with ViT-LSTM")
    gr.Markdown("Upload a short video to predict the action.")

    video_input = gr.File(label="Upload a video")
    output_text = gr.Textbox(label="Prediction")

    predict_btn = gr.Button("Predict Action")
    predict_btn.click(fn=predict_action, inputs=video_input, outputs=output_text)

demo.launch()