Update app.py

app.py

@@ -9,28 +9,21 @@ import gradio as gr
 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() 
+model = torch.load(model_path)
+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_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()
-
+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"
-    except Exception as e:
-        print("Error during prediction:", str(e)) 
-        return f"Prediction Error: {str(e)}"
 
 gr.Interface(
     fn=predict,