import gradio as gr
import torch
from torchvision import transforms
from PIL import Image

# 1. Setup Device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 2. Load the model
# Hugging Face Spaces will look for 'model.pt' in the same folder
model = torch.load("model.pt", map_location=device,weights_only=False)
model.eval()

# 3. Define the Prediction Logic
def predict_signature(inp_img):
    if inp_img is None:
        return "Please upload an image."
    
    # Convert to RGB (handles RGBA or Grayscale uploads)
    img = Image.fromarray(inp_img.astype('uint8'), 'RGB')
    
    # Transformation pipeline (Matching your training code)
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    
    img_tensor = transform(img).unsqueeze(0).to(device)
    
    with torch.no_grad():
        output = model(img_tensor)
        # Apply Softmax to get probabilities
        probs = torch.nn.functional.softmax(output, dim=1)
        confidences = {
            "Forged": float(probs[0][0]),
            "Original": float(probs[0][1])
        }
        
    return confidences

# 4. Create the Gradio Interface
interface = gr.Interface(
    fn=predict_signature,
    inputs=gr.Image(),
    outputs=gr.Label(num_top_classes=2),
    title="ResNet-34 Signature Verification",
    description="Upload a signature image to verify if it is an **Original** or a **Forgery**. This model was fine-tuned on the CEDAR dataset.",
    
)

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