return other results

app.py

@@ -251,11 +251,18 @@ async def download_file(file_name: str):
     
     return FileResponse(file_path, filename=sanitized_name)
 
+def softmax(x):
+    exp_x = np.exp(x - np.max(x))  # Subtract max for numerical stability
+    return exp_x / np.sum(exp_x)
+
 def predict_malware(img_array: np.ndarray) -> str:
     """Make prediction using the preloaded model"""
     try:
         prediction = app.state.model.predict(img_array)
-        return MAL_CLASSES[np.argmax(prediction)]
+        return {
+            "result": MAL_CLASSES[np.argmax(prediction)],
+            "all_results": dict(zip(keys, softmax(prediction)))
+        }
     except Exception as e:
         logger.error(f"Prediction error: {str(e)}")
         raise HTTPException(
@@ -291,7 +298,7 @@ async def analyse(file_name: str):
     try:
         img_array = await process_image(file_path)
         result = predict_malware(img_array)
-        return {"result": result}
+        return result
     except HTTPException as he:
         raise he
     except Exception as e:
@@ -361,7 +368,7 @@ async def analyse_bin(file_name: str):
         convert_binary_to_image(file_path, temp_image, width)
         img_array = await process_image(temp_image)
         result = predict_malware(img_array)
-        return {"result": result}
+        return result
         
     except HTTPException as he:
         raise he