Second check of all

qim3d/io/_convert.py

@@ -7,6 +7,7 @@ import numpy as np
 import tifffile as tiff
 import zarr
 from tqdm import tqdm
+import zarr.core
 
 from qim3d.utils._misc import stringify_path
 from qim3d.io._saving import save
@@ -67,7 +68,7 @@ class Convert:
             else:
                 raise ValueError("Invalid path")
 
-    def convert_tif_to_zarr(self, tif_path: str, zarr_path: str):
+    def convert_tif_to_zarr(self, tif_path: str, zarr_path: str) -> zarr.core.Array:
         """Convert a tiff file to a zarr file
 
         Args:
@@ -97,7 +98,7 @@ class Convert:
 
         return z
 
-    def convert_zarr_to_tif(self, zarr_path: str, tif_path: str):
+    def convert_zarr_to_tif(self, zarr_path: str, tif_path: str) -> None:
         """Convert a zarr file to a tiff file
 
         Args:
@@ -110,7 +111,7 @@ class Convert:
         z = zarr.open(zarr_path)
         save(tif_path, z)
 
-    def convert_nifti_to_zarr(self, nifti_path: str, zarr_path: str):
+    def convert_nifti_to_zarr(self, nifti_path: str, zarr_path: str) -> zarr.core.Array:
         """Convert a nifti file to a zarr file
 
         Args:
@@ -139,7 +140,7 @@ class Convert:
 
         return z
 
-    def convert_zarr_to_nifti(self, zarr_path: str, nifti_path: str, compression: bool = False):
+    def convert_zarr_to_nifti(self, zarr_path: str, nifti_path: str, compression: bool = False) -> None:
         """Convert a zarr file to a nifti file
 
         Args:
@@ -153,7 +154,7 @@ class Convert:
         save(nifti_path, z, compression=compression)
         
 
-def convert(input_path: str, output_path: str, chunk_shape: tuple = (64, 64, 64)):
+def convert(input_path: str, output_path: str, chunk_shape: tuple = (64, 64, 64)) -> None:
     """Convert a file to another format without loading the entire file into memory
 
     Args:

qim3d/io/_loading.py

@@ -29,6 +29,8 @@ from qim3d.utils._system import Memory
 from qim3d.utils._progress_bar import FileLoadingProgressBar
 import trimesh
 
+from typing import Optional, Dict
+
 dask.config.set(scheduler="processes") 
 
 
@@ -100,7 +102,7 @@ class DataLoader:
 
         return vol
 
-    def load_h5(self, path: str):
+    def load_h5(self, path: str) -> tuple[np.ndarray, Optional[Dict]]:
         """Load an HDF5 file from the specified path.
 
         Args:
@@ -183,7 +185,7 @@ class DataLoader:
         else:
             return vol
 
-    def load_tiff_stack(self, path: str):
+    def load_tiff_stack(self, path: str) -> np.ndarray|np.memmap:
         """Load a stack of TIFF files from the specified path.
 
         Args:
@@ -237,7 +239,7 @@ class DataLoader:
 
         return vol
 
-    def load_txrm(self, path: str):
+    def load_txrm(self, path: str) -> tuple[dask.array|np.ndarray, Optional[Dict]]:
         """Load a TXRM/XRM/TXM file from the specified path.
 
         Args:
@@ -766,7 +768,7 @@ def load(
     force_load: bool = False,
     dim_order: tuple = (2, 1, 0),
     **kwargs,
-):
+) -> np.ndarray:
     """
     Load data from the specified file or directory.
 
@@ -854,7 +856,7 @@ def load(
 
     return data
 
-def load_mesh(filename: str):
+def load_mesh(filename: str) -> trimesh.Trimesh:
     """
     Load a mesh from an .obj file using trimesh.
 

qim3d/io/_ome_zarr.py

