| from fastapi import FastAPI, UploadFile, File, Query |
| import tifffile |
| import numpy as np |
| from cellpose import models |
| from fastapi.responses import FileResponse |
| from huggingface_hub import snapshot_download |
| import os |
| from io import BytesIO |
| import tempfile |
|
|
| app = FastAPI( |
| title="Cellpose API", |
| docs_url="/docs", |
| redoc_url=None |
| ) |
|
|
| @app.get("/") |
| def root(): |
| return {"status": "ok"} |
| |
| MODEL_REPO = "Shalani08/cellpose_3.1.1_finetuned" |
|
|
| repo_dir = snapshot_download( |
| repo_id=MODEL_REPO, |
| repo_type="model", |
| revision="8a020c0", |
| token=os.environ.get("HF_TOKEN") |
| ) |
|
|
| model_dir = f"{repo_dir}/ddq_model" |
|
|
| model = models.CellposeModel( |
| pretrained_model=model_dir, |
| gpu=False |
| ) |
|
|
|
|
|
|
| @app.post("/segment") |
| async def segment( |
| image: UploadFile = File(...), |
| diameter: float = Query(0, description="Cell diameter (0 = auto)"), |
| channels: str = Query("0,0", description="Cellpose channels"), |
| flow_threshold: float = Query(0.4), |
| cellprob_threshold: float = Query(0.0) |
| ): |
| img_bytes = await image.read() |
| img = tifffile.imread(BytesIO(img_bytes)) |
|
|
| if img.ndim == 3 and img.shape[-1] == 3: |
| img = img[..., 0] |
|
|
| ch = [int(c) for c in channels.split(",")] |
|
|
| masks, _, _ = model.eval( |
| img, |
| diameter=diameter, |
| channels=ch, |
| flow_threshold=flow_threshold, |
| cellprob_threshold=cellprob_threshold |
| ) |
|
|
| tmp = tempfile.NamedTemporaryFile(suffix=".tif", delete=False) |
| tifffile.imwrite(tmp.name, masks.astype(np.uint16)) |
|
|
| return FileResponse(tmp.name, media_type="image/tiff") |
|
|
