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

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

# Load model
model_hg = torch.load(model_path)
model_hg.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_path):
    
    try:
        image = Image.open(image_path).convert("RGB")  
        image = transform(image).unsqueeze(0) 
        with torch.no_grad():
            output = model_hg(image)
            print("Model output:", output)
            class_idx = torch.argmax(output, dim=1).item()

        return "Diagram" if class_idx == 0 else "Not Diagram"
    except Exception as e:
        print("Error during prediction:", str(e)) 
        return f"Prediction Error: {str(e)}"

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