refactor: change variable names for readability

Author: okunator

cellseg_models_pytorch/decoders/decoder.py

@@ -14,8 +14,8 @@ def __init__(
         enc_channels: Tuple[int, ...],
         out_channels: Tuple[int, ...] = (256, 128, 64, 32, 16),
         style_channels: int = None,
-        n_layers: Tuple[int, ...] = (1, 1, 1, 1, 1),
-        n_blocks: Tuple[Tuple[int, ...], ...] = ((2,), (2,), (2,), (2,), (2,)),
+        n_conv_layers: Tuple[int, ...] = (1, 1, 1, 1, 1),
+        n_conv_blocks: Tuple[Tuple[int, ...], ...] = ((2,), (2,), (2,), (2,), (2,)),
         long_skip: str = "unet",
         n_transformers: Tuple[int, ...] = None,
         n_transformer_blocks: Tuple[Tuple[int], ...] = ((1,), (1,), (1,), (1,), (1,)),
@@ -32,9 +32,9 @@ def __init__(
                 Number of channels at each decoder layer output.
             style_channels : int, default=None
                 Number of style vector channels. If None, style vectors are ignored.
-            n_layers : Tuple[int, ...], default=(1, 1, 1, 1, 1)
+            n_conv_layers : Tuple[int, ...], default=(1, 1, 1, 1, 1)
                 The number of conv layers inside each of the decoder stages.
-            n_blocks : Tuple[Tuple[int, ...], ...] = ((2, ),(2, ),(2, ),(2, ),(2, ))
+            n_conv_blocks : Tuple[Tuple[int, ...], ...] =((2, ),(2, ),(2, ),(2, ),(2, ))
                 The number of blocks inside each conv-layer at each decoder stage.
             long_skip : str, default="unet"
                 long skip method to be used. One of: "unet", "unetpp", "unet3p",
@@ -76,14 +76,14 @@ def __init__(
         # Build decoder
         for i in range(depth - 1):
             # number of conv layers
-            n_conv_layers = None
-            if n_layers is not None:
-                n_conv_layers = n_layers[i]
+            n_clayers = None
+            if n_conv_layers is not None:
+                n_clayers = n_conv_layers[i]
 
             # number of conv blocks inside each layer
-            n_conv_blocks = None
-            if n_blocks is not None:
-                n_conv_blocks = n_blocks[i]
+            n_cblocks = None
+            if n_conv_blocks is not None:
+                n_cblocks = n_conv_blocks[i]
 
             # number of transformer layers
             n_tr_layers = None
@@ -102,8 +102,8 @@ def __init__(
                 skip_channels=skip_channels,
                 style_channels=style_channels,
                 long_skip=long_skip,
-                n_layers=n_conv_layers,
-                n_blocks=n_conv_blocks,
+                n_conv_layers=n_clayers,
+                n_conv_blocks=n_cblocks,
                 n_transformers=n_tr_layers,
                 n_transformer_blocks=n_tr_blocks,
                 **stage_params[i] if stage_params is not None else {"k": None},

cellseg_models_pytorch/decoders/decoder_stage.py

@@ -17,8 +17,8 @@ def __init__(
         dec_dims: Tuple[int, ...],
         skip_channels: Tuple[int, ...],
         style_channels: int = None,
-        n_layers: int = 1,
-        n_blocks: Tuple[int, ...] = (2,),
+        n_conv_layers: int = 1,
+        n_conv_blocks: Tuple[int, ...] = (2,),
         short_skips: Tuple[str, ...] = ("residual",),
         expand_ratios: Tuple[float, float] = ((1.0, 1.0),),
         block_types: Tuple[Tuple[str, ...], ...] = (("basic", "basic"),),
@@ -32,14 +32,15 @@ def __init__(
         kernel_sizes: Tuple[Tuple[int, ...]] = ((3, 3),),
         groups: Tuple[Tuple[int, ...]] = ((1, 1),),
         biases: Tuple[Tuple[bool, ...]] = ((False, False),),
+        layer_residual: bool = False,
         upsampling: str = "fixed-unpool",
         long_skip: str = "unet",
         merge_policy: str = "sum",
-        layer_residual: bool = False,
         skip_params: Optional[Dict[str, Any]] = None,
         n_transformers: Optional[int] = None,
         n_transformer_blocks: Optional[Tuple[int, ...]] = (1,),
-        self_attentions: Optional[Tuple[Tuple[str, ...], ...]] = (("basic",),),
+        transformer_blocks: Optional[Tuple[Tuple[str, ...], ...]] = (("exact",),),
+        transformer_computations: Optional[Tuple[Tuple[str, ...], ...]] = (("basic",),),
         transformer_biases: Optional[Tuple[Tuple[bool, ...], ...]] = ((False,),),
         transformer_dropouts: Optional[Tuple[Tuple[float, ...], ...]] = ((0.0,),),
         transformer_params: Optional[List[Dict[str, Any]]] = None,
@@ -68,48 +69,59 @@ def __init__(
                 `long_skip` == None.
             style_channels : int, default=None
                 Number of style vector channels. If None, style vectors are ignored.
-                Also, ignored if `n_layers` is None.
-            n_layers : int, default=1
+                Also, ignored if `n_conv_layers` is None.
+            n_conv_layers : int, default=1
                 The number of conv layers inside one decoder stage.
-            n_blocks : Tuple[int, ...], default=(2,)
+            n_conv_blocks : Tuple[int, ...], default=(2,)
                 Number of conv-blocks inside each conv layer. The tuple-length has to
-                match `n_layers`. Ignored if `n_layers` is None.
+                match `n_conv_layers`. Ignored if `n_conv_layers` is None.
             short_skips : str, default=("residual", )
                 The short skip methods used inside the conv layers. Ignored if
-                `n_layers` is None.
+                `n_conv_layers` is None.
             expand_ratios : Tuple[float, ...], default=((1.0, 1.0),):
                 Expansion/Squeeze ratios for the out channels of each conv block.
-                The tuple-length has to match `n_layers`. Ignored if `n_layers` is None.
+                The tuple-length has to match `n_conv_layers`. Ignored if
+                `n_conv_layers` is None.
             block_types : Tuple[Tuple[str, ...], ...], default=(("basic", "basic"), )
                 The type of the convolution blocks in the conv blocks inside the layers.
-                The tuple-length has to match `n_layers`. Ignored if `n_layers` is None.
+                The tuple-length has to match `n_conv_layers`. Ignored if
+                `n_conv_layers` is None.
             normalizations : Tuple[Tuple[str, ...], ...], default: (("bn", "bn"), )
                 Normalization methods used in the conv blocks inside the conv layers.
-                The tuple-length has to match `n_layers`. Ignored if `n_layers` is None.
+                The tuple-length has to match `n_conv_layers`. Ignored if
+                `n_conv_layers` is None.
             activations : Tuple[Tuple[str, ...], ...], default: (("relu", "relu"), )
                 Activation methods used inside the conv layers.
-                The tuple-length has to match `n_layers`. Ignored if `n_layers` is None.
+                The tuple-length has to match `n_conv_layers`. Ignored if
+                `n_conv_layers` is None.
             attentions : Tuple[Tuple[str, ...], ...], default: ((None, "se"), )
                 Attention methods used inside the conv layers.
-                The tuple-length has to match `n_layers`. Ignored if `n_layers` is None.
+                The tuple-length has to match `n_conv_layers`. Ignored if
+                `n_conv_layers` is None.
             preactivates Tuple[Tuple[bool, ...], ...], default: ((False, False), )
                 Boolean flags for the conv layers to use pre-activation.
-                The tuple-length has to match `n_layers`. Ignored if `n_layers` is None.
+                The tuple-length has to match `n_conv_layers`. Ignored if
+                `n_conv_layers` is None.
             preattends Tuple[Tuple[bool, ...], ...], default: ((False, False), )
                 Boolean flags for the conv layers to use pre-activation.
-                The tuple-length has to match `n_layers`. Ignored if `n_layers` is None.
+                The tuple-length has to match `n_conv_layers`. Ignored if
+                `n_conv_layers` is None.
             use_styles : Tuple[Tuple[bool, ...], ...], default=((False, False), )
                 Boolean flags for the conv layers to add style vectors at each block.
-                The tuple-length has to match `n_layers`. Ignored if `n_layers` is None.
+                The tuple-length has to match `n_conv_layers`. Ignored if
+                `n_conv_layers` is None.
             kernel_sizes : Tuple[int, ...], default=((3, 3),)
                 The size of the convolution kernels in each conv block.
-                The tuple-length has to match `n_layers`. Ignored if `n_layers` is None.
+                The tuple-length has to match `n_conv_layers`. Ignored if
+                `n_conv_layers` is None.
             groups : int, default=((1, 1),)
                 Number of groups for the kernels in each convolution blocks.
-                The tuple-length has to match `n_layers`. Ignored if `n_layers` is None.
+                The tuple-length has to match `n_conv_layers`. Ignored if
+                `n_conv_layers` is None.
             biases : bool, default=((False, False),)
                 Include bias terms in the convolution blocks.
-                The tuple-length has to match `n_layers`. Ignored if `n_layers` is None.
+                The tuple-length has to match `n_conv_layers`. Ignored if
+                `n_conv_layers` is None.
             upsampling : str, default="fixed-unpool"
                 Name of the upsampling method.
             long_skip : str, default="unet"
@@ -128,10 +140,16 @@ def __init__(
             n_transformer_blocks : int, default=(2, ), optional
                 Number of multi-head self attention blocks used in the transformer
                 layers. Ignored if `n_transformers` is None.
-            self_attentions : Tuple[Tuple[str, ...], ...], default=(("basic",),)
-                The self-attention mechanisms used in the transformer layers.
+            transformer_blocks : Tuple[Tuple[str, ...], ...], default=(("basic",),)
+                The name of the SelfAttentionBlocks in the TransformerLayer(s).
                 Allowed values: "basic", "slice", "flash". Ignored if `n_transformers`
-                is None.
+                is None. Length of the tuple has to equal `n_transformer_blocks`.
+                Allowed names: ("exact", "linformer").
+            transformer_computations : Tuple[Tuple[str, ...],...], default=(("basic",),)
+                The way of computing the attention matrices in the SelfAttentionBlocks
+                in the TransformerLayer(s). Length of the tuple has to equal
+                `n_transformer_blocks`. Allowed styles: "basic". "slice", "flash",
+                "memeff", "slice-memeff".
             transformer_biases : Tuple[Tuple[bool, ...], ...], default=((False,),)
                 Flags, whether to use biases in the transformer layers. Ignored if
                 `n_transformers` is None.
@@ -146,13 +164,13 @@ def __init__(
         Raises
         ------
             ValueError:
-                If lengths of the conv layer tuple args are not equal to `n_layers`.
+                If lengths of the conv layer tuple args are not equal to `n_conv_layers`
                 If lengths of the transformer layer tuple args are not equal to
                 `n_transformers`.
         """
         super().__init__()
 
-        self.n_layers = n_layers
+        self.n_conv_layers = n_conv_layers
         self.n_transformers = n_transformers
         self.long_skip = long_skip
         self.stage_ix = stage_ix
@@ -176,19 +194,20 @@ def __init__(
 
         # Set up n layers of conv blocks
         layer = None  # placeholder
-        if n_layers is not None:
+        if n_conv_layers is not None:
 
             # check that the conv-layer tuple-args are not illegal.
             self._check_tuple_args(
                 "conv-layer related",
-                "n_layers",
-                n_layers,
+                "n_conv_layers",
+                n_conv_layers,
                 all_args=locals(),
                 skip_args=(
                     skip_channels,
                     dec_channels,
                     dec_dims,
-                    self_attentions,
+                    transformer_blocks,
+                    transformer_computations,
                     n_transformer_blocks,
                     transformer_biases,
                     transformer_dropouts,
@@ -197,12 +216,12 @@ def __init__(
 
             # set up the conv-layers.
             self.conv_layers = nn.ModuleDict()
-            for i in range(n_layers):
+            for i in range(n_conv_layers):
                 n_in_feats = self.skip.out_channels if i == 0 else layer.out_channels
                 layer = ConvLayer(
                     in_channels=n_in_feats,
                     out_channels=self.out_channels,
-                    n_blocks=n_blocks[i],
+                    n_blocks=n_conv_blocks[i],
                     layer_residual=layer_residual,
                     style_channels=style_channels,
                     short_skip=short_skips[i],
@@ -225,7 +244,9 @@ def __init__(
             self.out_channels = layer.out_channels
 
         # set in_channels for final operations
-        in_channels = self.skip.out_channels if n_layers is None else self.out_channels
+        in_channels = (
+            self.skip.out_channels if n_conv_layers is None else self.out_channels
+        )
 
         if n_transformers is not None:
 
@@ -239,7 +260,7 @@ def __init__(
                     skip_channels,
                     dec_channels,
                     dec_dims,
-                    n_blocks,
+                    n_conv_blocks,
                     short_skips,
                     expand_ratios,
                     block_types,
@@ -262,7 +283,8 @@ def __init__(
                 tr = Transformer2D(
                     in_channels=in_channels,
                     n_blocks=n_transformer_blocks[i],
-                    block_types=self_attentions[i],
+                    block_types=transformer_blocks[i],
+                    computation_types=transformer_computations[i],
                     biases=transformer_biases[i],
                     dropouts=transformer_dropouts[i],
                     **transformer_params
@@ -272,7 +294,7 @@ def __init__(
                 self.transformers[f"tr_layer_{i + 1}"] = tr
 
         # add a channel pooling layer at the end if no conv-layers are set up
-        if n_layers is None:
+        if n_conv_layers is None:
             self.ch_pool = ChannelPool(
                 in_channels=in_channels,
                 out_channels=self.out_channels,
@@ -338,7 +360,7 @@ def forward(
         x = x[0] if self.long_skip == "unetpp" else x
 
         # conv layers
-        if self.n_layers is not None:
+        if self.n_conv_layers is not None:
             for conv_layer in self.conv_layers.values():
                 x = conv_layer(x, style)  # (B, out_channels, H, W)
 
@@ -348,7 +370,7 @@ def forward(
                 x = transformer(x)  # (B, long_skip_channels/out_channels, H, W)
 
         # channel pool if conv-layers are skipped.
-        if self.n_layers is None:
+        if self.n_conv_layers is None:
             x = self.ch_pool(x)  # (B, out_channels, H, W)
 
         return x, extra_skips

cellseg_models_pytorch/models/cellpose/cellpose.py

@@ -174,8 +174,8 @@ def __init__(
                 out_channels=out_channels,
                 style_channels=style_channels,
                 long_skip=long_skip,
-                n_layers=n_layers,
-                n_blocks=n_blocks,
+                n_conv_layers=n_layers,
+                n_conv_blocks=n_blocks,
                 stage_params=dec_params[decoder_name],
             )
             self.add_module(f"{decoder_name}_decoder", decoder)

cellseg_models_pytorch/models/hovernet/hovernet.py

@@ -168,8 +168,8 @@ def __init__(
                 style_channels=style_channels,
                 long_skip=long_skip,
                 merge_policy=merge_policy,
-                n_layers=n_layers,
-                n_blocks=n_blocks,
+                n_conv_layers=n_layers,
+                n_conv_blocks=n_blocks,
                 stage_params=dec_params[decoder_name],
             )
             self.add_module(f"{decoder_name}_decoder", decoder)

cellseg_models_pytorch/models/stardist/stardist.py

@@ -169,8 +169,8 @@ def __init__(
                 out_channels=out_channels,
                 style_channels=style_channels,
                 long_skip=long_skip,
-                n_layers=n_layers,
-                n_blocks=n_blocks,
+                n_conv_layers=n_layers,
+                n_conv_blocks=n_blocks,
                 stage_params=dec_params[decoder_name],
             )
             self.add_module(f"{decoder_name}_decoder", decoder)

cellseg_models_pytorch/transforms/functional/hovernet.py

@@ -78,7 +78,7 @@ def gen_hv_maps(inst_map: np.ndarray, min_size: int = 5) -> np.ndarray:
         inst = inst[y1:y2, x1:x2]
 
         # instance center of mass, rounded to nearest pixel
-        inst_com = list(ndi.measurements.center_of_mass(inst))
+        inst_com = list(ndi.center_of_mass(inst))
         inst_com[0] = int(inst_com[0] + 0.5)
         inst_com[1] = int(inst_com[1] + 0.5)