import torch
import torch.nn.functional as F
import torchvision.transforms as transforms
from PIL import Image
from huggingface_hub import hf_hub_download
import gradio as gr

# Download model from Hugging Face Hub
model_path = hf_hub_download(repo_id="Ayamohamed/DiaClassification", filename="model.pth")

# Load model
model = torch.load(model_path,weights_only=False)
model.eval() 

transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])  
])
def predict(image):
    image = transform(image).unsqueeze(0) 
    with torch.no_grad():
        output = model(image)
        probabilities = F.softmax(output, dim=1)
        class_idx = torch.argmax(probabilities, dim=1).item()
        return "Diagram" if class_idx == 0 else "Not Diagram"

gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title="Diagram Classifier"
).launch(share=True)