skoots.lib#

skoots.lib.embedding_to_prob#

skoots.lib.embedding_to_prob.baked_embed_to_prob(embedding, baked_skeletons, sigma, eps=1e-16)[source]#

N Dimensional embedding to probability with a baked skeleton array

Calculates a probability \(\phi\) based on a euclidean distance between a spatial embedding \(E_i\) and a baked skeleton pixel \(S_i\).

\[\phi(E_i, S_i) =exp\left(\sum_{k \in [x,y,z]} \frac{(E_{ki} - S_{ki})^2}{-2\sigma^2_k} \right)\]

In three spatial dimmensions, this expands to

\[\phi(E_i, S_i) =exp\left(\frac{(E_{xi} - S_{xi})^2}{-2\sigma^2_x} + \frac{(E_{yi} - S_{yi})^2}{-2\sigma_y^2} + \frac{(E_{zi} - S_{zi})^2}{-2\sigma^2_z}\right)\]
Shapes:

  • embedding: \((B_{in}, 2, X_{in}, Y_{in})\) or \((B_{in}, 3, X_{in}, Y_{in}, Z_{in})\)

  • baked_skeletons: \((B_{in}, 2, X_in, Y_{in})\) or \((B_{in}, 3, X_{in}, Y_in, Z_{in})\)

  • sigma: \((2)\) or \((3)\)

  • returns: \((B_{in}, 1, X_{in}, Y_{in})\) or \((B_{in}, 1, X_{in}, Y_{in}, Z_{in})\)

Parameters:

  • embedding (Tensor) – embedding tensor

  • baked_skeletons (Tensor) – a baked skeleton tensor

  • sigma (Tensor) – Standard deviation of the gaussian. Larger values give higher probability further away.

  • eps (float) – small value for numerical stability

Return type:

Tensor

Returns:

Probability matrix

skoots.lib.vector_to_embedding#

skoots.lib.vector_to_embedding._vec2embed2D(scale, vector)[source]#

2D vector to embedding

Parameters:

  • scale (Tensor) – The offest in XY of the vectors.

  • vector (Tensor) – [B, C=2, X, Y] the vector matrix predicted by the unet

Return type:

Tensor

Returns:

Pixel Spatial Embeddings (i.e. vector + pixel_indicies)

skoots.lib.vector_to_embedding._vec2embed3D(scale, vector, n=1, decay=1.0)[source]#

2D or 3D vector to embedding Could be a faster way to do this with strides but idk… :type scale: Tensor :param scale: [N=2/3] :type vector: Tensor :param vector: [B, C, X, Y, Z?] :type n: int :param n: number of times to apply vectors :rtype: Tensor :return:

skoots.lib.vector_to_embedding._vec2embed3D_graphable(static_scale, vector, static_mesh)[source]#

3D vector to embedding which uses static inputs for cuda graphs

Could be a faster way to do this with strides but idk…

Parameters:

  • scale – Tensor with shape (3)

  • vector (Tensor) – [B, C, X, Y, Z]

Return type:

Tensor

Returns:

embedding vector

skoots.lib.vector_to_embedding.get_vector_mesh(shape, device)[source]#

generates a 3d mesh from a vector

Return type:

Tensor

skoots.lib.vector_to_embedding.vec2embedND(scale, vector)[source]#

Generic N dimmensional vector to embedding

Could be a faster way to do this with strides but idk…

Parameters:

  • scale – [N=2/3]

  • vector – [B, C, X, Y, Z?]

Returns:

skoots.lib.vector_to_embedding.vector_to_embedding(scale, vector, N=1, decay=1.0)[source]#

Converts a 2D or 3D vector field to a spatial embedding by adding the vector at any position to its own position.

vector is a 2D or 3D vector field of shape \((B, 2, X, Y)\) for 2D or \((B, 3, X, Y, Z)\) for 3D. Each vector “\(v\)” lies within the range -1 and 1 and is scaled by scale “\(s\)”. The scaled vector is then added to its own position to form a spatial embedding “\(\phi\)”:

Formally:
\[ \begin{align}\begin{aligned}i,j,k \in \mathbb{Z}_{≥0} \\v_{i,j,k} \in [-1, 1] \\s = [s_i, s_j, s_k]\\\phi_{i,j,k} = v_{i,j,k} * s + [i, j, k]\end{aligned}\end{align} \]
Shapes:

  • scale: \((2)\) or \((3)\)

  • vector: \((B_{in}, 2, X_{in}, Y_{in})\) or \((B_{in}, 3, X_{in}, Y_{in}, Z_{in})\)

  • Returns: \((B_{in}, 2, X_{in}, Y_{in})\) or \((B_{in}, 3, X_{in}, Y_{in}, Z_{in})\)

Parameters:

  • scale (Tensor) – Scaling factors for each vector spatial dimension

  • vector (Tensor) – Vector field predicted by a neural network

  • N (int) – Number of iterations to apply the vectors.

  • decay (float) – vector strength decay after each iteration. Default 1.0

Return type:

Tensor

Returns:

Pixel spatial embeddings

skoots.lib.flood_fill#

skoots.lib.flood_fill._in_place_replace(x, to_replace, replace_with)[source]#

Performs an in place replacement of values in tensor x.

Checks each location in x for a value in to_replace. if a value is in to_replace, the value is swapped with the associated value in replace_with.

Parameters:

  • x (ndarray) – input nd.array

  • to_replace (ndarray) – array of values to be replaced in x

  • replace_with (ndarray) – array of values by which should be replaced

Return type:

None

Returns:

skoots.lib.flood_fill.connected_components(graph)[source]#

Finds all connected components in a graph of id values by performing depth first search.

Parameters:

graph (Dict[int, List[int]]) – input graph where each key is a node, and each value is a list of edges

Return type:

List[List[int]]

