Spaces:
Sleeping
Sleeping
File size: 1,244 Bytes
971b356 d20b3ac 971b356 637c47d 971b356 dbd5e78 971b356 dbd5e78 30784af dbd5e78 971b356 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 |
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)
|