Update app.py

app.py

@@ -259,9 +259,10 @@ def predict_malware(img_array: np.ndarray) -> str:
     """Make prediction using the preloaded model"""
     try:
         prediction = app.state.model.predict(img_array)
+        probabilities = softmax(prediction).flatten().tolist()
         return {
             "result": MAL_CLASSES[np.argmax(prediction)],
-            "all_results": dict(zip(MAL_CLASSES, softmax(prediction).tolist()))
+            "all_results": dict(zip(MAL_CLASSES, probabilities))
         }
     except Exception as e:
         logger.error(f"Prediction error: {str(e)}")