return link to image generated from ai

app.py

@@ -261,7 +261,7 @@ def predict_malware(img_array: np.ndarray) -> str:
         prediction = app.state.model.predict(img_array)
         return {
             "result": MAL_CLASSES[np.argmax(prediction)],
-            "all_results": dict(zip(keys, softmax(prediction)))
+            "all_results": dict(zip(MAL_CLASSES, softmax(prediction)))
         }
     except Exception as e:
         logger.error(f"Prediction error: {str(e)}")
@@ -283,30 +283,30 @@ async def process_image(file_path: str, target_size: tuple = (64, 64)) -> np.nda
             detail="Invalid image file format"
         )
 
-@app.get("/analyse/{file_name}")
-async def analyse(file_name: str):
-    """Analyze image files"""
-    sanitized_name = Path(file_name).name
-    file_path = os.path.join(UPLOAD_DIR, sanitized_name)
+# @app.get("/analyse/{file_name}")
+# async def analyse(file_name: str):
+#     """Analyze image files"""
+#     sanitized_name = Path(file_name).name
+#     file_path = os.path.join(UPLOAD_DIR, sanitized_name)
     
-    if not os.path.exists(file_path):
-        raise HTTPException(
-            status_code=status.HTTP_404_NOT_FOUND,
-            detail="File not found"
-        )
+#     if not os.path.exists(file_path):
+#         raise HTTPException(
+#             status_code=status.HTTP_404_NOT_FOUND,
+#             detail="File not found"
+#         )
     
-    try:
-        img_array = await process_image(file_path)
-        result = predict_malware(img_array)
-        return result
-    except HTTPException as he:
-        raise he
-    except Exception as e:
-        logger.error(f"Analysis error: {str(e)}")
-        raise HTTPException(
-            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-            detail=f"Analysis failed: {str(e)}"
-        )
+#     try:
+#         img_array = await process_image(file_path)
+#         result = predict_malware(img_array)
+#         return result
+#     except HTTPException as he:
+#         raise he
+#     except Exception as e:
+#         logger.error(f"Analysis error: {str(e)}")
+#         raise HTTPException(
+#             status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+#             detail=f"Analysis failed: {str(e)}"
+#         )
 
 def get_image_width(file_path: str) -> int:
     """Determine image width based on file size"""
@@ -349,6 +349,12 @@ def convert_binary_to_image(binary_path: str, output_path: str, width: int):
             detail="Invalid binary file format"
         )
 
+@app.get("/image")
+async def image(file_name: str):
+    if os.path.exists(file_name):
+        return FileResponse(file_name, media_type="image/png")
+    return JSONResponse(content={"error": "Image not found"})
+
 @app.get("/analysebin/{file_name}")
 async def analyse_bin(file_name: str):
     """Analyze binary files by converting to images"""
@@ -368,6 +374,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)
+        result['image'] = temp_image
         return result
         
     except HTTPException as he:
@@ -378,9 +385,6 @@ async def analyse_bin(file_name: str):
             status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
             detail=f"Binary analysis failed: {str(e)}"
         )
-    finally:
-        if os.path.exists(temp_image):
-            os.remove(temp_image)
 
 @app.get("/analyse/yara/{file_name}")
 async def analyse_yara(file_name: str):