|
|
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") |
|
|
|
