qim3d/ml/_data.py

@@ -4,7 +4,9 @@ from PIL import Image
 from qim3d.utils._logger import log
 import torch
 import numpy as np
-
+from typing import Optional, Callable
+import torch.nn as nn
+from ._data import Augmentation
 
 class Dataset(torch.utils.data.Dataset):
     """
@@ -36,7 +38,7 @@ class Dataset(torch.utils.data.Dataset):
             transform=albumentations.Compose([ToTensorV2()]))
         image, target = dataset[idx]
     """
-    def __init__(self, root_path: str, split="train", transform=None):
+    def __init__(self, root_path: str, split: str = "train", transform: Optional[Callable] = None):
         super().__init__()
 
         # Check if split is valid
@@ -58,7 +60,7 @@ class Dataset(torch.utils.data.Dataset):
     def __len__(self):
         return len(self.sample_images)
 
-    def __getitem__(self, idx):
+    def __getitem__(self, idx: int):
         image_path = self.sample_images[idx]
         target_path = self.sample_targets[idx]
 
@@ -76,7 +78,7 @@ class Dataset(torch.utils.data.Dataset):
 
 
     # TODO: working with images of different sizes
-    def check_shape_consistency(self,sample_images):
+    def check_shape_consistency(self,sample_images: tuple[str]):
         image_shapes= []
         for image_path in sample_images:
             image_shape = self._get_shape(image_path)
@@ -99,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.
@@ -131,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, 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.
 
@@ -172,7 +174,13 @@ def prepare_datasets(path: str, val_fraction: float, model, augmentation):
     return train_set, val_set, test_set
 
 
-def prepare_dataloaders(train_set, val_set, test_set, batch_size, shuffle_train = True, num_workers = 8, pin_memory = False):  
+def prepare_dataloaders(train_set: torch.utils.data, 
+                        val_set: torch.utils.data, 
+                        test_set: torch.utils.data, 
+                        batch_size: int, 
+                        shuffle_train: bool = True, 
+                        num_workers: int = 8, 
+                        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,24 +3,24 @@
 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
 
 from tqdm.auto import tqdm
 from tqdm.contrib.logging import logging_redirect_tqdm
-
+from models._unet import Hyperparameters
 
 def train_model(
-    model,
-    hyperparameters,
-    train_loader,
-    val_loader,
-    eval_every=1,
-    print_every=5,
-    plot=True,
-    return_loss=False,
-):
+    model: torch.nn.Module,
+    hyperparameters: Hyperparameters,
+    train_loader: torch.utils.data.DataLoader,
+    val_loader: torch.utils.data.DataLoader,
+    eval_every: int = 1,
+    print_every: int = 5,
+    plot: bool = True,
+    return_loss: bool = False,
+) -> tuple[tuple[float], tuple[float]]:
     """Function for training Neural Network models.
 
     Args:
@@ -137,7 +137,7 @@ def train_model(
         return train_loss, val_loss
 
 
-def model_summary(dataloader, model):
+def model_summary(dataloader: torch.utils.data.DataLoader, model: torch.nn.Module) -> ModelStatistics:
     """Prints the summary of a PyTorch model.
 
     Args:
@@ -160,7 +160,7 @@ def model_summary(dataloader, model):
     return model_s
 
 
-def inference(data, model):
+def inference(data: torch.utils.data.Dataset, model: torch.nn.Module) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """Performs inference on input data using the specified model.
 
     Performs inference on the input data using the provided model. The input data should be in the form of a list,
@@ -242,7 +242,7 @@ def inference(data, model):
     return inputs, targets, preds
 
 
-def volume_inference(volume, model, threshold=0.5):
+def volume_inference(volume: np.ndarray, model: torch.nn.Module, threshold:float = 0.5) -> np.ndarray:
     """
     Compute on the entire volume
     Args:

qim3d/ml/models/_unet.py

