Misc classification typing (#106)

* Misc classification typing

* Use numpy/dask array type

* Use types.array in cube_generator

src/cellfinder_core/classify/classify.py

@@ -1,20 +1,20 @@
-from typing import Callable, Optional
+from typing import Callable, List, Optional
 
 import numpy as np
 from imlib.general.system import get_num_processes
 from tensorflow import keras
 
-from cellfinder_core import logger
+from cellfinder_core import logger, types
 from cellfinder_core.classify.cube_generator import CubeGeneratorFromFile
 from cellfinder_core.classify.tools import get_model
 from cellfinder_core.train.train_yml import models
 
 
 def main(
     points,
-    signal_array,
-    background_array,
-    n_free_cpus,
+    signal_array: types.array,
+    background_array: types.array,
+    n_free_cpus: int,
     voxel_sizes,
     network_voxel_sizes,
     batch_size,
@@ -27,7 +27,7 @@ def main(
     max_workers=3,
     *,
     callback: Optional[Callable[[int], None]] = None,
-):
+) -> List:
     """
     Parameters
     ----------

src/cellfinder_core/classify/cube_generator.py

@@ -1,4 +1,5 @@
 from random import shuffle
+from typing import Optional, Tuple
 
 import numpy as np
 import tensorflow as tf
@@ -8,6 +9,7 @@
 from skimage.io import imread
 from tensorflow.keras.utils import Sequence
 
+from cellfinder_core import types
 from cellfinder_core.classify.augment import AugmentationParameters, augment
 
 # TODO: rename, as now using dask arrays -
@@ -33,26 +35,26 @@ class CubeGeneratorFromFile(Sequence):
     def __init__(
         self,
         points,
-        signal_array,
-        background_array,
+        signal_array: types.array,
+        background_array: types.array,
         voxel_sizes,
         network_voxel_sizes,
-        batch_size=16,
-        cube_width=50,
-        cube_height=50,
-        cube_depth=20,
-        channels=2,  # No other option currently
-        classes=2,
-        extract=False,
-        train=False,
-        augment=False,
-        augment_likelihood=0.1,
-        flip_axis=[0, 1, 2],
-        rotate_max_axes=[1, 1, 1],  # degrees
+        batch_size: Optional[int] = 16,
+        cube_width: Optional[int] = 50,
+        cube_height: Optional[int] = 50,
+        cube_depth: Optional[int] = 20,
+        channels: Optional[int] = 2,  # No other option currently
+        classes: Optional[int] = 2,
+        extract: Optional[bool] = False,
+        train: Optional[bool] = False,
+        augment: Optional[bool] = False,
+        augment_likelihood: Optional[float] = 0.1,
+        flip_axis: Tuple[int, int, int] = (0, 1, 2),
+        rotate_max_axes: Tuple[float, float, float] = (1, 1, 1),  # degrees
         # scale=[0.5, 2],  # min, max
-        translate=[0.05, 0.05, 0.05],
-        shuffle=False,
-        interpolation_order=2,
+        translate: Tuple[float, float, float] = (0.05, 0.05, 0.05),
+        shuffle: Optional[bool] = False,
+        interpolation_order: Optional[int] = 2,
     ):
         self.points = points
         self.signal_array = signal_array

src/cellfinder_core/classify/resnet.py

@@ -1,4 +1,7 @@
+from typing import Dict, List, Literal, Tuple, Union
+
 from tensorflow.keras import Model
+from tensorflow.keras.initializers import Initializer
 from tensorflow.keras.layers import (
     Activation,
     Add,
@@ -15,15 +18,19 @@
 #####################################################################
 # Define the types of ResNet
 
-resnet_unit_blocks = {
+layer_type = Literal[
+    "18-layer", "34-layer", "50-layer", "101-layer", "152-layer"
+]
+
+resnet_unit_blocks: Dict[layer_type, List[int]] = {
     "18-layer": [2, 2, 2, 2],
     "34-layer": [3, 4, 6, 3],
     "50-layer": [3, 4, 6, 3],
     "101-layer": [3, 4, 23, 3],
     "152-layer": [3, 6, 36, 3],
 }
 
-network_residual_bottleneck = {
+network_residual_bottleneck: Dict[layer_type, bool] = {
     "18-layer": False,
     "34-layer": False,
     "50-layer": True,
@@ -34,31 +41,17 @@
 
 
 def build_model(
-    shape=(50, 50, 20, 2),
-    network_depth="18-layer",
+    shape: Tuple[int, int, int, int] = (50, 50, 20, 2),
+    network_depth: layer_type = "18-layer",
     optimizer=None,
-    learning_rate=0.0005,  # higher rates don't always converge
+    learning_rate: float = 0.0005,  # higher rates don't always converge
     loss="categorical_crossentropy",
     metrics=["accuracy"],
-    number_classes=2,
-    axis=3,
-    starting_features=64,
+    number_classes: int = 2,
+    axis: int = 3,
+    starting_features: int = 64,
     classification_activation="softmax",
-):
-    """
-
-    :param shape:
-    :param network_depth:
-    :param optimizer:
-    :param learning_rate:
-    :param loss:
-    :param metrics:
-    :param number_classes:
-    :param int axis: Default: 3. Assumed channels are last
-    :param starting_features: # increases in each set of residual units
-    :param classification_activation:
-    :return:
-    """
+) -> Model:
     blocks, bottleneck = get_resnet_blocks_and_bottleneck(network_depth)
 
     inputs = Input(shape)
@@ -94,7 +87,7 @@ def build_model(
     return model
 
 
-def get_resnet_blocks_and_bottleneck(network_depth):
+def get_resnet_blocks_and_bottleneck(network_depth: layer_type):
     """
     Parses dicts, and returns how many resnet blocks are in each unit, along
     with whether they are bottlneck blocks or not
@@ -304,14 +297,14 @@ def f(x):
 
 def get_shortcut(
     inputs,
-    resnet_unit_label,
-    block_id,
-    features,
-    stride,
-    use_bias=False,
-    kernel_initializer="he_normal",
-    bn_epsilon=1e-5,
-    axis=3,
+    resnet_unit_label: int,
+    block_id: int,
+    features: int,
+    stride: int,
+    use_bias: bool = False,
+    kernel_initializer: Union[str, Initializer] = "he_normal",
+    bn_epsilon: float = 1e-5,
+    axis: int = 3,
 ):
     """
     Create shortcut. For none-bottleneck residual units, this is just the
@@ -351,7 +344,7 @@ def get_shortcut(
         return inputs
 
 
-def get_stride(resnet_unit_id, block_id):
+def get_stride(resnet_unit_id: int, block_id: int) -> int:
     """
     Determines the convolution stride.
 

src/cellfinder_core/classify/tools.py

@@ -1,3 +1,6 @@
+import os
+from typing import Optional
+
 import numpy as np
 import tensorflow as tf
 
@@ -6,8 +9,8 @@
 
 
 def get_model(
-    existing_model=None,
-    model_weights=None,
+    existing_model: Optional[os.PathLike] = None,
+    model_weights: Optional[os.PathLike] = None,
     network_depth=None,
     learning_rate=0.0001,
     inference=False,

src/cellfinder_core/types.py

@@ -0,0 +1,6 @@
+from typing import Union
+
+import dask.array as da
+import numpy as np
+
+array = Union[da.Array, np.ndarray]