Data Transforms API Flow#

This gives the rough flow of data outlining the data augmetnation procedure to train a SKOOTS model.

_images/skoots_transforms_api_flow.pdf _images/skoots_transforms_api_flow_inverted.pdf
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

pin_memory()[source]

Pins underlying memory allowing faster transfer to GPU

Return type:

dataset

subtract_square_sum(other)[source]

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

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.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.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