@@ -119,13 +119,13 @@ class Hyperparameters:
 
     def __init__(
         self,
-        model,
-        n_epochs=10,
-        learning_rate=1e-3,
-        optimizer="Adam",
-        momentum=0,
-        weight_decay=0,
-        loss_function="Focal",
+        model: nn.Module,
+        n_epochs: int = 10,
+        learning_rate: float = 1e-3,
+        optimizer: str = "Adam",
+        momentum: float = 0,
+        weight_decay: float = 0,
+        loss_function: str = "Focal",
     ):
 
         # TODO: implement custom loss_functions? then add a check to see if loss works for segmentation.
@@ -168,13 +168,13 @@ class Hyperparameters:
 
     def model_params(
         self,
-        model,
-        n_epochs,
-        optimizer,
-        learning_rate,
-        weight_decay,
-        momentum,
-        loss_function,
+        model: nn.Module,
+        n_epochs: int,
+        optimizer: str,
+        learning_rate: float,
+        weight_decay: float,
+        momentum: float,
+        loss_function: str,
     ):
 
         optim = self._optimizer(model, optimizer, learning_rate, weight_decay, momentum)
@@ -188,7 +188,12 @@ class Hyperparameters:
         return hyper_dict
 
     # selecting the optimizer
-    def _optimizer(self, model, optimizer, learning_rate, weight_decay, momentum):
+    def _optimizer(self, 
+        model: nn.Module, 
+        optimizer: str, 
+        learning_rate: float, 
+        weight_decay: float, 
+        momentum: float):
         from torch.optim import Adam, SGD, RMSprop
 
         if optimizer == "Adam":
@@ -212,7 +217,7 @@ class Hyperparameters:
         return optim
 
     # selecting the loss function
-    def _loss_functions(self, loss_function):
+    def _loss_functions(self, loss_function: str):
         from monai.losses import FocalLoss, DiceLoss, DiceCELoss
         from torch.nn import BCEWithLogitsLoss
 

qim3d/processing/_layers.py

@@ -2,7 +2,14 @@ import numpy as np
 from slgbuilder import GraphObject 
 from slgbuilder import MaxflowBuilder
 
-def segment_layers(data:np.ndarray, inverted:bool = False, n_layers:int = 1, delta:float = 1, min_margin:int = 10, max_margin:int = None, wrap:bool = False):
+def segment_layers(data: np.ndarray, 
+                   inverted: bool = False, 
+                   n_layers: int = 1, 
+                   delta: float = 1, 
+                   min_margin: int = 10, 
+                   max_margin: int = None, 
+                   wrap: bool = False
+                   ) -> list:
     """
     Works on 2D and 3D data.
     Light one function wrapper around slgbuilder https://github.com/Skielex/slgbuilder to do layer segmentation
@@ -82,7 +89,7 @@ def segment_layers(data:np.ndarray, inverted:bool = False, n_layers:int = 1, del
 
     return segmentations
 
-def get_lines(segmentations:list|np.ndarray) -> list:
+def get_lines(segmentations:list[np.ndarray]) -> list:
     """
     Expects list of arrays where each array is 2D segmentation with only 2 classes. This function gets the border between those two
     so it could be plotted. Used with qim3d.processing.segment_layers

qim3d/processing/_local_thickness.py

