Update main.py

main.py

@@ -13,9 +13,11 @@ from scipy.optimize import curve_fit, minimize
 from ultralytics import YOLO
 from fastapi import FastAPI, UploadFile, File, Query, Form
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi.concurrency import run_in_threadpool
 import uvicorn
 from skimage import color
 import hashlib
+import asyncio
 
 # OOM PREVENTION
 torch.set_num_threads(1)
@@ -796,36 +798,41 @@ app.add_middleware(
 
 processor = WatermelonProcessor(MODEL_PATH)
 
+# --- CONCURRENCY & QUEUE MANAGEMENT ---
+dev_lock = asyncio.Lock()
+gen_lock = asyncio.Lock()
+dev_queue_count = 0
+gen_queue_count = 0
+
 @app.get("/")
-def read_root(): return {"status": "Phenotyping API is awake and running!"}
+def read_root(): return {"status": "Watermelon API is awake and running!"}
+
+@app.get("/queue_status")
+def get_queue_status():
+    return {"dev_queue": dev_queue_count, "gen_queue": gen_queue_count}
 
 @app.post("/process_single")
 async def process_single(
     file: UploadFile = File(...), 
     include_image: bool = Query(True), 
-    password: str = Form("")
+    password: str = Form(""),
+    username: str = Form("")
 ):
+    global dev_queue_count, gen_queue_count
     request_t = time.perf_counter()
     
-    # --- SECURE PASSWORD CHECK ---
     expected_hash = "9139eb3676d5dfafced7613f044d86d9e7c84f40a04c83ddce062878621315d0"
     if hashlib.sha256(password.encode('utf-8')).hexdigest() != expected_hash:
-        return ProcessResult(success=False, message="Unauthorized: Incorrect password.", filename=file.filename, processing_ms=int(round((time.perf_counter() - request_t) * 1000))).__dict__
+        return ProcessResult(success=False, message="Unauthorized.", filename=file.filename).__dict__
 
     contents, img = None, None
 
     try:
         contents = await file.read()
-        if not contents:
-            res = ProcessResult(success=False, message="Empty upload.", filename=file.filename)
-            res.processing_ms = int(round((time.perf_counter() - request_t) * 1000))
-            return res.__dict__
+        if not contents: return ProcessResult(success=False, message="Empty.", filename=file.filename).__dict__
 
         img = cv2.imdecode(np.frombuffer(contents, np.uint8), cv2.IMREAD_COLOR)
-        if img is None:
-            res = ProcessResult(success=False, message="Could not decode image.", filename=file.filename)
-            res.processing_ms = int(round((time.perf_counter() - request_t) * 1000))
-            return res.__dict__
+        if img is None: return ProcessResult(success=False, message="Decode error.", filename=file.filename).__dict__
 
         scale_ratio = 1.0
         h, w = img.shape[:2]
@@ -833,19 +840,31 @@ async def process_single(
             scale_ratio = MAX_IMAGE_SIZE / float(max(h, w))
             img = cv2.resize(img, (int(w * scale_ratio), int(h * scale_ratio)), interpolation=cv2.INTER_AREA)
 
-        res = processor.process_image(img, file.filename, scale_ratio, include_image=include_image)
+        # --- CPU CORE ROUTING & QUEUE LOGIC ---
+        is_dev = (username.strip().lower() == 'devtest')
+        
+        # run_in_threadpool prevents OpenCV/YOLO from freezing the API so /queue_status can still answer
+        if is_dev:
+            dev_queue_count += 1
+            try:
+                async with dev_lock:
+                    res = await run_in_threadpool(processor.process_image, img, file.filename, scale_ratio, include_image)
+            finally:
+                dev_queue_count -= 1
+        else:
+            gen_queue_count += 1
+            try:
+                async with gen_lock:
+                    res = await run_in_threadpool(processor.process_image, img, file.filename, scale_ratio, include_image)
+            finally:
+                gen_queue_count -= 1
+
         res.processing_ms = int(round((time.perf_counter() - request_t) * 1000))
         return res.__dict__
 
     except Exception as exc:
         traceback.print_exc()
-        res = ProcessResult(
-            success=False,
-            message=f"Server error: {type(exc).__name__}: {exc}",
-            filename=file.filename,
-            processing_ms=int(round((time.perf_counter() - request_t) * 1000)),
-        )
-        return res.__dict__
+        return ProcessResult(success=False, message=str(exc), filename=file.filename).__dict__
 
     finally:
         del img, contents