lacss.ops

Ops on bounding-boxes

All functions here are degisned to work as either a numpy op or a jax op depending on the data type of the input.

box_area(box)

Computes area of boxes. Args: box: a float Tensor with [..., N, 2d].

Returns:

Type	Description
ArrayLike	a float Tensor with [..., N]

box_intersection(gt_boxes, boxes)

Compute pairwise intersection areas between boxes.

Parameters:

Name	Type	Description	Default
gt_boxes	ArrayLike	[..., N, 2d]	required
boxes	ArrayLike	[..., M, 2d]	required

Returns:

Type	Description
ArrayLike	a float Tensor with shape [..., N, M] representing pairwise intersections.

box_iou_similarity(gt_boxes, boxes)

Computes pairwise intersection-over-union between box collections.

Parameters:

Name	Type	Description	Default
gt_boxes	ArrayLike	a float Tensor with [..., N, 2d].	required
boxes	ArrayLike	a float Tensor with [..., M, 2d].	required

Returns:

Type	Description
ArrayLike	a Tensor with shape [..., N, M] representing pairwise iou scores.

distance_iou_loss(gt_boxes, boxes)

Loss_distance_iou = 1 - IoU + ho2(B, B_GT) / c2 The correction term is "distance_sq between box center" / "distance_sq of enclosing box cornor"

Parameters:

Name	Type	Description	Default
gt_boxes	ArrayLike	a float Tensor with [..., N, 2d].	required
boxes	ArrayLike	a float Tensor with [..., N, 2d].	required

Returns:

Type	Description
ArrayLike	a Tensor with shape [..., N] representing pairwise loss.

distance_similarity(pred_locations, gt_locations)

Compute distance similarity matrix

pairwise similarity = 1 / distance ^2

Parameters:

Name	Type	Description	Default
pred_locations	ArrayLike	[N, d] use -1 to mask out invalid locations	required
gt_locations	ArrayLike	[K, d] use -1 to mask out invalid locations	required

Returns:

Name	Type	Description
similarity_matrix	Array	[N, k]

feature_matching(features_a, features_b, threshold, *, similarity_fn=None)

Match predicted location to gt locations

Parameters:

Name	Type	Description	Default
features_a	ArrayLike	r [N, d], points or bboxes	required
features_b	ArrayLike	[K, d], points or bboxes	required
threshold	float	float. Min similarity score for match	required

Other Parameters:

Name	Type	Description
similarity_fn		optional custom function for computing similarity score fn(feautes_a, feature_b) -> similarity matrix

Returns:

Name	Type	Description
matches	Array	[N], indices of the matches row in b
indicators	Array	[N] bool

generalized_iou_loss(gt_boxes, boxes)

Loss_GIoU = 1 - IoU + |C - B union B_GT| / |C| where C is the smallest enclosing box for both B and B_GT. The resulting value has a gradient for non-overlapping boxes. See Zheng et al. [AAAI 2020]

Parameters:

Name	Type	Description	Default
gt_boxes	ArrayLike	a float Tensor with [..., N, 2d].	required
boxes	ArrayLike	a float Tensor with [..., N, 2d].	required

Returns:

Type	Description
ArrayLike	a Tensor with shape [..., N] representing pairwise loss.

iou_loss(gt_boxes, boxes)

IOU loss = 1 - IOU Args: gt_boxes: a float Tensor with [..., N, 2d]. boxes: a float Tensor with [..., N, 2d].

Returns:

Type	Description
ArrayLike	a Tensor with shape [..., N] representing pairwise loss.

match_and_replace(gt_features, pred_features, threshold, *, similarity_fn=None)

replacing gt_locations with pred_locations if the close enough 1. Each pred_location is matched to the closest gt_location 2. For each gt_location, pick the matched pred_location with highest score 3. if the picked pred_location is within threshold distance, replace the gt_location with the pred_location

yxhw_iou_similarity(yxhw_a, yxhw_b)

Computes pairwise IOU on bbox in yxhw format

Parameters:

Name	Type	Description	Default
gt_boxes		a float Tensor with [..., N, 2d].	required
boxes		a float Tensor with [..., M, 2d].	required

Returns:

Type	Description
	a Tensor with shape [..., N, M] representing pairwise iou scores.

sorbel_edges(images)

Returns a tensor holding Sobel edge maps.

Examples:

>>> image = random.uniform(key, shape=[3, 28, 28])
>>> sobel = sobel_edges(image)
>>> sobel_y = sobel[0, :, :, :] # sobel in y-direction
>>> sobel_x = sobel[1, :, :, :] # sobel in x-direction

Parameters:

Name	Type	Description	Default
image		[n, h, w]	required

Returns:

Type	Description
Array	Tensor holding edge maps for each channel. [2, n, h, w]

sorbel_edges_3d(images)

Returns a tensor holding Sobel edge maps in 3d.

Parameters:

Name	Type	Description	Default
images	ArrayLike	[n, d, h, w]	required

Returns:

Type	Description
Array	Tensor holding edge maps for each channel. [3, n, d, h, w]

sub_pixel_crop_and_resize(img, bbox, output_shape, out_of_bound_value=0)

Retrieve image values of a bbox. Resize output to output_shape. Used for ROI-Align.

Parameters:

Name	Type	Description	Default
img	ArrayLike	Array of shape [H, W, ...] (2D) or [D, H, W, ...] (3D)	required
bbox	ArrayLike	[y0, x0, y1, x1] (2D) or [z0, y0, x0, z1, y1, x1] (3D)	required
output_shape	tuple[int]	[h, w] (2D) or [d, h, w] (3D)	required
out_of_bound_value	float	optional float constant, defualt 0.	0

Returns:

Name	Type	Description
values	Array	[h, w, ...] or [d, h, w, ...]

sub_pixel_samples(img, locs, out_of_bound_value=0, edge_indexing=False)

Retrieve image values as non-integer locations by interpolation

Parameters:

Name	Type	Description	Default
img	ArrayLike	Array of shape [D1,D2,..,Dk, ...]	required
locs	ArrayLike	Array of shape [d1,d2,..,dn, k]	required
out_of_bound_value	float	optional float constant, defualt 0.	0
edge_indexing	bool	if True, the index for the top/left pixel is 0.5, otherwise 0. Default is False	False

Returns:

Name	Type	Description
values	Array	[d1,d2,..,dn, ...], float

non_max_suppression(scores, boxes, max_output_size, threshold=0.5, min_score=0, return_selection=False, similarity_func=None)

non-maximum suppression for either bboxes or points.

Assumption:

* The boxes/points are sorted by scores

The overal design of the algorithm is to handle boxes tile-by-tile:

Parameters:

Name	Type	Description	Default
scores	ArrayLike	[N]	required
boxes	ArrayLike	[N, C] C=4 for boxes, C=2/3 for locations	required
max_output_size	int	a positive scalar integer	required
threshold	float	threshold of similarity score to supress	0.5
min_score	float	min score to be selected, default 0	0
return_selection	bool	whether also return the boolean indicator	False
similarity_func	callable | None	Optionally provide a custom callable to compute similarity score	None

Returns:

Name	Type	Description
nms_scores	Array | tuple[Array]	[M]. M = max_output_size
nms_proposals	Array | tuple[Array]	[M, C].
selection	Array | tuple[Array]	[N] a boolean indicator of selection status of original input only if return_selection is True

Various functions deals with segmentation pathces

All functions here takes unbatched input. Use vmap to convert to batched data

bboxes_of_patches(pred, threshold=0, *, image_shape=None, is2d=None)

Compute the instance bboxes from model predictions

Parameters:

Name	Type	Description	Default
pred	dict	A model prediction dictionary:	required
threshold	float	for segmentation, default 0	0

Keyward Args

image_shape: if not None, the boxes are clipped to be within the bound is2d: whether to force 2d output

Returns:

Name	Type	Description
bboxes	ndarray	[n, 4] or [n, 6]

coords_of_patches(preds, image_shape)

get the zyx coordinates of segmentations

Parameters:

Name	Type	Description	Default
preds	dict	the model prediction dictionary	required
image_shape	tuple[int, ...]	the original input image shape. (H, W) for 2d and (D, H, W) for 3d	required

a tuple of (coordinates, boolean_masks)

Name	Type	Description
coordinates	Array	integer tensor of shape: (3,) + preds['segmentations'].shape
boolan_masks	Array	boolean tensor indicating whether the coordinate is a real one of padding.

crop_and_resize_patches(pred, bboxes, *, target_shape=(48, 48), convert_logits=False)

crop and rescale all instances to a target_size

Parameters:

Name	Type	Description	Default
pred	dict	model predictions	required
bboxes		optionally supply bboxes for cropping	required

Keyward Args: target_shape: output shape. usually a 3-tuple but can be a 2-tuple if input is a 2D image. convert_logits: whether to convert the logits to probability

Returns:

Type	Description
	Array [N] + target_shape.

gather_patches(source, locations, patch_size, *, padding_value=0)

gather patches from an array

Parameters:

Name	Type	Description	Default
source	Array	[D1..dk, ...]	required
locations	ArrayLike	[N, k] top-left coords of the patches, negative (out-of-bound) is ok	required
patch_size	int | tuple[int]	size of the patch.	required

Keyward Args: padding_value: value for out-of-bound locations

Returns:

Type	Description
	N patches [N, d1..dk, ...]

patches_to_label(pred, input_size, *, mask=None, score_threshold=0.5, threshold=0)

convert patch output to the image label.

Parameters:

Name	Type	Description	Default
pred	DataDict	A model prediction dictionary	required
input_size	Shape	shape of the input image. (H, W) or (D, H, W)	required

Keyward Args: mask: boolean indicators masking out unwanted instances. Default is None (all cells) score_threshold: otional, min_score to be included. Default is .5. threshold: segmentation threshold.

Returns:

Name	Type	Description
label	Array	of the dimension input_size

rescale_patches(pred, scale, *, transform_logits=True)

Rescale/resize instance outputs in a sub-pixel accurate way. If the input image was rescaled, this function take care of rescaling the predictions to the orginal coodinates.

Parameters:

Name	Type	Description	Default
pred	DataDict	A model prediction dictionary	required
scale	float	The scaling value. The function does not take a noop shortcut even if scale is 1.	required

A tuple of three arrays

Name	Type	Description
patches	Array	a 3D array of the rescaled segmentation patches. The array shape should be different from the orignal patches in model predition.
yy	Array	The y-coordinates of the patches in mesh-grid format
xx	Array	The x-coordinates of the patches in mesh-grid format