@@ -10,7 +10,7 @@ def local_thickness(
     image: np.ndarray,
     scale: float = 1,
     mask: Optional[np.ndarray] = None,
-    visualize=False,
+    visualize: bool = False,
     **viz_kwargs
 ) -> np.ndarray:
     """Wrapper for the local thickness function from the [local thickness package](https://github.com/vedranaa/local-thickness)

qim3d/processing/_structure_tensor.py

@@ -12,7 +12,7 @@ def structure_tensor(
     rho: float = 6.0,
     base_noise: bool = True,
     full: bool = False,
-    visualize=False,
+    visualize: bool = False,
     **viz_kwargs
 ) -> Tuple[np.ndarray, np.ndarray]:
     """Wrapper for the 3D structure tensor implementation from the [structure_tensor package](https://github.com/Skielex/structure-tensor/).

qim3d/segmentation/_connected_components.py

@@ -4,7 +4,7 @@ from qim3d.utils._logger import log
 
 
 class CC:
-    def __init__(self, connected_components, num_connected_components):
+    def __init__(self, connected_components: np.ndarray, num_connected_components: int):
         """
         Initializes a ConnectedComponents object.
 
@@ -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/utils/_doi.py

@@ -4,7 +4,7 @@ import requests
 from qim3d.utils._logger import log
 
 
-def _validate_response(response):
+def _validate_response(response: requests.Response) -> bool:
     # Check if we got a good response
     if not response.ok:
         log.error(f"Could not read the provided DOI ({response.reason})")
@@ -13,7 +13,7 @@ def _validate_response(response):
     return True
 
 
-def _doi_to_url(doi):
+def _doi_to_url(doi: str) -> str:
     if doi[:3] != "http":
         url = "https://doi.org/" + doi
     else:
@@ -22,7 +22,7 @@ def _doi_to_url(doi):
     return url
 
 
-def _make_request(doi, header):
+def _make_request(doi: str, header: str) -> requests.Response:
     # Get url from doi
     url = _doi_to_url(doi)
 
@@ -35,7 +35,7 @@ def _make_request(doi, header):
     return response
 
 
-def _log_and_get_text(doi, header):
+def _log_and_get_text(doi, header) -> str:
     response = _make_request(doi, header)
 
     if response and response.encoding:
@@ -50,13 +50,13 @@ def _log_and_get_text(doi, header):
         return text
 
 
-def get_bibtex(doi):
+def get_bibtex(doi: str):
     """Generates bibtex from doi"""
     header = {"Accept": "application/x-bibtex"}
 
     return _log_and_get_text(doi, header)
 
-def cusom_header(doi, header):
+def custom_header(doi: str, header: str) -> str:
     """Allows a custom header to be passed
 
     For example:
@@ -67,7 +67,7 @@ def cusom_header(doi, header):
     """
     return _log_and_get_text(doi, header)
 
-def get_metadata(doi):
+def get_metadata(doi: str) -> dict:
     """Generates a metadata dictionary from doi"""
     header = {"Accept": "application/vnd.citationstyles.csl+json"}
     response = _make_request(doi, header)
@@ -76,7 +76,7 @@ def get_metadata(doi):
 
     return metadata
 
-def get_reference(doi):
+def get_reference(doi: str) -> str:
     """Generates a metadata dictionary from doi and use it to build a reference string"""
 
     metadata = get_metadata(doi)
@@ -84,7 +84,7 @@ def get_reference(doi):
 
     return reference_string
 
-def build_reference_string(metadata):
+def build_reference_string(metadata: dict) -> str:
     """Generates a reference string from metadata"""
     authors = ", ".join([f"{author['family']} {author['given']}" for author in metadata['author']])
     year = metadata['issued']['date-parts'][0][0]

qim3d/utils/_misc.py

@@ -12,7 +12,7 @@ import difflib
 import qim3d
 
 
-def get_local_ip():
+def get_local_ip() -> str:
     """Retrieves the local IP address of the current machine.
 
     The function uses a socket to determine the local IP address.
@@ -42,7 +42,7 @@ def get_local_ip():
     return ip_address
 
 
-def port_from_str(s):
+def port_from_str(s: str) -> int:
     """
     Generates a port number from a given string.
 
@@ -65,7 +65,7 @@ def port_from_str(s):
     return int(hashlib.sha1(s.encode("utf-8")).hexdigest(), 16) % (10**4)
 
 
-def gradio_header(title, port):
+def gradio_header(title: str, port: int) -> None:
     """Display the header for a Gradio server.
 
     Displays a formatted header containing the provided title,
@@ -99,7 +99,7 @@ def gradio_header(title, port):
     ouf.showlist(details, style="box", title="Starting gradio server")
 
 
-def sizeof(num, suffix="B"):
+def sizeof(num: float, suffix: str = "B") -> str:
     """Converts a number to a human-readable string representing its size.
 
     Converts the given number to a human-readable string representing its size in
@@ -131,7 +131,7 @@ def sizeof(num, suffix="B"):
     return f"{num:.1f} Y{suffix}"
 
 
-def find_similar_paths(path):
+def find_similar_paths(path: str) -> list[str]:
     parent_dir = os.path.dirname(path) or "."
     parent_files = os.listdir(parent_dir) if os.path.isdir(parent_dir) else ""
     valid_paths = [os.path.join(parent_dir, file) for file in parent_files]
@@ -165,14 +165,14 @@ def get_file_size(file_path: str) -> int:
     return file_size
 
 
-def stringify_path(path):
+def stringify_path(path: os.PathLike) -> str:
     """Converts an os.PathLike object to a string"""
     if isinstance(path, os.PathLike):
         path = path.__fspath__()
     return path
 
 
-def get_port_dict():
+def get_port_dict() -> dict:
     # Gets user and port
     username = getpass.getuser()
     url = f"https://platform.qim.dk/qim-api/get-port/{username}"
@@ -189,7 +189,7 @@ def get_port_dict():
     return port_dict
 
 
-def get_css():
+def get_css() -> str:
 
     current_directory = os.path.dirname(os.path.abspath(__file__))
     parent_directory = os.path.abspath(os.path.join(current_directory, os.pardir))
@@ -201,7 +201,7 @@ def get_css():
     return css_content
 
 
-def downscale_img(img, max_voxels=512**3):
+def downscale_img(img: np.ndarray, max_voxels: int = 512**3) -> np.ndarray:
     """Downscale image if total number of voxels exceeds 512³.
 
     Args:
@@ -226,7 +226,7 @@ def downscale_img(img, max_voxels=512**3):
     return zoom(img, zoom_factor, order=0)
 
 
-def scale_to_float16(arr: np.ndarray):
+def scale_to_float16(arr: np.ndarray) -> np.ndarray:
     """
     Scale the input array to the float16 data type.
 

qim3d/utils/_ome_zarr.py

 from zarr.util import normalize_chunks, normalize_dtype, normalize_shape
 import numpy as np
 
-def get_chunk_size(shape:tuple, dtype):
+def get_chunk_size(shape:tuple, dtype: tuple) -> tuple[int, ...]:
     """
     How the chunk size is computed in zarr.storage.init_array_metadata which is ran in the chain of functions we use 
     in qim3d.io.export_ome_zarr function
@@ -20,7 +20,7 @@ def get_chunk_size(shape:tuple, dtype):
     return chunks
 
 
-def get_n_chunks(shapes:tuple, dtypes:tuple):
+def get_n_chunks(shapes: tuple, dtypes: tuple) -> int:
     """
     Estimates how many chunks we will use in advence so we can pass this number to a progress bar and track how many
     have been already written to disk

qim3d/utils/_progress_bar.py

@@ -98,7 +98,7 @@ class FileLoadingProgressBar(ProgressBar):
         super().__init__( tqdm_kwargs, repeat_time)
         self.process = psutil.Process()
 
-    def get_new_update(self):
+    def get_new_update(self) -> int:
         counters = self.process.io_counters()
         try:
             memory = counters.read_chars
@@ -107,7 +107,7 @@ class FileLoadingProgressBar(ProgressBar):
         return memory
 
 class OmeZarrExportProgressBar(ProgressBar):
-    def __init__(self,path:str, n_chunks:int, reapeat_time="auto"):
+    def __init__(self,path: str, n_chunks: int, reapeat_time: str = "auto"):
         """
         Context manager to track the exporting of OmeZarr files.
 
@@ -152,7 +152,7 @@ class OmeZarrExportProgressBar(ProgressBar):
         self.last_update = 0
 
     def get_new_update(self):
-        def file_count(folder_path:str):
+        def file_count(folder_path: str) -> int:
             """
             Goes recursively through the folders and counts how many files are there, 
             Doesn't count metadata json files

qim3d/utils/_server.py

@@ -14,7 +14,7 @@ class CustomHTTPRequestHandler(SimpleHTTPRequestHandler):
         self.send_header("Access-Control-Allow-Headers", "X-Requested-With, Content-Type")
         super().end_headers()
 
-    def list_directory(self, path):
+    def list_directory(self, path: str|os.PathLike):
         """Helper to produce a directory listing, includes hidden files."""
         try:
             file_list = os.listdir(path)
@@ -49,7 +49,7 @@ class CustomHTTPRequestHandler(SimpleHTTPRequestHandler):
         # Write the encoded HTML directly to the response
         self.wfile.write(encoded)
 
-def start_http_server(directory, port=8000):
+def start_http_server(directory: str, port: int = 8000) -> HTTPServer:
     """
     Starts an HTTP server serving the specified directory on the given port with CORS enabled.
     

qim3d/utils/_system.py

@@ -36,7 +36,7 @@ class Memory:
             round(self.free_pct, 1),
         )
 
