Update app.py

app.py

@@ -28,7 +28,7 @@ def predict(image):
             print("Model output:", output)
             class_idx = torch.argmax(output, dim=1).item()
 
-        return "Diagram" if class_idx == 1 else "Not Diagram"
+        return "Diagram" if class_idx == 0 else "Not Diagram"
 
     except Exception as e:
         print("Error during prediction:", str(e))