Returns:

list of all connected notes

skoots.lib.flood_fill.dfs(connected, node, graph, visited)[source]#

depth first search for finding connected components of a graph

Return type:

List[int]

skoots.lib.flood_fill.efficient_flood_fill(skeleton)[source]#

Efficiently floods a binary skeleton mask in place by first flood filling small regions, then merging connected components later. Avoids memory copies when possible. Returns a skeleton mask where each connected component has a unique label, however these labels may not be sequential. I.e. unique(skeleton) -> [4, 16, 23, 24, 96]

Parameters:

skeleton (Tensor) – binary skeleton mask to flood fill

Return type:

Tensor

Returns:

Flood filled tensor

skoots.lib.flood_fill.flood_all(x, id)[source]#

Finds all features (connected components in an ndarray) and gives it a unique label from 1 to N for N components

Parameters:

  • x (Tensor) – torch.bool tensor

  • id (int) – previous max id value

Return type:

Tuple[Tensor, int, Dict[int, Tensor]]

Returns:

skoots.lib.flood_fill.get_adjacent_labels(x, y)[source]#

calculates which masks of a signle object have two labels (due to the border)

Parameters:

  • x (Tensor) –

  • y (Tensor) –

Return type:

List[Tuple[int, int]]

Returns:

skoots.lib.flood_fill.replace(x, collisions)[source]#

Performs an in place replacement of values in the input tensor \(x\). This function calls a just-in-time compiled numba kernel which parallelizes the replacement.

Roughly performs this algorith:
for i in range(x):
for (to_replace, replace_with) in collisions:
if x[i] == to_replace:

x[i] = replace_with

return x

Parameters:

  • x (Tensor) – Input torch.Tensor of any size

  • collisions (List[Tuple[int, int]]) – List of collisions where c[0] is the value to replace, and c[1] is the value to replace with

Return type:

Tensor

Returns:

Original tensor with modified memory

skoots.lib.skeleton#

skoots.lib.skeleton._bake_skeleton_torch(masks, skeletons, anisotropy=(1.0, 1.0, 1.0), average=True, device='cpu')[source]#

For each pixel \(p_ik\) of object \(k\) at index \(i\in[x,y,z]\) in masks, returns a baked skeleton where the value at each index is the closest skeleton point \(s_{jk}\) of any instance \(k\).

This should reflect the ACTUAL spatial distance of your dataset for best results…These models tend to like XY embedding vectors more than Z. For anisotropic datasets, you should roughly provide the anisotropic correction factor of each voxel. For instance anisotropy of (1.0, 1.0, 5.0) means that the Z dimension is 5x larger than XY.

Formally, the value at each position \(i\in[x,y,z]\) of the baked skeleton tensor \(S\) is the minimum of the euclidean distance function \(f(a, b)\) and the skeleton point of any instance:

\[S_{i} = min \left( f(i, s_{k})\right)\ for\ k \in [1, 2, ..., N]\]
Shapes:

  • masks: \((1, X_{in}, Y_{in}, Z_{in})\)

  • skeletons: \((3, N_i)\)

  • anisotropy: \((3)\)

  • returns: \((3, X_{in}, Y_{in}, Z_{in})\)

Parameters:

  • masks (Tensor) – Ground Truth instance mask of shape [1, X, Y, Z] of objects where each pixel is an integer id value.

  • skeletons (Dict[int, Tensor]) – Dict of skeleton indicies where each key is a unique instance of an object in mask. - Each skeleton has a shape [3, N] where N is the number of pixels constituting the skeleton

  • anisotropy (List[float]) – Anisotropic correction factor for min distance calculation

  • average (bool) – Average the skeletons such that there is a smooth transition form one area to the next

  • device (str) – torch.Device by which to run calculations

Return type:

Tensor

Returns:

Baked skeleton

skoots.lib.skeleton._bake_skeleton_triton(masks, skeletons, anisotropy, average=True)[source]#

Launches a triton kernel to perform an in place distance calculation

Parameters:

  • masks (Tensor) –

  • skeletons (Dict[int, Tensor]) – [N, 3]

  • anisotropy (List[float]) –

  • average (bool) –

Returns:

skoots.lib.skeleton._skeleton_to_mask(skeletons, shape, kernel_size=(15, 15, 1), n=2)[source]#

Converts a skeleton Dict to a skeleton mask which can simply be regressed against via Dice loss or whatever…

Shapes:

  • skeletons: [N, 3]

  • shape \((3)\)

  • returns: math:(1, X_{in}, Y_{in}, Z_{in})

Parameters:

  • skeletons (Dict[int, Tensor]) – Dict of skeletons

  • shape (Tuple[int, int, int]) – Shape of final mask

Return type:

Tensor

Returns:

Maks of embedded skeleton px

skoots.lib.skeleton.average_baked_skeletons(baked_skeleton, kernel_size=3)[source]#

Takes a baked skeleton computed by skoots.lib.skeleton.bake_skeleton and averages all the values such that there is a smooth transition from one location to another.

Shapes:

  • baked_skeleton \((B_{in}, 3, X_{in}, Y_{in}, Z_{in})\)

  • returns \((B_{in}, 3, X_{in}, Y_{in}, Z_{in})\)

Parameters:

  • baked_skeleton (Tensor) – Baked skeleton tensor

  • kernel_size (int) – Kernel size for smoothing

Return type:

Tensor

Returns:

smoothed skeleton: Smoothed Tensor

skoots.lib.skeleton.bake_skeleton(masks, skeletons, anisotropy=(1.0, 1.0, 1.0), average=True, device='cpu', return_distance=False)[source]#

For each pixel \(p_ik\) of object \(k\) at index \(i\in[x,y,z]\) in masks, returns a baked skeleton where the value at each index is the closest skeleton point \(s_{jk}\) of any instance \(k\).

