Update app.py

app.py

@@ -14,13 +14,25 @@ transform = transforms.Compose([
     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"
+    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 == 1 else "Not Diagram"
+
+    except Exception as e:
+        print("Error during prediction:", str(e)) 
+        return f"Prediction Error: {str(e)}"
 
 gr.Interface(
     fn=predict,