fix: sanitize image route input

main.py

@@ -1,5 +1,5 @@
 import numpy as np
-from fastapi import FastAPI, File, UploadFile
+from fastapi import FastAPI, File, UploadFile, HTTPException
 import tensorflow as tf
 from PIL import Image
 from io import BytesIO
@@ -8,6 +8,7 @@ import cv2
 from datetime import datetime
 from fastapi.responses import FileResponse
 from fastapi.middleware.cors import CORSMiddleware
+from pathlib import Path
 app = FastAPI()
 
 app.add_middleware(
@@ -129,14 +130,26 @@ async def predict_v2(file: UploadFile = File(...)):
     plot_detected_rectangles(cv2.imread("input/" + filename), detections, "output/" + filename)
     return {
         "message": "Detection and classification completed successfully",
-        "data": "output/" + filename,
+        "out": filename,
         "class_names": class_names
     }
 
+IMAGE_DIR = Path("output")
+
 @app.get("/image/")
 async def get_image(image_name: str):
-    # Assume the images are stored in a directory named "images"
-    image_path = f"output/{image_name}"
+    # Sanitize the image_name to prevent directory traversal attacks
+    if "../" in image_name:
+        raise HTTPException(status_code=400, detail="Invalid image name")
+
+    # Construct the image path
+    image_path = IMAGE_DIR / image_name
+
+    # Check if the image exists
+    if not image_path.exists() or not image_path.is_file():
+        raise HTTPException(status_code=404, detail="Image not found")
+
+    # Return the image file
     return FileResponse(image_path)
 
 @app.post("/predict")