This should reflect the ACTUAL spatial distance of your dataset for best results…These models tend to like XY embedding vectors more than Z. For anisotropic datasets, you should roughly provide the anisotropic correction factor of each voxel. For instance anisotropy of (1.0, 1.0, 5.0) means that the Z dimension is 5x larger than XY.

Formally, the value at each position \(i\in[x,y,z]\) of the baked skeleton tensor \(S\) is the minimum of the euclidean distance function \(f(a, b)\) and the skeleton point of any instance:

\[S_{i} = min \left( f(i, s_{k})\right)\ for\ k \in [1, 2, ..., N]\]
Shapes:

  • masks: \((1, X_{in}, Y_{in}, Z_{in})\)

  • skeletons: \((3, N_i)\)

  • anisotropy: \((3)\)

  • returns: \((3, X_{in}, Y_{in}, Z_{in})\)

Parameters:

  • masks (Tensor) – Ground Truth instance mask of shape [1, X, Y, Z] of objects where each pixel is an integer id value.

  • skeletons (Dict[int, Tensor]) – Dict of skeleton indicies where each key is a unique instance of an object in mask. - Each skeleton has a shape [3, N] where N is the number of pixels constituting the skeleton

  • anisotropy (Tuple[float, float, float]) – Anisotropic correction factor for min distance calculation

  • average (bool) – Average the skeletons such that there is a smooth transition form one area to the next

  • device (str) – torch.Device by which to run calculations

  • return_distance (bool) – if true and bake_skeletons is dispatching the triton kernel, returns the distance to each closest skeleton

Return type:

Union[Tensor, Tuple[Tensor, Tensor]]

Returns:

Baked skeleton

skoots.lib.skeleton.index_skeleton_by_embed(skeleton, embed)[source]#

Returns an instance mask by indexing skeleton with an embedding tensor For memory efficiency, skeleton is only ever Referenced! Never copied (I hope)

Shapes:

  • skeleton: \((B_{in}=1, 1, X_{in}, Y_{in}, Z_{in})\)

  • embed: \((B_{in}=1, 3, X_{in}, Y_{in}, Z_{in})\)

Parameters:

  • skeleton (Tensor) – Skeleton of a single instance

  • embed (Tensor) – Embedding

Return type:

Tensor

Returns:

torch.int instance mask

skoots.lib.skeleton.next_power_of_2(x)[source]#
skoots.lib.skeleton.skeleton_to_mask(skeletons, shape, device=None, radius=7, flank_radius=3)[source]#

Converts a skeleton Dict to a skeleton mask which can simply be regressed against via Dice loss or similar…

Shapes:

  • skeletons: [N, 3]

  • shape \((3)\)

  • returns: math:(1, X_{in}, Y_{in}, Z_{in})

Parameters:

  • skeletons (Dict[int, Tensor]) – Dict of skeletons

  • shape (Tuple[int, int, int]) – Shape of final mask

  • device (Union[device, str, None]) – output device

Return type:

Tensor

Returns:

Mask of embedded skeleton px

skoots.lib.morphology#

skoots.lib.morphology._compute_zero_padding(kernel_size)[source]#

Utility function that computes zero padding tuple. Adapted from Kornia

Return type:

Tuple[int, int, int]

skoots.lib.morphology._get_binary_kernel2d(window_size, device)[source]#

Creates a symmetric binary kernel to extract the patches. If the window size is HxWxD will create a (H*W)xHxW kernel.

Adapted from a 2D Kornia implementation

Return type:

Tensor

skoots.lib.morphology._get_binary_kernel3d(window_size, device)[source]#

Creates a symmetric binary kernel to extract the patches. If the window size is HxWxD will create a (H*W)xHxW kernel.

Adapted from a 2D Kornia implementation

Return type:

Tensor

skoots.lib.morphology._get_gaussian_kernel1d(kernel_size, sigma)[source]#
Return type:

Tensor

skoots.lib.morphology._get_gaussian_kernel2d(kernel_size, sigma, dtype, device)[source]#
Return type:

Tensor

skoots.lib.morphology._get_gaussian_kernel3d(kernel_size, sigma, dtype, device)[source]#
Return type:

Tensor

skoots.lib.morphology.binary_dilation(image)[source]#

Performs binary dilation on a 5D Tensor.

Shapes:

  • input: \((B, C, X, Y, Z)\)

  • output: \((C, C, X, Y, Z)\)

Parameters:

image (Tensor) – binary image

Return type:

Tensor

Returns:

dilated image

skoots.lib.morphology.binary_dilation_2d(image)[source]#

Performs binary dilation on a 5D Tensor.

Shapes:

  • input: \((B, C, X, Y, Z)\)

  • output: \((C, C, X, Y, Z)\)

Parameters:

image (Tensor) – binary image

Return type:

Tensor

Returns:

dilated image

skoots.lib.morphology.binary_erosion(image)[source]#

Performs binary erosion on a 5D Tensor.

Shapes:

  • input: \((B, C, X, Y, Z)\)

  • output: \((B, C, X, Y, Z)\)

Parameters:

image (Tensor) – binary image

Return type:

Tensor

Returns:

eroded image

skoots.lib.morphology.dilate(input)[source]#
Return type:

Tensor

skoots.lib.morphology.gauss_filter(input, kernel, sigma)[source]#

gaussian filter of a 3D tensor

Parameters:

  • input (Tensor) – (B, C, X, Y, Z)

  • kernel (List[int]) – [int, int, int]

  • sigma (List[float]) – [float, float, float]

Return type:

Tensor

Returns:

blured image

skoots.lib.morphology.mean_filter(input)[source]#
Return type:

Tensor

skoots.lib.morphology.median_filter(input)[source]#
Return type:

Tensor

skoots.lib.utils#