-def _test_disk_speed(file_size_bytes=1024, ntimes=10):
+def _test_disk_speed(file_size_bytes: int = 1024, ntimes: int =10) -> tuple[float, float, float, float]:
     '''
     Test the write and read speed of the disk by writing a file of a given size
     and then reading it back.
@@ -95,7 +95,7 @@ def _test_disk_speed(file_size_bytes=1024, ntimes=10):
     return avg_write_speed, write_speed_std, avg_read_speed, read_speed_std
 
 
-def disk_report(file_size=1024 * 1024 * 100, ntimes=10):
+def disk_report(file_size: int = 1024 * 1024 * 100, ntimes: int = 10) -> None:
     '''
     Report the average write and read speed of the disk.
 

qim3d/viz/__init__.py

@@ -9,7 +9,7 @@ from ._data_exploration import (
     chunks,
     histogram,
 )
-from .itk_vtk_viewer import itk_vtk, Installer, NotInstalledError
+from .itk_vtk_viewer import itk_vtk
 from ._k3d import volumetric, mesh
 from ._local_thickness import local_thickness
 from ._structure_tensor import vectors

qim3d/viz/_cc.py

@@ -2,18 +2,18 @@ import matplotlib.pyplot as plt
 import numpy as np
 import qim3d
 from qim3d.utils._logger import log
