tiling
sleap_nn.data.tiling
¶
Shared tiling primitives for splitting large images into overlapping tiles.
This module holds the low-level building blocks used to tile large frames into
fixed-size square tiles for training and inference. :func:generate_tile_grid
computes the deterministic grid of snapped tile origins in input-pixel space
used at inference time. The remaining primitives support random tile sampling
during training:
- :func:
frame_foreground_centersextracts valid keypoint (or centroid) locations to bias tiles toward foreground. - :func:
draw_tile_origindraws a single (possibly foreground-biased) tile origin for a frame. - :func:
extract_tilecrops a tile (with an optional geometric-augmentation halo) and returns tile-local instance coordinates. - :func:
tile_sample_seed, :func:tiling_worker_init_fn, :class:_FrameLRUand :class:FrameGroupedTileSamplerprovide the deterministic seeding, frame caching and frame-grouped batching used by the tiling dataset.
Classes:
| Name | Description |
|---|---|
FrameGroupedTileSampler |
Yield flat sample indices grouped into contiguous per-frame blocks. |
Functions:
| Name | Description |
|---|---|
draw_tile_origin |
Draw a single (optionally foreground-biased) tile origin for a frame. |
extract_tile |
Crop a fixed-size tile from a frame and return tile-local instances. |
frame_foreground_centers |
Extract valid foreground keypoint (or centroid) locations for a frame. |
generate_tile_grid |
Compute snapped square-tile top-left origins covering an image. |
tile_sample_seed |
Derive a deterministic per-sample seed for tile sampling. |
tiling_worker_init_fn |
Seed |
FrameGroupedTileSampler
¶
Bases: Sampler[int]
Yield flat sample indices grouped into contiguous per-frame blocks.
Each block holds the flat sample indices that belong to a single frame.
Blocks are kept contiguous (so all tiles of a frame are consumed together,
enabling frame caching), while the order of blocks can be shuffled. With
block_align enabled, each block is padded (by repeating its own indices)
up to a multiple of batch_size so batches never straddle frames. Blocks
are sharded across DDP replicas as whole units.
Methods:
| Name | Description |
|---|---|
__init__ |
Initialize the sampler. |
__iter__ |
Iterate flat sample indices, block by block. |
__len__ |
Return the number of indices yielded for the current epoch. |
set_epoch |
Set the current epoch to reseed shuffling. |
Source code in sleap_nn/data/tiling.py
458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 | |
__init__(frame_blocks, batch_size, shuffle, seed=0, block_align=True, shuffle_within_block=False, num_replicas=1, rank=0)
¶
Initialize the sampler.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
frame_blocks
|
List[List[int]]
|
List of blocks, each a list of flat sample indices for one frame. |
required |
batch_size
|
int
|
Batch size used for block-alignment padding. |
required |
shuffle
|
bool
|
If |
required |
seed
|
int
|
Base seed for shuffling. Defaults to |
0
|
block_align
|
bool
|
If |
True
|
shuffle_within_block
|
bool
|
If |
False
|
num_replicas
|
int
|
Number of DDP replicas the blocks are sharded across.
Defaults to |
1
|
rank
|
int
|
This process's replica rank. Defaults to |
0
|
Source code in sleap_nn/data/tiling.py
__iter__()
¶
Iterate flat sample indices, block by block.
Yields:
| Type | Description |
|---|---|
|
Flat sample indices with blocks kept contiguous, padded for block alignment, and (optionally) shuffled in order and within block. |
Source code in sleap_nn/data/tiling.py
__len__()
¶
Return the number of indices yielded for the current epoch.
Returns:
| Type | Description |
|---|---|
int
|
The total padded length summed over this rank's blocks, consistent
with :meth: |
Source code in sleap_nn/data/tiling.py
set_epoch(epoch)
¶
Set the current epoch to reseed shuffling.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
epoch
|
int
|
The epoch index. |
required |
draw_tile_origin(centers, frame_hw, tile_size, sample_k, samples_per_frame, tile_fg_fraction, center_jitter, rng, pos_ratio=1.0)
¶
Draw a single (optionally foreground-biased) tile origin for a frame.
The origin is returned unclamped; :func:extract_tile is responsible for
handling origins that fall partially (or fully) outside the frame via
zero-padding.
Foreground slots are assigned deterministically by index: the last
tile_fg_fraction of the samples_per_frame slots are foreground slots
(force_fg is True for them). Foreground slots draw a tile centered on
a random valid center (plus jitter); all other slots draw a uniformly random
origin. When there are no valid centers or pos_ratio is 0.0, every
slot falls back to a uniformly random origin.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
centers
|
Tensor
|
|
required |
frame_hw
|
Tuple[int, int]
|
Frame size as a |
required |
tile_size
|
int
|
Side length of the square tile in pixels. |
required |
sample_k
|
int
|
Index of this sample within the frame, in |
required |
samples_per_frame
|
int
|
Total number of tiles drawn per frame. |
required |
tile_fg_fraction
|
float
|
Fraction of the per-frame slots that are foreground slots (the trailing slots by index). |
required |
center_jitter
|
float
|
Jitter magnitude as a fraction of |
required |
rng
|
Generator
|
A seeded :class: |
required |
pos_ratio
|
float
|
If |
1.0
|
Returns:
| Type | Description |
|---|---|
Tuple[int, int]
|
The |
Source code in sleap_nn/data/tiling.py
extract_tile(image, instances, tile_origin, tile_size, *, apply_geometric=False, geometric_kwargs=None, symmetric_inds=None, rng_seed=None)
¶
Crop a fixed-size tile from a frame and return tile-local instances.
Two code paths are provided:
- Fast path (
apply_geometric=False, default): a pure tensor slice with constant-zero padding. The tile pixels are byte-identical to the source pixels (no resampling), even for origins that fall partially outside the frame. - Halo path (
apply_geometric=True): a larger halo (sideceil(tile_size * sqrt(2))) is extracted so that a subsequent geometric augmentation (e.g. rotation) has valid context, then the augmented halo is cropped back down totile_sizecentered on the halo center.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
image
|
Tensor
|
Sized-frame image of shape |
required |
instances
|
Tensor
|
Sized-frame instance keypoints of shape |
required |
tile_origin
|
Tuple[int, int]
|
The |
required |
tile_size
|
int
|
Side length of the square tile in pixels. |
required |
apply_geometric
|
bool
|
If |
False
|
geometric_kwargs
|
Optional[Dict]
|
Extra keyword arguments forwarded to
:func: |
None
|
symmetric_inds
|
Optional[Sequence[Tuple[int, int]]]
|
Symmetric node-index pairs forwarded to the geometric augmentation (for left/right flips). |
None
|
rng_seed
|
Optional[int]
|
If not |
None
|
Returns:
| Type | Description |
|---|---|
Tuple[Tensor, Tensor]
|
A tuple |
Source code in sleap_nn/data/tiling.py
262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 | |
frame_foreground_centers(instances, use_centroid=False)
¶
Extract valid foreground keypoint (or centroid) locations for a frame.
These locations are used by :func:draw_tile_origin to bias sampled tiles
toward regions of the frame that actually contain animals.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
instances
|
Tensor
|
Sized-frame instance keypoints of shape |
required |
use_centroid
|
bool
|
If |
False
|
Returns:
| Type | Description |
|---|---|
Tensor
|
A |
Source code in sleap_nn/data/tiling.py
generate_tile_grid(image_hw, tile_size, overlap, output_stride, max_stride=1, min_overlap_fraction=0.25)
¶
Compute snapped square-tile top-left origins covering an image.
Origins are computed independently for the height and width axes and then
combined via a Cartesian product in row-major (y then x) order. Each
origin is snapped to a multiple of output_stride and constrained so that
tiles never overrun the frame while still covering the far edges.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
image_hw
|
Tuple[int, int]
|
Image size as a |
required |
tile_size
|
int
|
Side length of the (square) tile in pixels. |
required |
overlap
|
int
|
Requested overlap between neighboring tiles in pixels. This is
raised to |
required |
output_stride
|
int
|
Tile origins are snapped down to multiples of this value. This is typically the output stride of the model so that tiles align to the network's prediction grid. |
required |
max_stride
|
int
|
Coarse stride used to snap the step between neighboring
tiles when it is at least as large as |
1
|
min_overlap_fraction
|
float
|
Minimum overlap expressed as a fraction of
|
0.25
|
Returns:
| Type | Description |
|---|---|
List[Tuple[int, int]]
|
A list of |
Notes
The union of the resulting tiles covers the frame up to the last origin
that can be placed on the output_stride grid without overrunning.
When (image_dim - tile_size) is a multiple of output_stride (the
typical case) this is the entire frame including the right/bottom edge.
Source code in sleap_nn/data/tiling.py
tile_sample_seed(base_seed, epoch, video_idx, frame_idx, sample_k, salt=0)
¶
Derive a deterministic per-sample seed for tile sampling.
Uses :class:numpy.random.SeedSequence so that the returned seed is a
stable, well-mixed function of all inputs (independent of worker/process).
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
base_seed
|
int
|
Global base seed for the run. |
required |
epoch
|
int
|
Current training epoch. |
required |
video_idx
|
int
|
Index of the source video. |
required |
frame_idx
|
int
|
Index of the source frame within the video. |
required |
sample_k
|
int
|
Index of the sample within the frame. |
required |
salt
|
int
|
Extra disambiguating value (e.g. to derive independent streams).
Defaults to |
0
|
Returns:
| Type | Description |
|---|---|
int
|
A 32-bit unsigned integer seed as a Python int. |
Source code in sleap_nn/data/tiling.py
tiling_worker_init_fn(worker_id)
¶
Seed numpy and torch per dataloader worker for tiling.
Reads base_seed off the worker's dataset (defaulting to 0) and mixes
it with the worker id so each worker gets an independent, reproducible
stream.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
worker_id
|
int
|
The dataloader worker id (as passed by PyTorch). |
required |