skoots.lib.utils._crop2d(img, x, y, w, h)[source]#

torch scriptable function which crops an image

Parameters:

  • img (Tensor) – torch.Tensor image of shape [C, X, Y, Z]

  • x (int) – x coord of crop box

  • y (int) – y coord of crop box

  • z – z coord of crop box

  • w (int) – width of crop box

  • h (int) – height of crop box

  • d – depth of crop box

Return type:

Tensor

Returns:

skoots.lib.utils._crop3d(img, x, y, z, w, h, d)[source]#

torch scriptable function which crops an image

Parameters:

  • img (Tensor) – torch.Tensor image of shape [C, X, Y, Z]

  • x (int) – x coord of crop box

  • y (int) – y coord of crop box

  • z (int) – z coord of crop box

  • w (int) – width of crop box

  • h (int) – height of crop box

  • d (int) – depth of crop box

Return type:

Tensor

Returns:

skoots.lib.utils._get_cached_disk_coords(device)[source]#
Return type:

Tensor

skoots.lib.utils.calculate_indexes(pad_size, eval_image_size, image_shape, padded_image_shape)[source]#

This calculates indexes for the complete evaluation of an arbitrarily large image by unet. each index is offset by eval_image_size, but has a width of eval_image_size + pad_size * 2. Unet needs padding on each side of the evaluation to ensure only full convolutions are used in generation of the final mask. If the algorithm cannot evenly create indexes for padded_image_shape, an additional index is added at the end of equal size.

Parameters:

  • pad_size (int) – int corresponding to the amount of padding on each side of the padded image

  • eval_image_size (int) – int corresponding to the shape of the image to be used for the final mask

  • image_shape (int) – int Shape of image before padding is applied

  • padded_image_shape (int) – int Shape of image after padding is applied

Return type:

List[List[int]]

Returns:

List of lists corresponding to the indexes

skoots.lib.utils.cantor2(a, b)[source]#

Hashes two combination of tensors together

Return type:

Tensor

skoots.lib.utils.cantor3(a, b, c)[source]#

Hashes three combination of tensors together

Return type:

Tensor

skoots.lib.utils.cfg_to_bism_model(cfg)[source]#

utility function to get a bism model from cfg

Return type:

Module

skoots.lib.utils.crop_to_identical_size(a, b)[source]#

Crops Tensor a to the shape of Tensor b, then crops Tensor b to the shape of Tensor a.

Parameters:

  • a (Tensor) – torch.

  • b (Tensor) –

Return type:

Tuple[Tensor, Tensor]

Returns:

skoots.lib.utils.get_cached_ball_coords(device)[source]#
Return type:

Tensor

skoots.lib.utils.get_cached_disk_coords(device, radius, flank_radius)[source]#
Return type:

Tensor

skoots.lib.utils.get_dtype_offset(dtype='uint16', image_max=None)[source]#

Returns the scaling factor such that such that :math:`

rac{image}{f} in [0, …, 1]`

Supports: uint16, uint8, uint12, float64

Return type:

int

skoots.lib.utils.identical_rows(a, b)[source]#

Given two matrices of identical size, determines the indices of identical rows.

Parameters:

  • a (Tensor) – [N, 3] torch.Tensor

  • b (Tensor) – [N, 3] torch.Tensor

Return type:

Tensor

Returns:

Indicies of identical rows

skoots.lib.mp_utils#

For code used in distributed training. Basically just found a stackoverflow implementation and got it to work.

https://stackoverflow.com/questions/1365265/on-localhost-how-do-i-pick-a-free-port-number

skoots.lib.mp_utils.cleanup(rank)[source]#

Destroy a given process group, and deinitialize the distributed package

skoots.lib.mp_utils.find_free_port()[source]#

https://stackoverflow.com/questions/1365265/on-localhost-how-do-i-pick-a-free-port-number

skoots.lib.mp_utils.get_batch(batch, rank)[source]#
Return type:

Tuple[Tensor, Tensor]

skoots.lib.mp_utils.set_sharing_strategy(new_strategy=None)[source]#

https://pytorch.org/docs/stable/multiprocessing.html https://discuss.pytorch.org/t/how-does-one-setp-up-the-set-sharing-strategy-strategy-for-multiprocessing/113302 https://stackoverflow.com/questions/66426199/how-does-one-setup-the-set-sharing-strategy-strategy-for-multiprocessing-in-pyto

skoots.lib.mp_utils.setup_process(rank, world_size, port, backend='gloo')[source]#

Initialize the distributed environment (for each process).

