MouseLand / cellpose

a generalist algorithm for cellular segmentation with human-in-the-loop capabilities
https://www.cellpose.org/
BSD 3-Clause "New" or "Revised" License
1.24k stars 359 forks source link

[Question] Scaling of cellprob and flows #935

Closed FrickTobias closed 1 month ago

FrickTobias commented 1 month ago

TL;DR

How should I compute flows from a manually corrected labeled mask correctly?

My current approach produces floats in range (-1, 1), which is different from when I just load and print a non-corrected file.

Code for how I concluded ranges are different between recomputed and loaded flows:

# load flow from CellPose segmented file
flow = np.load(img0_seg.npy, allow_pickle)["flows"][4][:2]

# inspect range
flow.min() # will print roughly -5
flow.max() # will print roughly 5

# Load corrected masks from cellpose gui corrections
corrected_masks = np.load(img0_corrected_seg.npy, allow_pickle)["masks"]

# recompute flows
corrected_flows = dynamics.labels_to_flows([corrected_masks])[0][-2:]

# inspect range
corrected_flows.min() # will print -1
corrected_flows.max() # will print 1

Context of why I want to recompute the flows

I'm trying out some segmentation repairing by adjusting cellprob and flows with my own downstream models. However, I'm having issues with doing the post processing, more specifically when calculating the labeled masks, where I get no cells as output.

I think this is because CellPose somehow scales cellprob and flows, but I'm not able to find it in the source code.

PS. Possibly this is because I'm not fully understanding the difference between dP and flows, which I currently understand as being the same. Or maybe I'm using the wrong dynamics function to go from cellprob and flow to labeled masks.

# train model to reproduce recomputed flows and cellprob
cellpose_repair_model = train_model()

# from saved CellPose segmented file
cellprobs_and_flows = [
    np.load(img0_seg.npy, allow_pickle)["flows"][4], 
    np.load(img1_seg.npy, allow_pickle)["flows"][4],
    np.load(img2_seg.npy, allow_pickle)["flows"][4]
]

# convert to tensor
cellprobs_and_flows_tensor = torch.tensor(cellprobs_and_flows)

# run model (outputs same dimensions and channels as input)
repaired_cellprobs_and_flows = cellpose_repair_model(cellprobs_and_flows)

for repaired_cellprob_and_flow in repaired_cellprobs_and_flows:

    # split into flows and cellprob
    repaired_flow = repaired_cellprobs_and_flow[:2]
    repaired_cellprob = repaired_cellprobs_and_flow[2]

    # compute labels for masks 
    masks, p = dynamics.compute_masks(repaired_flow, repaired_cellprob)

    # will print 0
    print(masks.max())

Manually modifying the range to (-5, 5) helps but doesn't solve everything

    # compute labels for masks 
    masks, p = dynamics.compute_masks(repaired_flow*5, repaired_cellprob)

    # will print something more similar to what it's supposed to be but still misses some masks
    print(masks.max())

Repaired flows visually looks fine

flows[0] flows[1] cellprob

carsen-stringer commented 1 month ago

You are correct, we scale dP (the output of the model) by a factor of 5 for the dynamics -- but a cellpose model cannot take as input flows, that seems to be your issue. The cellpose model will mess with the flows and filter them in strange ways

FrickTobias commented 1 month ago

Okay thank you for the confirmation! It helps a lot!