@@ -82,7 +82,7 @@ class OMEScaler(
 
 
         """
-        def resize_zoom(vol, scale_factors, order, scaled_shape):
+        def resize_zoom(vol: dask.array, scale_factors, order, scaled_shape):
 
             # Get the chunksize needed so that all the blocks match the new shape
             # This snippet comes from the original OME-Zarr-python library
@@ -182,7 +182,7 @@ class OMEScaler(
 
 def export_ome_zarr(
     path: str,
-    data: np.ndarray,
+    data: np.ndarray|dask.array,
     chunk_size: int = 256,
     downsample_rate: int = 2,
     order: int = 1,

qim3d/io/_sync.py

@@ -28,7 +28,7 @@ class Sync:
 
             return False
 
-    def check_destination(self, source: str, destination: str, checksum: bool = False, verbose: bool = True):
+    def check_destination(self, source: str, destination: str, checksum: bool = False, verbose: bool = True) -> list[str]:
         """Check if all files from 'source' are in 'destination'
 
         This function compares the files in the 'source' directory to those in
@@ -80,7 +80,7 @@ class Sync:
 
         return diff_files
 
-    def compare_dirs(self, source: str, destination: str, checksum: bool = False, verbose: bool = True):
+    def compare_dirs(self, source: str, destination: str, checksum: bool = False, verbose: bool = True) -> None:
         """Checks whether 'source' and 'destination' directories are synchronized.
 
         This function compares the contents of two directories
@@ -168,7 +168,7 @@ class Sync:
             )
         return
 
-    def count_files_and_dirs(self, path: str, verbose: bool = True):
+    def count_files_and_dirs(self, path: str, verbose: bool = True) -> tuple[int, int]:
         """Count the number of files and directories in the given path.
 
         This function recursively counts the number of files and

qim3d/ml/_augmentations.py

@@ -20,10 +20,10 @@ class Augmentation:
     """
     
     def __init__(self, 
-                 resize = 'crop', 
-                 transform_train = 'moderate', 
-                 transform_validation = None,
-                 transform_test = None,
+                 resize: str = 'crop', 
+                 transform_train: str = 'moderate', 
+                 transform_validation: str | None = None,
+                 transform_test: str | None = None,
                  mean: float = 0.5, 
                  std: float = 0.5
                 ):
@@ -38,7 +38,7 @@ class Augmentation:
         self.transform_validation = transform_validation
         self.transform_test = transform_test
     
-    def augment(self, im_h: int, im_w: int, level: bool = None):
+    def augment(self, im_h: int, im_w: int, level: bool | None = None):
         """
         Returns an albumentations.core.composition.Compose class depending on the augmentation level.
         A baseline augmentation is implemented regardless of the level, and a set of augmentations are added depending of the level.

qim3d/ml/_data.py

@@ -101,7 +101,7 @@ class Dataset(torch.utils.data.Dataset):
         return Image.open(str(image_path)).size
 
 
-def check_resize(im_height: int, im_width: int, resize: str, n_channels: int):
+def check_resize(im_height: int, im_width: int, resize: str, n_channels: int) -> tuple[int, int]:
     """
     Checks the compatibility of the image shape with the depth of the model.
     If the image height and width cannot be divided by 2 `n_channels` times, then the image size is inappropriate.
@@ -133,7 +133,7 @@ def check_resize(im_height: int, im_width: int, resize: str, n_channels: int):
     return h_adjust, w_adjust 
 
 
-def prepare_datasets(path: str, val_fraction: float, model: nn.Module, augmentation: Augmentation):
+def prepare_datasets(path: str, val_fraction: float, model: nn.Module, augmentation: Augmentation) -> tuple[torch.utils.data.Subset, torch.utils.data.Subset, torch.utils.data.Subset]:
     """
     Splits and augments the train/validation/test datasets.
 
@@ -180,7 +180,7 @@ def prepare_dataloaders(train_set: torch.utils.data,
                         batch_size: int, 
                         shuffle_train: bool = True, 
                         num_workers: int = 8, 
-                        pin_memory: bool = False):  
+                        pin_memory: bool = False) -> tuple[torch.utils.data.DataLoader, torch.utils.data.DataLoader, torch.utils.data.DataLoader]:  
     """
     Prepares the dataloaders for model training.
 

qim3d/ml/_ml_utils.py

@@ -3,7 +3,7 @@
 import torch
 import numpy as np
 
-from torchinfo import summary
+from torchinfo import summary, ModelStatistics
 from qim3d.utils._logger import log
 from qim3d.viz._metrics import plot_metrics
 
@@ -137,7 +137,7 @@ def train_model(
         return train_loss, val_loss
 
 
-def model_summary(dataloader: torch.utils.data.DataLoader, model: torch.nn.Module):
+def model_summary(dataloader: torch.utils.data.DataLoader, model: torch.nn.Module) -> ModelStatistics:
     """Prints the summary of a PyTorch model.
 
     Args:

qim3d/segmentation/_connected_components.py

@@ -23,7 +23,7 @@ class CC:
         """
         return self.cc_count
 
-    def get_cc(self, index=None, crop=False):
+    def get_cc(self, index: int|None =None, crop: bool=False) -> np.ndarray:
         """
         Get the connected component with the given index, if index is None selects a random component.
 
@@ -52,7 +52,7 @@ class CC:
         
         return volume
 
-    def get_bounding_box(self, index=None):
+    def get_bounding_box(self, index: int|None =None)-> list[tuple]:
         """Get the bounding boxes of the connected components.
 
         Args:

qim3d/viz/_data_exploration.py

@@ -298,7 +298,7 @@ def slices_grid(
     return fig
 
 
-def _get_slice_range(position: int, num_slices: int, n_total) -> np.ndarray:
+def _get_slice_range(position: int, num_slices: int, n_total: int) -> np.ndarray:
     """Helper function for `slices`. Returns the range of slices to be displayed around the given position."""
     start_idx = position - num_slices // 2
     end_idx = (
@@ -856,14 +856,14 @@ def histogram(
     log_scale: bool = False,
     despine: bool = True,
     show_title: bool = True,
-    color="qim3d",
-    edgecolor=None,
-    figsize=(8, 4.5),
-    element="step",
-    return_fig=False,
-    show=True,
+    color: str = "qim3d",
+    edgecolor: str|None = None,
+    figsize: tuple[float, float] = (8, 4.5),
+    element: str = "step",
+    return_fig: bool = False,
+    show: bool = True,
     **sns_kwargs,
-) -> Optional[matplotlib.figure.Figure]:
+) -> None|matplotlib.figure.Figure:
     """
     Plots a histogram of voxel intensities from a 3D volume, with options to show a specific slice or the entire volume.
 

qim3d/viz/_k3d.py

@@ -22,8 +22,8 @@ def volumetric(
     grid_visible: bool = False,
     color_map: str = 'magma',
     constant_opacity: bool = False,
-    vmin: float = None,
-    vmax: float = None,
+    vmin: float|None = None,
+    vmax: float|None = None,
     samples: int|str = "auto",
     max_voxels: int = 512**3,
     data_type: str = "scaled_float16",

qim3d/viz/_metrics.py

@@ -13,7 +13,7 @@ def plot_metrics(
     *metrics: tuple[dict[str, float]],
     linestyle: str = "-",
     batch_linestyle: str = "dotted",
-    labels: list = None,
+    labels: list|None = None,
     figsize: tuple = (16, 6),
     show: bool = False
 ):

qim3d/viz/_structure_tensor.py

@@ -19,9 +19,9 @@ def vectors(
     vec: np.ndarray,
     axis: int = 0,
     volume_cmap:str = 'grey',
-    vmin:float = None,
-    vmax:float = None,
-    slice_idx: Optional[Union[int, float]] = None,
+    vmin: float|None = None,
+    vmax: float|None = None,
+    slice_idx: Union[int, float]|None = None,
     grid_size: int = 10,
     interactive: bool = True,
     figsize: Tuple[int, int] = (10, 5),