rollaback to fastapi

app.py

@@ -1,96 +1,122 @@
-import gradio as gr
-from cellpose import models, utils
+from fastapi import FastAPI, UploadFile, File, Response
+from cellpose import models, utils, io
 import numpy as np
-import os
 import cv2
-from cellpose.plot import mask_overlay
+import os
+import multiprocessing
+import time
+import sys
+import logging
+
+# --- LOGGING SETUP ---
+io.logger_setup()
+logging.basicConfig(level=logging.INFO, format='%(asctime)s %(message)s', stream=sys.stdout)
+logger = logging.getLogger(__name__)
 
-# --- SETUP ---
-# Point to baked-in models (if you are using a Docker container)
+# Point to baked-in models
 os.environ["CELLPOSE_LOCAL_MODELS_PATH"] = "/app/models"
 
-# 1. LOAD MODEL GLOBALLY
-# We load the model once at startup. This fixes the "slow" issue.
-# Gradio handles concurrency automatically, so we don't need manual multiprocessing.
-print(">>> Loading Cellpose Model...")
-model = models.CellposeModel(model_type='cpsam', gpu=False)
-print(">>> Model Loaded.")
-
-def segment(image, diameter, flow_threshold, cellprob_threshold, channels_str):
-    """
-    Args:
-        image: Numpy array from Gradio (H, W, 3)
-        diameter: Number
-        flow_threshold: Number
-        cellprob_threshold: Number
-        channels_str: String "0,0"
-    """
-    if image is None:
-        return None, "No image provided"
-
-    # 2. PARSE INPUTS
-    try:
-        chan_list = [int(c) for c in channels_str.split(',')]
-    except:
-        chan_list = [0, 0]
+app = FastAPI()
 
-    print(">>> Running Inference...")
-    
-    # 3. RUN INFERENCE
-    # Gradio handles threading, so we just run the eval directly.
-    masks, flows, styles = model.eval(
-        image, 
-        diameter=float(diameter) if diameter > 0 else None,
-        channels=chan_list,
-        flow_threshold=float(flow_threshold),
-        cellprob_threshold=float(cellprob_threshold)
-    )
+# Global process management
+current_process = None
+result_queue = multiprocessing.Queue()
 
-    # 4. GENERATE TEXT OUTPUT (Your original CSV format)
-    outlines = utils.outlines_list(masks)
-    response_lines = []
-    for outline in outlines:
-        roi_coords = []
-        for point in outline:
-            x, y = point
-            roi_coords.append(f"{x},{y}")
-        response_lines.append(",".join(roi_coords))
-    
-    text_result = "\n".join(response_lines)
+# --- WORKER FUNCTION ---
+def run_segmentation(img_data, diameter, flow, prob, channels, queue):
+    try:
+        print(">>> START SEGMENTATION JOB")
+        sys.stdout.flush()
 
-    # 5. GENERATE VISUAL OUTPUT
-    # Since we are using Gradio, let's return an image with masks drawn on it
-    # so you can verify the results visually.
-    overlay_image = mask_overlay(image, masks)
+        # Load Model (Fast because cached)
+        model = models.CellposeModel(model_type='cpsam', gpu=False)
 
-    return overlay_image, text_result
+        # Decode Image
+        nparr = np.frombuffer(img_data, np.uint8)
+        img = cv2.imdecode(nparr, cv2.IMREAD_UNCHANGED)
 
-# --- UI DEFINITION ---
-with gr.Blocks(title="Cellpose Segmentation API") as app:
-    gr.Markdown("## Cellpose Segmentation (CPSAM)")
-    
-    with gr.Row():
-        with gr.Column():
-            input_img = gr.Image(label="Input Image", type="numpy")
-            # Settings
-            diam = gr.Number(label="Diameter (0 = auto)", value=0)
-            flow = gr.Slider(label="Flow Threshold", minimum=0.0, maximum=1.0, value=0.4)
-            prob = gr.Slider(label="Cellprob Threshold", minimum=-6.0, maximum=6.0, value=0.0)
-            chans = gr.Textbox(label="Channels", value="0,0")
-            run_btn = gr.Button("Segment", variant="primary")
+        # Run Inference
+        chan_list = [int(c) for c in channels.split(',')]
+        masks, flows, styles = model.eval(
+            img, 
+            diameter=diameter if diameter > 0 else None,
+            channels=chan_list,
+            flow_threshold=flow,
+            cellprob_threshold=prob
+        )
+
+        # Calculate Outlines
+        outlines = utils.outlines_list(masks)
         
-        with gr.Column():
-            output_img = gr.Image(label="Visual Result")
-            output_text = gr.Textbox(label="ROI Coordinates (CSV)", show_copy_button=True, lines=10)
-
-    # Connect logic
-    run_btn.click(
-        fn=segment,
-        inputs=[input_img, diam, flow, prob, chans],
-        outputs=[output_img, output_text]
+        # Format output
+        response_lines = []
+        for outline in outlines:
+            roi_coords = []
+            for point in outline:
+                x, y = point
+                roi_coords.append(f"{x},{y}")
+            response_lines.append(",".join(roi_coords))
+            
+        result_text = "\n".join(response_lines)
+        
+        queue.put({"status": "success", "data": result_text})
+        print("<<< END SEGMENTATION JOB")
+        sys.stdout.flush()
+
+    except Exception as e:
+        print(f"!!! WORKER ERROR: {e}")
+        queue.put({"status": "error", "message": str(e)})
+
+@app.post("/segment")
+async def segment(
+    image: UploadFile = File(...),
+    diameter: float = 0.0,
+    flow_threshold: float = 0.4,
+    cellprob_threshold: float = 0.0,
+    channels: str = "0,0"
+):
+    global current_process
+
+    # 1. KILL EXISTING JOB (If a new request comes in)
+    if current_process is not None and current_process.is_alive():
+        logger.info(f"⚠️ Stopping previous job (PID: {current_process.pid}) for new request")
+        current_process.terminate()
+        current_process.join()
+        while not result_queue.empty():
+            result_queue.get()
+
+    # 2. READ DATA
+    img_data = await image.read()
+
+    # 3. START WORKER
+    current_process = multiprocessing.Process(
+        target=run_segmentation,
+        args=(img_data, diameter, flow_threshold, cellprob_threshold, channels, result_queue)
     )
+    current_process.start()
+    
+    start_time = time.time()
+    SERVER_TIMEOUT = 300  # 5 Minutes (Seconds)
 
-# Launch with queue enabled (handles the multiprocessing logic automatically)
-if __name__ == "__main__":
-    app.queue(max_size=10) # Manages traffic automatically
-    app.launch(server_name="0.0.0.0", server_port=7860)
+    # 4. MONITOR LOOP
+    while current_process.is_alive():
+        # A. Check Timeout
+        elapsed = time.time() - start_time
+        if elapsed > SERVER_TIMEOUT:
+            logger.error("⏰ JOB TIMED OUT. Killing process.")
+            current_process.terminate()
+            current_process.join()
+            return Response(content="Server Timeout: Processing took longer than 5 minutes.", status_code=504)
+
+        # B. Check for Result
+        try:
+            # Poll queue every 0.5s
+            result = result_queue.get(timeout=0.5)
+            if result["status"] == "success":
+                return Response(content=result["data"], media_type="text/plain")
+            else:
+                return Response(content=result["message"], status_code=500)
+        except:
+            continue
+    
+    return Response(content="Process failed silently", status_code=500)