-
+from qim3d.segmentation._connected_components import CC
 
 def plot_cc(
-    connected_components,
+    connected_components: CC,
     component_indexs: list | tuple = None,
-    max_cc_to_plot=32,
-    overlay=None,
-    crop=False,
-    show=True,
-    cmap: str = "viridis",
-    vmin: float = None,
-    vmax: float = None,
+    max_cc_to_plot: int = 32,
+    overlay: np.ndarray = None,
+    crop: bool = False,
+    display_figure: bool = True,
+    color_map: str = "viridis",
+    value_min: float = None,
+    value_max: float = None,
     **kwargs,
 ) -> list[plt.Figure]:
     """
@@ -75,7 +75,7 @@ def plot_cc(
                 cc = connected_components.get_cc(component, crop=False)
                 overlay_crop = np.where(cc == 0, 0, overlay)
             fig = qim3d.viz.slices_grid(
-                overlay_crop, show=show, cmap=cmap, vmin=vmin, vmax=vmax, **kwargs
+                overlay_crop, display_figure=display_figure, color_map=color_map, value_min=value_min, value_max=value_max, **kwargs
             )
         else:
             # assigns discrete color map to each connected component if not given
@@ -89,7 +89,7 @@ def plot_cc(
 
         figs.append(fig)
 
-    if not show:
+    if not display_figure:
         return figs
 
     return

qim3d/viz/_data_exploration.py

@@ -9,7 +9,9 @@ from typing import List, Optional, Union
 
 import dask.array as da
 import ipywidgets as widgets
+import matplotlib.figure
 import matplotlib.pyplot as plt
+import matplotlib
 from IPython.display import SVG, display
 import matplotlib
 import numpy as np
@@ -29,7 +31,7 @@ def slices_grid(
     color_map: str = "magma",
     value_min: float = None,
     value_max: float = None,
-    image_size=None,
+    image_size: int = None,
     image_height: int = 2,
     image_width: int = 2,
     display_figure: bool = False,
@@ -38,7 +40,7 @@ def slices_grid(
     color_bar: bool = False,
     color_bar_style: str = "small",
     **matplotlib_imshow_kwargs,
-) -> plt.Figure:
+) -> matplotlib.figure.Figure:
     """Displays one or several slices from a 3d volume.
 
     By default if `slice_positions` is None, slices_grid plots `num_slices` linearly spaced slices.
