app.py

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

model_path = hf_hub_download(repo_id="Ayamohamed/DiaClassification", filename="dia_none_classifier_full.pth")
model = torch.load(model_path, map_location=torch.device("cpu"), 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 = local_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)