Update main.py

main.py

@@ -1,26 +1,36 @@
 from fastapi import FastAPI, UploadFile, File
 from fastapi.responses import FileResponse
+from starlette.middleware.base import BaseHTTPMiddleware
 import cv2
 import subprocess
 import uuid
 import os
+import shutil
+import tempfile
 
 app = FastAPI()
 
+# Define permissões restritas para os arquivos criados
+os.umask(0o077)
+
 def detect_product_in_frame(frame):
     return (100, 50, 400, 400)
 
 @app.post("/process_video/")
 async def process_video(file: UploadFile = File(...)):
-    input_filename = f"/tmp/input_{uuid.uuid4()}.mp4"
-    output_filename = f"/tmp/output_{uuid.uuid4()}.mp4"
+    # Cria um diretório temporário exclusivo para a requisição
+    temp_dir = tempfile.mkdtemp()
+    input_filename = os.path.join(temp_dir, f"input_{uuid.uuid4()}.mp4")
+    output_filename = os.path.join(temp_dir, f"output_{uuid.uuid4()}.mp4")
     
+    # Salva o vídeo de entrada no diretório temporário
     with open(input_filename, "wb") as f:
         f.write(await file.read())
 
     cap = cv2.VideoCapture(input_filename)
     success, frame = cap.read()
     if not success:
+        shutil.rmtree(temp_dir)
         return {"error": "Não foi possível ler o vídeo."}
     
     roi_box = detect_product_in_frame(frame)
@@ -35,9 +45,31 @@ async def process_video(file: UploadFile = File(...)):
             '-c:a', 'copy',
             output_filename
         ]
-        subprocess.run(crop_cmd)
+        try:
+            subprocess.run(crop_cmd, check=True)
+        except subprocess.CalledProcessError:
+            shutil.rmtree(temp_dir)
+            return {"error": "Erro ao processar o vídeo com FFmpeg."}
 
-        os.remove(input_filename)
-        return FileResponse(output_filename, media_type="video/mp4", filename="cropped_video.mp4")
+        # Verifica se o arquivo de saída foi criado e prepara a resposta
+        if os.path.exists(output_filename):
+            response = FileResponse(output_filename, media_type="video/mp4", filename="cropped_video.mp4")
+            response.headers["X-Delete-Dir"] = temp_dir
+            return response
+        else:
+            shutil.rmtree(temp_dir)
+            return {"error": "O arquivo de saída não foi gerado."}
     else:
+        shutil.rmtree(temp_dir)
         return {"error": "Não foi possível detectar o produto para realizar o crop."}
+
+# Middleware para limpar o diretório temporário após a resposta
+class CleanupMiddleware(BaseHTTPMiddleware):
+    async def dispatch(self, request, call_next):
+        response = await call_next(request)
+        temp_dir_to_delete = response.headers.get("X-Delete-Dir")
+        if temp_dir_to_delete and os.path.exists(temp_dir_to_delete):
+            shutil.rmtree(temp_dir_to_delete)
+        return response
+
+app.add_middleware(CleanupMiddleware)