app.py

import gradio as gr
import cv2
import torch
import numpy as np
from transformers import CLIPProcessor, CLIPVisionModel
from PIL import Image
from torch import nn
import requests
import matplotlib.pyplot as plt
from huggingface_hub import hf_hub_download

# ... (rest of your code remains the same)

def process_image_classification(image):
    model, processor, reverse_mapping, device = load_model()
    
    # Convert image to PIL Image
    image = Image.fromarray(image)
    
    inputs = processor(images=image, return_tensors="pt")
    pixel_values = inputs.pixel_values.to(device)
    
    with torch.no_grad():
        logits, attentions = model(pixel_values, output_attentions=True)
        probs = torch.nn.functional.softmax(logits, dim=-1)
        prediction = torch.argmax(probs).item()
    
    # Generate attention map
    attention_map = get_attention_map(attentions)
    
    visualization = apply_heatmap(image, attention_map)
    
    card_name = reverse_mapping[prediction]
    confidence = probs[0][prediction].item()
    
    # Convert back to RGB for matplotlib display
    visualization_rgb = cv2.cvtColor(visualization, cv2.COLOR_BGR2RGB)
    
    return visualization_rgb, card_name, confidence

def gradio_interface():
    gr_interface = gr.Interface(
        fn=process_image_classification,
        inputs=gr.inputs.Image(type="numpy"),
        outputs=[
            gr.outputs.Image(label="Heatmap Plot"),
            gr.outputs.Textbox(label="Predicted Card"),
            gr.outputs.Textbox(label="Confidence")
        ],
        title="Uno Card Recognizer",
        description="Upload an image or use your webcam to recognize an Uno card."
    )
    gr_interface.launch()

if __name__ == "__main__":
    gradio_interface()