Update main.py

main.py

@@ -1,5 +1,6 @@
 from fastapi import FastAPI, UploadFile, File
 from fastapi.responses import FileResponse
+from starlette.middleware.base import BaseHTTPMiddleware
 import cv2
 import subprocess
 import uuid
@@ -21,6 +22,7 @@ async def process_video(file: UploadFile = File(...)):
     cap = cv2.VideoCapture(input_filename)
     success, frame = cap.read()
     if not success:
+        os.remove(input_filename)
         return {"error": "Não foi possível ler o vídeo."}
     
     roi_box = detect_product_in_frame(frame)
@@ -36,8 +38,21 @@ async def process_video(file: UploadFile = File(...)):
             output_filename
         ]
         subprocess.run(crop_cmd)
-
         os.remove(input_filename)
-        return FileResponse(output_filename, media_type="video/mp4", filename="cropped_video.mp4")
+
+        response = FileResponse(output_filename, media_type="video/mp4", filename="cropped_video.mp4")
+        response.headers["X-Delete-File"] = output_filename
+        return response
     else:
+        os.remove(input_filename)
         return {"error": "Não foi possível detectar o produto para realizar o crop."}
+
+class CleanupMiddleware(BaseHTTPMiddleware):
+    async def dispatch(self, request, call_next):
+        response = await call_next(request)
+        file_to_delete = response.headers.get("X-Delete-File")
+        if file_to_delete and os.path.exists(file_to_delete):
+            os.remove(file_to_delete)
+        return response
+
+app.add_middleware(CleanupMiddleware)