Update main.py

main.py

@@ -1,14 +1,25 @@
 import logging
-from fastapi import FastAPI, UploadFile, File
+import os
+import shutil
+import tempfile
+import base64
+from typing import List
+
+from fastapi import FastAPI, UploadFile, File, Form, HTTPException
 from fastapi.responses import HTMLResponse, JSONResponse
 from fastapi.staticfiles import StaticFiles
 from fastapi.middleware.cors import CORSMiddleware
-import base64
+
 import mmcv
-from model.run_inference import infer_images
+from model.run_inference import infer_images  # <- update this to accept model=...
 
 logging.basicConfig(level=logging.INFO)
 
+ALLOWED_MODELS = {
+    "regnetx4.0gf+dterd",
+    "regnetx4.0gf+petr",
+}
+
 app = FastAPI()
 app.add_middleware(
     CORSMiddleware,
@@ -19,6 +30,7 @@ app.add_middleware(
 )
 app.mount("/static", StaticFiles(directory="static"), name="static")
 
+
 @app.get("/", response_class=HTMLResponse)
 async def root():
     with open("static/index.html", "r", encoding="utf-8") as f:
@@ -26,26 +38,51 @@ async def root():
 
 
 @app.post("/infer")
-async def run_inference(images: list[UploadFile] = File(...)):
+async def run_inference(
+    model: str = Form(...),
+    images: List[UploadFile] = File(...),
+):
+    model = model.strip().lower()
+    if model not in ALLOWED_MODELS:
+        raise HTTPException(status_code=400, detail=f"Invalid model '{model}'. Allowed: {sorted(ALLOWED_MODELS)}")
+
+    if len(images) != 6:
+        raise HTTPException(status_code=400, detail=f"Expected 6 images, received {len(images)}")
+
+    tmpdir = tempfile.mkdtemp(prefix="bev_infer_")
     img_paths = []
-    for upload in images:
-        data = await upload.read()
-        # Drop any alpha channel, force 3-channel BGR
-        bgr = mmcv.imfrombytes(data, flag="color")
-        tmp = f"/tmp/{upload.filename}"
-        mmcv.imwrite(bgr, tmp)
-        img_paths.append(tmp)
 
     try:
-        bev_paths = infer_images(img_paths)
+        for idx, upload in enumerate(images):
+            data = await upload.read()
+            bgr = mmcv.imfrombytes(data, flag="color")
+            if bgr is None:
+                raise HTTPException(status_code=400, detail=f"File '{upload.filename}' is not a valid image.")
+
+            out_path = os.path.join(tmpdir, f"cam_{idx}.png")
+            mmcv.imwrite(bgr, out_path)
+            img_paths.append(out_path)
+
+        logging.info("Starting inference with model=%s on %d images", model, len(img_paths))
+
+        bev_paths = infer_images(img_paths, model=model)
+
+        output = []
+        for p in bev_paths:
+            with open(p, "rb") as f:
+                b64 = base64.b64encode(f.read()).decode("utf-8")
+            output.append({"bev_image": b64})
+
+        return JSONResponse(content=output)
+
+    except HTTPException:
+
+        raise
     except Exception as e:
         logging.exception("inference failed")
         return JSONResponse(status_code=500, content={"error": str(e)})
-
-    output = []
-    for p in bev_paths:
-        with open(p, "rb") as f:
-            b64 = base64.b64encode(f.read()).decode("utf-8")
-        output.append({"bev_image": b64})
-
-    return JSONResponse(content=output)
+    finally:
+        try:
+            shutil.rmtree(tmpdir)
+        except Exception:
+            logging.warning("Failed to clean tmpdir %s", tmpdir)