@@ -296,7 +298,7 @@ def slices_grid(
     return fig
 
 
-def _get_slice_range(position: int, num_slices: int, n_total):
+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 = (
@@ -324,7 +326,7 @@ def slicer(
     image_width: int = 3,
     display_positions: bool = False,
     interpolation: Optional[str] = None,
-    image_size=None,
+    image_size: int = None,
     color_bar: bool = False,
     **matplotlib_imshow_kwargs,
 ) -> widgets.interactive:
@@ -401,8 +403,8 @@ def slicer_orthogonal(
     image_width: int = 3,
     display_positions: bool = False,
     interpolation: Optional[str] = None,
-    image_size=None,
-):
+    image_size: int = None,
+)-> widgets.interactive:
     """Interactive widget for visualizing orthogonal slices of a 3D volume.
 
     Args:
@@ -461,7 +463,7 @@ def fade_mask(
     color_map: str = "magma",
     value_min: float = None,
     value_max: float = None,
-):
+)-> widgets.interactive:
     """Interactive widget for visualizing the effect of edge fading on a 3D volume.
 
     This can be used to select the best parameters before applying the mask.
@@ -596,7 +598,7 @@ def fade_mask(
     return slicer_obj
 
 
-def chunks(zarr_path: str, **kwargs):
+def chunks(zarr_path: str, **kwargs)-> widgets.interactive:
     """
     Function to visualize chunks of a Zarr dataset using the specified visualization method.
 
@@ -854,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,
-):
+) -> 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/_detection.py

@@ -6,7 +6,7 @@ from IPython.display import clear_output, display
 import qim3d
 
 
-def circles(blobs, vol, alpha=0.5, color="#ff9900", **kwargs):
+def circles(blobs: tuple[float,float,float,float], vol: np.ndarray, alpha: float = 0.5, color: str = "#ff9900", **kwargs)-> widgets.interactive:
     """
     Plots the blobs found on a slice of the volume.
 

qim3d/viz/_k3d.py

@@ -15,18 +15,18 @@ from qim3d.utils._misc import downscale_img, scale_to_float16
 
 
 def volumetric(
-    img,
-    aspectmode="data",
-    show=True,
-    save=False,
-    grid_visible=False,
-    color_map='magma',
-    constant_opacity=False,
-    vmin=None,
-    vmax=None,
-    samples="auto",
-    max_voxels=512**3,
-    data_type="scaled_float16",
+    img: np.ndarray,
+    aspectmode: str = "data",
+    show: bool = True,
+    save: bool = False,
+    grid_visible: bool = False,
+    color_map: str = 'magma',
+    constant_opacity: bool = False,
+    vmin: float|None = None,
+    vmax: float|None = None,
+    samples: int|str = "auto",
+    max_voxels: int = 512**3,
+    data_type: str = "scaled_float16",
     **kwargs,
 ):
     """
@@ -175,13 +175,13 @@ def volumetric(
 
 
 def mesh(
-    verts,
-    faces,
-    wireframe=True,
-    flat_shading=True,
-    grid_visible=False,
-    show=True,
-    save=False,
+    verts: np.ndarray,
+    faces: np.ndarray,
+    wireframe: bool = True,
+    flat_shading: bool = True,
+    grid_visible: bool = False,
+    show: bool = True,
+    save: bool = False,
     **kwargs,
 ):
     """

qim3d/viz/_layers2d.py

@@ -8,7 +8,7 @@ import numpy as np
 from PIL import Image
 
 
-def image_with_lines(image:np.ndarray, lines: list, line_thickness:float|int) -> Image:
+def image_with_lines(image: np.ndarray, lines: list, line_thickness: float) -> Image:
     """
     Plots the image and plots the lines on top of it. Then extracts it as PIL.Image and in the same size as the input image was.
     Paramters: