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):
    try:
        print("Received image:", image)  
        image = transform(image).unsqueeze(0)
        print("Transformed image shape:", image.shape)

        # Model inference
        with torch.no_grad():
            output = model(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)