Bug: KeyError: 'binary_map' when using StarDist model with sliding window inference #75
Description
A KeyError: 'binary_map' occurs specifically when using the StarDist and CPPNet model with the sliding window predictor (predict_sliding_win). The error does not occur when using other models, such as CellPose, under the same conditions.
This appears to be caused by a typo in the result aggregation logic in predictor.py. This buggy code path is only triggered when a model provides a nuc_binary output head, which the library's default StarDist model does, but the CellPose model does not.
The error is caused by an incorrect key access when handling the nuc_binary output during the stitching process of the sliding window inference.
In cellseg_models_pytorch/inference/predictor.py, the following line attempts to access a non-existent key:
if nuc_binary is not None:
soft_masks["binary_map"].aux_map = (nuc_binary / recovery)[...]
The soft_masks dictionary does not have a top-level key named "binary_map". The binary map is an attribute of the SoftInstanceOutput object stored under the "nuc" key.
The line should be corrected to access the binary_map attribute of soft_masks["nuc"]:
if nuc_binary is not None:
soft_masks["nuc"].binary_map = (nuc_binary / recovery)[...]
This ensures the aggregated binary prediction is correctly assigned. Please check this issue. :D @okunator