Coverage for src / bioimageio / core / custom_ops / cellpose_flow_dynamics.py: 0%

1""" 

2Built-in custom op: cellpose_flow_dynamics 

3========================================== 

4 

5Decodes Cellpose / Cellpose-SAM model outputs into instance label images 

6using flow-dynamics integration and connected-component labelling. 

7 

8Usage in rdf.yaml 

9----------------- 

10:: 

11 

12 postprocessing: 

13 - id: custom 

14 callable: cellpose_flow_dynamics 

15 source: <path/to/cellpose_flow_dynamics.py> 

16 sha256: <sha256 of the file> 

17 kwargs: # all optional — defaults shown 

18 cellprob_threshold: 0.0 

19 flow_threshold: 0.4 

20 do_3D: false 

21 

22Expected model outputs (in rdf.yaml declaration order): 

23 0 - flow_y : vertical flow field (H x W), float32 

24 1 - flow_x : horizontal flow field (H x W), float32 

25 2 - cellprob : cell probability map (H x W), float32, sigmoid-activated 

26 

27Returns: 

28 labels : instance label image (H x W), int32 

29 0 = background, 1..N = individual object instances 

30 

31References 

32---------- 

33Stringer et al. (2021) "Cellpose: a generalist algorithm for cellular 

34segmentation." Nature Methods 18, 100-106. 

35https://doi.org/10.1038/s41592-020-01018-x 

36 

37Pachitariu & Stringer (2022) "Cellpose 2.0: how to train your own model." 

38Nature Methods 19, 1634-1641. 

39https://doi.org/10.1038/s41592-022-01663-4 

40""" 

41 

42from typing import Any, cast 

43 

44import numpy as np 

45from numpy.typing import NDArray 

46 

47 

48class cellpose_flow_dynamics: 

49 """Cellpose flow-dynamics postprocessing as a callable class. 

50 

51 Instantiated once with configuration kwargs; called once per image. 

52 

53 Example:: 

54 

55 op = cellpose_flow_dynamics(cellprob_threshold=0.0, flow_threshold=0.4) 

56 labels = op(flow_y, flow_x, cellprob) 

57 """ 

58 

59 def __init__( 

60 self, 

61 cellprob_threshold: float = 0.0, 

62 flow_threshold: float = 0.4, 

63 do_3D: bool = False, 

64 ) -> None: 

65 super().__init__() 

66 self.cellprob_threshold = cellprob_threshold 

67 self.flow_threshold = flow_threshold 

68 self.do_3D = do_3D 

69 

70 def __call__(self, *arrays: "NDArray[Any]") -> "NDArray[np.int32]": 

71 """Decode flow fields into instance labels. 

72 

73 Args: 

74 *arrays: Model output tensors in rdf.yaml declaration order: 

75 arrays[0] = flow_y (vertical flow field) 

76 arrays[1] = flow_x (horizontal flow field) 

77 arrays[2] = cellprob (cell probability, sigmoid-activated) 

78 

79 Returns: 

80 Integer label image (H x W), int32. 0 = background. 

81 """ 

82 if len(arrays) < 3: 

83 n = len(arrays) 

84 raise ValueError( 

85 f"cellpose_flow_dynamics expects 3 output tensors (flow_y, flow_x, cellprob), got {n}." 

86 ) 

87 

88 flow_y, flow_x, cellprob = arrays[0], arrays[1], arrays[2] 

89 

90 try: 

91 from cellpose import dynamics # pyright: ignore[reportMissingTypeStubs] 

92 except ImportError as e: 

93 raise ImportError( 

94 "cellpose is required for cellpose_flow_dynamics. Install with: pip install cellpose" 

95 ) from e 

96 

97 flows: NDArray[Any] = np.stack([flow_y, flow_x], axis=0) # (2, H, W) 

98 result: Any = dynamics.compute_masks( # pyright: ignore[reportUnknownVariableType] 

99 flows, 

100 cellprob, 

101 cellprob_threshold=self.cellprob_threshold, 

102 flow_threshold=self.flow_threshold, 

103 do_3D=self.do_3D, 

104 ) 

105 masks: NDArray[Any] = cast(NDArray[Any], result[0]) 

106 return masks.astype(np.int32)