gloo: is a collective communications library (https://github.com/facebookincubator/gloo). My understanding is that it’s a library/API for process to communicate/coordinate with each other/master. It’s a backend library.

export NCCL_SOCKET_IFNAME=eth0 export NCCL_IB_DISABLE=1

https://stackoverflow.com/questions/61075390/about-pytorch-nccl-error-unhandled-system-error-nccl-version-2-4-8

https://pytorch.org/docs/stable/distributed.html#common-environment-variables

skoots.lib.mp_utils.test_setup()[source]#
skoots.lib.mp_utils.use_file_system_sharing_strategy()[source]#

when to many file descriptor error happens

https://discuss.pytorch.org/t/how-does-one-setp-up-the-set-sharing-strategy-strategy-for-multiprocessing/113302

skoots.lib.merge#

skoots.lib.merge.get_adjacent_labels(x, y)[source]#

calculates which masks of a signle object have two labels (due to the border)

Parameters:

  • x (Tensor) –

  • y (Tensor) –

Return type:

List[Tuple[int, int]]

Returns:

skoots.train#

skoots.train.dataloader#

class skoots.train.dataloader.BackgroundDataset(path, transforms=<function BackgroundDataset.<lambda>>, device='cpu', sample_per_image=1)[source]#
cpu()[source]#

alias for self.to(‘cpu’)

cuda()[source]#

alias for self.to(‘cuda:0’)

map(fn, key)[source]#

applies a fn to an internal datastructure, provided by key. valid keys: [‘image’, ‘background’, ‘skele_masks’, ‘skeletons’]

Return type:

BackgroundDataset

mean()[source]#
numel()[source]#
std()[source]#
subtract_square_sum(other)[source]#

returns the sum of the entire dataset, each px subtracted by other :type other: :param other: :return:

sum()[source]#
to(device)[source]#

Sends all data stored in the dataloader to a device.

Parameters:

device (str) – torch device for images, masks, and skeletons

Returns:

self

class skoots.train.dataloader.MultiDataset(*args)[source]#

A utility class for joining multiple datasets into one accessible class. Sometimes, you may subdivide your training data based on some criteria. The most common is size: data from folder data/train/train_alot must be sampled 100 times per epoch, while data from folder data/train/train_notsomuch might only want to be sampled 1 times per epoch.

You could construct a two skoots.train.dataloader.dataset objects for each and access both in a single MultiDataset class…

>>> from skoots.train.dataloader import dataset
>>>
>>> # has one image sampled 100 times
>>> data0 = dataset('data/train/train_alot', sample_per_image=100)
>>> print(len(data0))  # 100
>>>
>>> # has one image sampled once
>>> data1 = dataset('data/train/train_notsomuch', sample_per_image=1)
>>> print(len(data1))  # 1
>>>
>>> merged_data = MultiDataset(data0, data1)
>>> print(len(merged_data))  # 101, they've been merged!
Parameters:

args –

cpu()[source]#

alias for self.to(‘cpu’)

Return type:

MultiDataset

cuda()[source]#

alias for self.to(‘cuda:0’)

Return type:

MultiDataset

map(fn, key)[source]#
Return type:

MultiDataset

mean(with_invert=False)[source]#
numel(with_invert=False)[source]#
std(with_invert=False)[source]#
sum(with_invert=False)[source]#
to(device)[source]#

Sends all data stored in the dataloader to a device. Occurs for ALL wrapped datasets.

Parameters:

device (str) – torch device for images, masks, and skeletons

Return type:

MultiDataset

Returns:

self

skoots.train.dataloader._sub_sq_sum(x, other)[source]#

subtracts other, sqares the result for each val, and sums to one number.

Parameters:

  • x (ndarray) – np.ndarray

  • other – number

Returns:

number

class skoots.train.dataloader.dataset(path, transforms=<function dataset.<lambda>>, pad_size=100, device='cpu', sample_per_image=1)[source]#

Custom dataset for loading and accessing skoots training data. This class loads data based on filenames and specific extensions: ‘.tif’ (raw image), ‘.labels.tif’ (instance masks), ‘.skeletons.tif’ (precomputed skeletons). An example training data folder might contain the following:

data\
 └  train\
      │ train_data.tif
      │ train_data.labels.tif
      └ train_data.skeletons.tif
Parameters:

  • path (Union[List[str], str]) – Path to training data

  • transforms (Optional[Callable[[Dict[str, Tensor]], Dict[str, Tensor]]]) – A function which applies dataset augmentation on a data_dict

  • pad_size (Optional[int]) – padding to add to every image in the dataset

  • device (Optional[str]) – torch.device which to output all data on

  • sample_per_image (Optional[int]) – number of times each image/mask pair is sampled per iteration over a dataset

cpu()[source]#

alias for self.to(‘cpu’)

Return type:

dataset

cuda()[source]#

alias for self.to(‘cuda:0’)

Return type:

dataset

map(fn, key)[source]#

applies a fn to an internal datastructure, provided by key. valid keys: [‘image’, ‘background’, ‘skele_masks’, ‘skeletons’]

Return type:

BackgroundDataset

mean(with_invert=False)[source]#
Return type:

float | None

numel(with_invert=False)[source]#
Return type:

int

pin_memory()[source]#

Pins underlying memory allowing faster transfer to GPU

Return type:

dataset

std(with_invert=False)[source]#
Return type:

float | None

subtract_square_sum(other)[source]#

returns the sum of the entire dataset, each px subtracted by other :type other: :param other: :return:

sum(with_invert=False)[source]#
Return type:

int

to(device)[source]#

Sends all data stored in the dataloader to a device.

Parameters:

device (str) – torch device for images, masks, and skeletons

Returns:

self

skoots.train.dataloader.skeleton_colate(data_dict)[source]#

Colate function with defines how we batch training data. Unpacks a data_dict with keys: ‘image’, ‘masks’, ‘skele_masks’, ‘baked_skeleton’, ‘skeleton’ and puts them each into a Tensor. This should not be called outright, rather passed to a torch.DataLoader for automatic batching.

Parameters:

data_dict (List[Dict[str, Tensor]]) – Dictonary of augmented training data

Return type:

Tuple[Tensor, Tensor, List[Dict[str, Tensor]], Tensor, Tensor]

Returns:

Tuple of batched data

skoots.train.distributed#

skoots.train.engine#

skoots.train.engine.train(rank, port, world_size, base_model, cfg, logging_level)[source]#

skoots.train.generate_skeletons#

skoots.train.generate_skeletons._calculate_skeletons(mask, scale)[source]#

image of shape [X, Y, Z,] with int masks returns dict[int, Tensor]

Parameters:

  • mask (Tensor) –

  • scale (Tensor) –

Return type:

Dict[int, Tensor]

Returns:

skoots.train.generate_skeletons.calculate_skeletons(mask, scale)[source]#

Calculates the skeleton of each object in mask

Parameters:

mask (Tensor) – [C, X, Y, Z]

Return type:

Dict[int, Tensor]

Returns:

Dict[int, Tensor] dict of masks where int is the object id and Tensor is [3, K] skeletons

skoots.train.generate_skeletons.create_gt_skeletons(base_dir, mask_filter, scale)[source]#
skoots.train.generate_skeletons.save_train_test_split(mask, skeleton, z_split, base)[source]#

Splits a volume of binary masks and skeletons. You CANNOT naively just split the mask in two as skeletons of objects on the border might not be properly calculated.

Saves pickled Dict[int, Tensor] to base+’_train.skeletons.trch’ and base+’_validate.skeletons.trch’

Parameters:

  • mask (Tensor) – Instance masks

  • skeleton (Dict[int, Tensor]) – Dict of skeleton of EVERY object in mask

  • z_split (int) – Z index of the train test split

  • base (str) – base filepath by which to save.

Returns:

None

skoots.train.loss#

skoots.train.loss._dice(pred, ground_truth, eps)[source]#
class skoots.train.loss.dice[source]#

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(predicted, ground_truth, eps=1e-10)[source]#

Returns dice index of two torch.Tensors

Parameters:

  • predicted (Tensor) –

    [B, I, X, Y, Z] torch.Tensor - probabilities calculated from hcat.utils.embedding_to_probability

    where B: is batch size, I: instances in image

  • ground_truth (Tensor) – [B, I, X, Y, Z] torch.Tensor - segmentation mask for each instance (I).

  • smooth – float - Very small number to ensure numerical stability. Default 1e-10

Return type:

Tensor

Returns:

dice_loss: [1] torch.Tensor - Result of Loss Function Calculation

class skoots.train.loss.jaccard[source]#

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(predicted, ground_truth, eps=1e-10)[source]#

Returns jaccard index of two torch.Tensors

Parameters:

  • predicted (Tensor) –

    [B, I, X, Y, Z] torch.Tensor - probabilities calculated from hcat.utils.embedding_to_probability

    where B: is batch size, I: instances in image

  • ground_truth (Tensor) – [B, I, X, Y, Z] torch.Tensor - segmentation mask for each instance (I).

  • eps (float) – float - Very small number to ensure numerical stability. Default 1e-10

Return type:

Tensor

Returns:

jaccard_loss: [1] torch.Tensor - Result of Loss Function Calculation

class skoots.train.loss.soft_cldice(iter_=3, smooth=1.0)[source]#

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(predicted, ground_truth)[source]#

Calculates the soft-clDice metric on a true and predicted value

Parameters:

  • ground_truth (Tensor) –

  • predicted (Tensor) –

Return type:

Tensor

Returns:

skoots.train.loss.soft_dice(predicted, ground_truth, smooth=1)[source]#

Computes the soft dice metric

Parameters:

  • ground_truth (Tensor) –

  • predicted (Tensor) –

  • smooth (int) – smoothing factor to prevent division by zero

Return type:

Tensor

Returns:

class skoots.train.loss.soft_dice_cldice(iter_=3, alpha=0.5, smooth=1.0)[source]#

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(predicted, ground_truth)[source]#

Calculates a singular loss value combining soft-Dice and soft-clDice which can be used to train a neural network

Parameters:

  • predicted (Tensor) – Input tensor

  • ground_truth (Tensor) – Ground Truth Tensor

Return type:

Tensor

Returns:

Single value which to perform a backwards pass

skoots.train.loss.soft_dilate(img)[source]#

approximates morphological operations through max_pooling for 2D and 3D

Return type:

Tensor

skoots.train.loss.soft_erode(img)[source]#

approximates morphological operations through max_pooling for 2D and 3D

Return type:

Tensor

skoots.train.loss.soft_open(img)[source]#

approximates morphological operations through max_pooling for 2D and 3D

Return type:

Tensor

skoots.train.loss.soft_skeletonize(img, iter_)[source]#

Performs a soft-skeletonization by terativly performing “soft morphological operations”

Parameters:

  • img (Tensor) – Image to perform operation on

  • iter – Number of times to perform the operation

Return type:

Tensor

Returns:

Soft-skeleton

class skoots.train.loss.split(n_iter=2, alpha=2.0, device='cpu')[source]#

The “oh shit my skeletons have split” loss. This basically checks if an edge has crossed the middle of a GT object. If it does, it applies a crazy loss.

Approximates the distance function by just eroding and adding a bunch. Approximates the edge function by subtracting the prediction by the eroded

For speed, will only check for pixels \(n_{iter}\) away. So if :math:`n_{iter} = 3’ the maximum distance any pixel might be from an edge would be 3.

Formally:

if \(E\) is the edge function and \(\Phi\) is the distance function, we compute the loss \(L(s, p)\) where \(s\) is the ground truth skeleton, and \(p\) is the predicted skeleton where

\[L(s, p) = E(p)^{ lpha \Phi(s)} - 1\]
Parameters:

  • n_iter (int) – Number of times to perform erosion for distance calculation.

  • alpha (float) – Scale factor for exponential loss. Large values penalize breakages more.

  • device (str) – a torch.device - ‘cuda’ or ‘cpu’

static _split_loss(edges, distance, a)[source]#
forward(pred, gt)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class skoots.train.loss.tversky(alpha, beta, eps)[source]#

Returns dice index of two torch.Tensors

Parameters:

  • alpha (float) – float - Value which penalizes False Positive Values

  • beta (float) – float - Value which penalizes False Negatives

  • eps (float) – float - Numerical stability term

static _tversky(pred, gt, alpha, beta, eps=1e-08)[source]#

tversky loss on per image basis.

Args:

pred: [N, X, Y, Z] Tensor of predicted segmentation masks (N instances) gt: [N, X, Y, Z] Tensor of ground truth segmentation masks (N instances) alpha: Penalty to false positives beta: Penalty to false negatives eps: stability parameter

Returns:

forward(predicted, ground_truth)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses. :rtype: Tensor

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

skoots.train.loss.tversky_graphable(pred, gt, alpha, beta)[source]#

skoots.train.merged_transform#

class skoots.train.merged_transform.BackgroundTransformFromCfg(cfg, device)[source]#

Initialize internal Module state, shared by both nn.Module and ScriptModule.

_crop1(image, masks, skeletons)[source]#
forward(data_dict)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class skoots.train.merged_transform.TransformFromCfg(cfg, device, scale=255.0)[source]#

Initialize internal Module state, shared by both nn.Module and ScriptModule.

_affine(image, masks, skeletons)[source]#
_brightness(image, masks)[source]#
_contrast(image, masks)[source]#
_crop1(image, masks, skeletons)[source]#
_crop2(image, masks, skeletons)[source]#
_elastic(image, masks, skeletons)[source]#
_flipX(image, masks, skeletons)[source]#
_flipY(image, masks, skeletons)[source]#
_flipZ(image, masks, skeletons)[source]#
_identity(*args)[source]#
_invert(image, masks)[source]#
_noise(image, masks)[source]#
_normalize(image, masks)[source]#
forward(data_dict)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses. :rtype: Dict[str, Tensor]

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

post_crop_fn(fn)[source]#
post_fn(fn)[source]#
pre_fn(fn)[source]#
set_dataset_mean(mean)[source]#
set_dataset_std(std)[source]#
skoots.train.merged_transform._compiled_affine(image, angle, shear, scale, translate)[source]#
skoots.train.merged_transform._compiled_grid_sample(a, grid, align_corners=True, mode='nearest')[source]#
skoots.train.merged_transform._elastic_on_skeletons(skeleton, grid, shape)[source]#
skoots.train.merged_transform._get_affine_matrix(center, angle, translate, scale, shear, device)[source]#
Return type:

Tensor

skoots.train.merged_transform._get_box(mask, device, threshold)[source]#
Return type:

Tuple[Tensor, Tensor]

skoots.train.merged_transform._get_inverse_affine_matrix(center, angle, translate, scale, shear, inverted=True)[source]#
Return type:

List[float]

skoots.train.merged_transform.calc_centroid(mask, id)[source]#
Return type:

Tensor

skoots.train.merged_transform.elastic_deform(*args, skeleton, displacement_shape=(6, 6, 2), displacement_magnitude=(0.05, 0.05, 0.01))[source]#

Randomly creates a deformation grid and applies to image, mask, and skeletons

Parameters:

  • args – Tensors of shape [1, 1, X, Y, Z]

  • skeleton (Dict[int, Tensor]) – Dict[int, Tensor[3, N]]

  • displacement_shape (Tuple[int, int, int]) – Tuple of len 3

  • max_displacement – float between 0 and 1

Return type:

Tuple[Tensor, Any, Dict[int, Tensor]]

Returns:

skoots.train.setup#

For code used in distributed training.

skoots.train.setup.cleanup(rank)[source]#

Destroy a given process group, and deinitialize the distributed package

skoots.train.setup.find_free_port()[source]#

https://stackoverflow.com/questions/1365265/on-localhost-how-do-i-pick-a-free-port-number

skoots.train.setup.get_batch(batch, rank)[source]#
Return type:

Tuple[Tensor, Tensor]

skoots.train.setup.set_sharing_strategy(new_strategy=None)[source]#

https://pytorch.org/docs/stable/multiprocessing.html https://discuss.pytorch.org/t/how-does-one-setp-up-the-set-sharing-strategy-strategy-for-multiprocessing/113302 https://stackoverflow.com/questions/66426199/how-does-one-setup-the-set-sharing-strategy-strategy-for-multiprocessing-in-pyto

skoots.train.setup.setup_process(rank, world_size, port, backend='gloo')[source]#

Initialize the distributed environment (for each process).

gloo: is a collective communications library (https://github.com/facebookincubator/gloo). My understanding is that it’s a library/API for process to communicate/coordinate with each other/master. It’s a backend library.

export NCCL_SOCKET_IFNAME=eth0 export NCCL_IB_DISABLE=1

https://stackoverflow.com/questions/61075390/about-pytorch-nccl-error-unhandled-system-error-nccl-version-2-4-8

https://pytorch.org/docs/stable/distributed.html#common-environment-variables

skoots.train.setup.test_setup()[source]#
skoots.train.setup.use_file_system_sharing_strategy()[source]#

when to many file descriptor error happens

https://discuss.pytorch.org/t/how-does-one-setp-up-the-set-sharing-strategy-strategy-for-multiprocessing/113302

skoots.train.sigma#

class skoots.train.sigma.Sigma(adjustments, initial_sigma=[0.1, 0.1, 0.8], device='cpu')[source]#

Creates an object which inputs an epoch, and returns a torch.tensor of [sigma_x, sigma_y, sigma_z]

Sigma should reflect the Error of embedding vectors at each spatial dim.

Parameters:

  • adjustments (List[Dict[str, float]]) – lost of adjustments and when to apply them

  • initial_sigma (List[float]) – initial values of sigma at epoch=0

  • device – device to load sigma on (‘cpu’ or ‘cuda’)

skoots.train.sigma.init_sigma(cfg, device)[source]#
Return type:

Sigma

skoots.train.utils#

skoots.train.utils.mask_overlay(image, mask, thr=0.5)[source]#
Return type:

Tensor

skoots.train.utils.show_box_pred(image, output, thr=0.9)[source]#
skoots.train.utils.sum_loss(input)[source]#
Return type:

Optional[Tensor]

skoots.train.utils.update_bn(loader, model, device=None)[source]#

Updates BatchNorm running_mean, running_var buffers in the model.

It performs one pass over data in loader to estimate the activation statistics for BatchNorm layers in the model. Args:

loader (torch.utils.data.DataLoader): dataset loader to compute the

activation statistics on. Each data batch should be either a tensor, or a list/tuple whose first element is a tensor containing data.

model (torch.nn.Module): model for which we seek to update BatchNorm

statistics.

device (torch.device, optional): If set, data will be transferred to

device before being passed into model.

Example:
>>> loader, model = ...
>>> torch.optim.swa_utils.update_bn(loader, model)

Note

The update_bn utility assumes that each data batch in loader is either a tensor or a list or tuple of tensors; in the latter case it is assumed that model.forward() should be called on the first element of the list or tuple corresponding to the data batch.

skoots.train.utils.write_progress(writer, tag, epoch, images, masks, probability_map, vector, out, skeleton=None, predicted_skeleton=None, gt_skeleton=None)[source]#

skoots.train.erosion#

skoots.train.erosion._compute_zero_padding(kernel_size)[source]#

Utility function that computes zero padding tuple. Adapted from Kornia

Return type:

Tuple[int, int, int]

skoots.train.erosion._get_binary_kernel3d(window_size, device)[source]#

Creates a symetric binary kernel to extract the patches. If the window size is HxWxD will create a (H*W)xHxW kernel.

Adapted from Kornia

Return type:

Tensor

skoots.train.erosion.binary_erosion(image)[source]#
Return type:

Tensor

class skoots.train.erosion.erosion(kernel_targets=3, rate=0.5, device='cpu')[source]#
_is_dict(data_dict)[source]#
Return type:

Dict[str, Tensor]

_is_tensor(input)[source]#

Assumes [C, X, Y, Z]

Return type:

Tensor

skoots.validate#

skoots.validate.compare#

skoots.validate.lib#

skoots.validate.lib._iou_instance_dict(a, b)[source]#

Given two instance masks, compares each instance in b against a. Usually assumes A is the ground truth.

Parameters:

  • a (Tensor) – Mask A

  • b (Tensor) – Mask B

Return type:

Dict[int, Tensor]

Returns:

Dict of instances and every IOU for each instance

skoots.validate.lib.accuracies_from_iou(iou, thr=0.1)[source]#
Return type:

Tensor

skoots.validate.lib.box_iou(a, b)[source]#

Compute the IoU of the cartesian product of two sets of boxes.

Each box in each set shall be (x0, y0, z0, x0, y0, z0).

Shapes:

a: \((6, N)\). b: \((6, M)\). returns: :math: (N, M)

Parameters:

  • a (Tensor) – box 1

  • b (Tensor) – box 2

Return type:

Tensor

Returns:

iou of each box in 1 against all boxes in 2

skoots.validate.lib.calculate_accuracies_from_bbox(ground_truth, predictions, device=None, threshold=0.1)[source]#

Calculates True positive, False Positive, False Negative from data_dict of segmentation 3d bboxes

Parameters:

  • ground_truth (Dict[str, Tensor]) –

  • predictions (Dict[str, Tensor]) –

  • device (Optional[str]) –

  • threshold –

Returns:

skoots.validate.lib.f1_score(tp, fp, fn)[source]#
skoots.validate.lib.get_segmentation_errors(ground_truth, predicted)[source]#

Calculates the IoU of each object on a per-mask-basis.

Parameters:

  • ground_truth (Tensor) – mask 1 with N instances

  • predicted (Tensor) – mask 2 with M instances

Return type:

float

Returns:

NxM matrix of IoU’s

skoots.validate.lib.mask_dice(gt, pred)[source]#

Calculates the Dice Index of each object on a per-mask-basis.

Parameters:

  • gt (Tensor) – mask 1 with N instances

  • pred (Tensor) – mask 2 with M instances

Returns:

NxM matrix of IoU’s

skoots.validate.lib.mask_iou(gt, pred)[source]#

Calculates the IoU of each object on a per-mask-basis.

Parameters:

  • gt (Tensor) – mask 1 with N instances

  • pred (Tensor) – mask 2 with M instances

Returns:

NxM matrix of IoU’s

skoots.validate.lib.mask_soft_cldice(gt, pred)[source]#

Calculates the Dice Index of each object on a per-mask-basis.

Parameters:

  • gt (Tensor) – mask 1 with N instances

  • pred (Tensor) – mask 2 with M instances

Returns:

NxM matrix of IoU’s

skoots.validate.lib.mask_to_bbox(mask)[source]#

Calculates the 3D bbox for each instance of an instance segmentation mask.

Assumes each positive integer is an instance for class label. Returns a tensor of id labels and a tensor of bboxes bboxes are in format: [x0,y0,z0,x1,y1,z1]

Assigns a bbox to each unique lablel! Does not mean each label has a valid bbox!!!

Shapes:

  • mask: \((1, X_{in}, Y_{in}, Z_{in})\)

  • return[0]: Id labels: :math: (N)

  • return[1]: bboxes: :math: (6, N)

Parameters:

mask (Tensor) – Input instance segmentation mask

Return type:

Tuple[Tensor, Tensor]

Returns:

id labels and bboxes

skoots.validate.lib.sparse_mask_iou(a, b)[source]#

Calculates the IoU of each object on a per-mask-basis using sparse tensors.

Parameters:

  • a (Tensor) – mask 1 with N instances

  • b (Tensor) – mask 2 with M instances

Return type:

Tensor

Returns:

NxM matrix of IoU’s

skoots.validate.lib.valid_box_inds(boxes)[source]#

returns the inds of all valid boxes

Parameters:

boxes – [6, N]

Returns:

[N]

skoots.validate.stats#

skoots.validate.utils#

skoots.validate.utils.imread(image_path, pin_memory=False)[source]#

Imports an image from file and returns in torch format

Parameters:

  • image_path (str) – path to image

  • pin_memory (Optional[bool]) – saves torch tensor in pinned memory if true

Return type:

Tensor

Returns: