swarmsort.embeddings module

GPU-accelerated embedding extractors for SwarmSort standalone package. Provides cupytexture and cupytexture_mega embeddings with automatic CPU fallback.

class swarmsort.embeddings.CupyShapeEmbedding(use_gpu=True, hog_cells=3, hog_bins=9)

Bases: EmbeddingExtractor

Shape-focused embedding optimized for microscopy organism tracking.

Key design principles (from embedding_recommendation.md): - NO RESIZE: Processes original patch resolution to preserve fine details - SHAPE-FOCUSED: Prioritizes contour/boundary features over texture - ROTATION-INVARIANT: Handles organisms at any angle

Feature composition (96 dimensions): - HOG features (36 dims): Histogram of Oriented Gradients for local shape - Fourier descriptors (16 dims): Contour shape in frequency domain - Radial signature (32 dims): Distance from centroid to boundary - Hu moments (7 dims): Classic rotation/scale invariant moments - Shape stats (5 dims): Area, perimeter, circularity, aspect ratio, solidity

This embedding is specifically designed for grayscale microscopy images where organisms have similar internal textures but distinct boundary shapes.

property embedding_dim: int

Return embedding dimensionality.

extract(frame, bbox)

Extract shape-focused embedding for a single patch.

Return type:

ndarray

extract_batch(frame, bboxes)

Extract batch of shape embeddings.

Return type:

List[ndarray]

class swarmsort.embeddings.CupyTextureColorEmbedding(patch_size=32, use_gpu=True)

Bases: EmbeddingExtractor

CuPy-accelerated texture embedding with rich color features.

Combines the robust texture analysis of CupyTextureEmbedding with comprehensive color-based features for improved tracking and re-identification performance.

Feature composition (84 total features): - Grayscale texture features (32) - from base CupyTextureEmbedding - RGB color histograms (24) - 8 bins per channel - HSV color histograms (24) - 8 bins per channel - Color moments (4) - mean and std of RGB channels

This embedding is particularly effective for scenarios where both texture and color information are important for distinguishing between objects.

property embedding_dim: int

Return embedding dimensionality.

extract(frame, bbox)

Extract features for a single patch.

Return type:

ndarray

extract_batch(frame, bboxes)

Extract batch of embeddings with GPU/CPU selection.

Return type:

List[ndarray]

extract_batch_cpu(frame, bboxes)

CPU batch extraction.

Return type:

List[ndarray]

extract_batch_gpu(frame, bboxes)

GPU batch extraction.

Return type:

List[ndarray]

class swarmsort.embeddings.CupyTextureEmbedding(patch_size=32, use_gpu=True)

Bases: EmbeddingExtractor

Fast GPU-accelerated embedding for microorganism tracking. Features: 36 dimensions - Basic statistics (4): mean, std, min, max - Gradient features (4): gradient magnitudes and variations - Shape features (8): center of mass, moments, eccentricity, orientation - Radial features (8): intensity at different radii from center - Color features (6): HSV statistics - Texture features (6): Local Binary Pattern approximation

property embedding_dim: int

Return embedding dimensionality.

extract(frame, bbox)

Extract features for a single patch.

Return type:

ndarray

extract_batch(frame, bboxes)

Extract batch of embeddings with GPU/CPU selection.

Return type:

List[ndarray]

extract_batch_cpu(frame, bboxes)

CPU batch extraction.

Return type:

List[ndarray]

extract_batch_gpu(frame, bboxes)

GPU batch extraction using CuPy.

Return type:

List[ndarray]

class swarmsort.embeddings.EmbeddingExtractor(use_tensor=False)

Bases: object

Base class for embedding extractors in standalone package.

property embedding_dim: int

Return embedding dimensionality.

extract(frame, bbox)

Extract single embedding.

Return type:

ndarray

extract_batch(frame, bboxes)

Extract batch of embeddings.

Return type:

List[ndarray]

class swarmsort.embeddings.MegaCupyTextureEmbedding(patch_size=32, use_gpu=True)

Bases: EmbeddingExtractor

Mega Microbe Embedding: Combines the best of shape, color, and texture analysis. - Dims 0-7: Basic stats and gradients - Dims 8-17: Advanced shape features (moments, eccentricity, Hu) - Dims 18-23: HSV color statistics (including circular stats for hue) - Dims 24-33: Rotation-invariant Local Binary Patterns (LBP) - Dims 34-49: Robust radial & polar features (intensity, variation, DFT) - Dims 50-55: Multi-scale wavelet and entropy texture features - Dims 56-63: Reserved / Other advanced features

property embedding_dim: int

Return embedding dimensionality.

extract(frame, bbox)

Extract single embedding.

Return type:

ndarray

extract_batch(frame, bboxes)

Extract batch of embeddings with GPU/CPU selection.

Return type:

List[ndarray]

extract_batch_cpu(frame, bboxes)

Return type:

List[ndarray]

extract_batch_gpu(frame, bboxes)

Return type:

List[ndarray]

swarmsort.embeddings.compute_embedding_distance(emb1, emb2, metric='cosine')

Compute distance between two embeddings.

Parameters:

• emb1 (ndarray) – First embedding vector

• emb2 (ndarray) – Second embedding vector

• metric (str) – Distance metric to use - “cosine” (default): Cosine distance, standard for L2-normalized embeddings - “correlation”: Pearson correlation distance, robust to offsets

Return type:

float

Returns:

Distance in [0, 1] where 0 = identical, 1 = maximally different

swarmsort.embeddings.compute_embedding_distances_batch(emb, embs, metric='cosine')

Compute distances from one embedding to multiple embeddings.

Parameters:

• emb (ndarray) – Query embedding vector

• embs (List[ndarray]) – List of embedding vectors to compare against

• metric (str) – Distance metric (“cosine” or “correlation”)

Return type:

ndarray

Returns:

Array of distances

swarmsort.embeddings.correlation_distance_jit(emb1, emb2)

JIT-optimized correlation (Pearson) distance computation.

Correlation distance = (1 - Pearson_correlation) / 2 Returns value in [0, 1] where 0 = identical pattern, 1 = opposite pattern.

More robust to offset differences than cosine, but may discard useful mean information from well-normalized embeddings.

Return type:

float

swarmsort.embeddings.cosine_distance_jit(emb1, emb2)

JIT-optimized cosine distance computation.

Cosine distance = 1 - cosine_similarity Returns value in [0, 1] where 0 = identical, 1 = opposite.

This is the standard distance metric for L2-normalized embeddings.

Return type:

float

swarmsort.embeddings.extract_features_batch_jit(patches, features_out)

JIT-optimized batch feature extraction for microorganism patches. patches: (N, H, W, 3) - batch of BGR patches features_out: (N, feature_dim) - output features

swarmsort.embeddings.extract_mega_features_batch_jit(patches, features_out)

JIT-optimized batch feature extraction combining shape, color, and texture. - Shape: Moments, eccentricity, orientation - Color: HSV statistics - Texture: Rotation-invariant LBP, multi-scale wavelet, entropy

swarmsort.embeddings.get_embedding_extractor(name, **kwargs)

Get an embedding extractor by name.

Return type:

EmbeddingExtractor

swarmsort.embeddings.is_gpu_available()

Check if GPU acceleration is available.

Return type:

bool

swarmsort.embeddings.list_available_embeddings()

List all available embeddings.

Return type:

List[str]