sleap-nn 0.1.0__py3-none-any.whl → 0.1.0a1__py3-none-any.whl

sleap_nn-0.1.0a1.dist-info/RECORD

@@ -0,0 +1,65 @@
+sleap_nn/.DS_Store,sha256=HY8amA79eHkt7o5VUiNsMxkc9YwW6WIPyZbYRj_JdSU,6148
+sleap_nn/__init__.py,sha256=l5Lwiad8GOurqkAhMwWw8-UcpH6af2TnMURf-oKj_U8,1362
+sleap_nn/cli.py,sha256=U4hpEcOxK7a92GeItY95E2DRm5P1ME1GqU__mxaDcW0,21167
+sleap_nn/evaluation.py,sha256=3u7y85wFoBgCwOB2xOGTJIDrd2dUPWOo4m0s0oW3da4,31095
+sleap_nn/legacy_models.py,sha256=8aGK30DZv3pW2IKDBEWH1G2mrytjaxPQD4miPUehj0M,20258
+sleap_nn/predict.py,sha256=8QKjRbS-L-6HF1NFJWioBPv3HSzUpFr2oGEB5hRJzQA,35523
+sleap_nn/system_info.py,sha256=7tWe3y6s872nDbrZoHIdSs-w4w46Z4dEV2qCV-Fe7No,14711
+sleap_nn/train.py,sha256=XvVhzMXL9rNQLx1-6jIcp5BAO1pR7AZjdphMn5ZX-_I,27558
+sleap_nn/architectures/__init__.py,sha256=w0XxQcx-CYyooszzvxRkKWiJkUg-26IlwQoGna8gn40,46
+sleap_nn/architectures/common.py,sha256=MLv-zdHsWL5Q2ct_Wv6SQbRS-5hrFtjK_pvBEfwx-vU,3660
+sleap_nn/architectures/convnext.py,sha256=l9lMJDxIMb-9MI3ShOtVwbOUMuwOLtSQlxiVyYHqjvE,13953
+sleap_nn/architectures/encoder_decoder.py,sha256=f3DUFJo6RrIUposdC3Ytyblr5J0tAeZ_si9dm_m_PhM,28339
+sleap_nn/architectures/heads.py,sha256=5E-7kQ-b2gsL0EviQ8z3KS1DAAMT4F2ZnEzx7eSG5gg,21001
+sleap_nn/architectures/model.py,sha256=1_dsP_4T9fsEVJjDt3er0haMKtbeM6w6JC6tc2jD0Gw,7139
+sleap_nn/architectures/swint.py,sha256=S66Wd0j8Hp-rGlv1C60WSw3AwGyAyGetgfwpL0nIK_M,14687
+sleap_nn/architectures/unet.py,sha256=rAy2Omi6tv1MNW2nBn0Tw-94Nw_-1wFfCT3-IUyPcgo,11723
+sleap_nn/architectures/utils.py,sha256=L0KVs0gbtG8U75Sl40oH_r_w2ySawh3oQPqIGi54HGo,2171
+sleap_nn/config/__init__.py,sha256=l0xV1uJsGJfMPfWAqlUR7Ivu4cSCWsP-3Y9ueyPESuk,42
+sleap_nn/config/data_config.py,sha256=5a5YlXm4V9qGvkqgFNy6o0XJ_Q06UFjpYJXmNHfvXEI,24021
+sleap_nn/config/get_config.py,sha256=rjNUffKU9z-ohLwrOVmJNGCqwUM93eh68h4KJfrSy8Y,42396
+sleap_nn/config/model_config.py,sha256=XFIbqFno7IkX0Se5WF_2_7aUalAlC2SvpDe-uP2TttM,57582
+sleap_nn/config/trainer_config.py,sha256=ZMXxns6VYakgYHRhkM541Eje76DdaTdDi4FFPNjJtP4,28413
+sleap_nn/config/training_job_config.py,sha256=v12_ME_tBUg8JFwOxJNW4sDQn-SedDhiJOGz-TlRwT0,5861
+sleap_nn/config/utils.py,sha256=GgWgVs7_N7ifsJ5OQG3_EyOagNyN3Dx7wS2BAlkaRkg,5553
+sleap_nn/data/__init__.py,sha256=eMNvFJFa3gv5Rq8oK5wzo6zt1pOlwUGYf8EQii6bq7c,54
+sleap_nn/data/augmentation.py,sha256=Kqw_DayPth_DBsmaO1G8Voou_-cYZuSPOjSQWSajgRI,13618
+sleap_nn/data/confidence_maps.py,sha256=PTRqZWSAz1S7viJhxu7QgIC1aHiek97c_dCUsKUwG1o,6217
+sleap_nn/data/custom_datasets.py,sha256=SO-aNB1-bB9DL5Zw-oGYDsliBxwI4iKX_FmwgZjKOgQ,99975
+sleap_nn/data/edge_maps.py,sha256=75qG_7zHRw7fC8JUCVI2tzYakIoxxneWWmcrTwjcHPo,12519
+sleap_nn/data/identity.py,sha256=7vNup6PudST4yDLyDT9wDO-cunRirTEvx4sP77xrlfk,5193
+sleap_nn/data/instance_centroids.py,sha256=SF-3zJt_VMTbZI5ssbrvmZQZDd3684bn55EAtvcbQ6o,2172
+sleap_nn/data/instance_cropping.py,sha256=2dYq5OTwkFN1PdMjoxyuMuHq1OEe03m3Vzqvcs_dkPE,8304
+sleap_nn/data/normalization.py,sha256=5xEvcguG-fvAGObl4nWPZ9TEM5gvv0uYPGDuni34XII,2930
+sleap_nn/data/providers.py,sha256=0x6GFP1s1c08ji4p0M5V6p-dhT4Z9c-SI_Aw1DWX-uM,14272
+sleap_nn/data/resizing.py,sha256=YFpSQduIBkRK39FYmrqDL-v8zMySlEs6TJxh6zb_0ZU,5076
+sleap_nn/data/utils.py,sha256=rT0w7KMOTlzaeKWq1TqjbgC4Lvjz_G96McllvEOqXx8,5641
+sleap_nn/inference/__init__.py,sha256=eVkCmKrxHlDFJIlZTf8B5XEOcSyw-gPQymXMY5uShOM,170
+sleap_nn/inference/bottomup.py,sha256=NqN-G8TzAOsvCoL3bttEjA1iGsuveLOnOCXIUeFCdSA,13684
+sleap_nn/inference/identity.py,sha256=GjNDL9MfGqNyQaK4AE8JQCAE8gpMuE_Y-3r3Gpa53CE,6540
+sleap_nn/inference/paf_grouping.py,sha256=7Fo9lCAj-zcHgv5rI5LIMYGcixCGNt_ZbSNs8Dik7l8,69973
+sleap_nn/inference/peak_finding.py,sha256=L9LdYKt_Bfw7cxo6xEpgF8wXcZAwq5plCfmKJ839N40,13014
+sleap_nn/inference/predictors.py,sha256=U114RlgOXKGm5iz1lnTfE3aN9S0WCh6gWhVP3KVewfc,158046
+sleap_nn/inference/provenance.py,sha256=0BekXyvpLMb0Vv6DjpctlLduG9RN-Q8jt5zDm783eZE,11204
+sleap_nn/inference/single_instance.py,sha256=rOns_5TsJ1rb-lwmHG3ZY-pOhXGN2D-SfW9RmBxxzcI,4089
+sleap_nn/inference/topdown.py,sha256=Ha0Nwx-XCH_rebIuIGhP0qW68QpjLB3XRr9rxt05JLs,35108
+sleap_nn/inference/utils.py,sha256=JnaJK4S_qLtHkWOSkHf4oRZjOmgnU9BGADQnntgGxxs,4689
+sleap_nn/tracking/__init__.py,sha256=rGR35wpSW-n5d3cMiQUzQQ_Dy5II5DPjlXAoPw2QhmM,31
+sleap_nn/tracking/track_instance.py,sha256=9k0uVy9VmpleaLcJh7sVWSeFUPXiw7yj95EYNdXJcks,1373
+sleap_nn/tracking/tracker.py,sha256=_WT-HFruzyOsvcq3AtLm3vnI9MYSwyBmq-HlQvj1vmU,41955
+sleap_nn/tracking/utils.py,sha256=uHVd_mzzZjviVDdLSKXJJ1T96n5ObKvkqIuGsl9Yy8U,11276
+sleap_nn/tracking/candidates/__init__.py,sha256=1O7NObIwshM7j1rLHmImbFphvkM9wY1j4j1TvO5scSE,49
+sleap_nn/tracking/candidates/fixed_window.py,sha256=D80KMlTnenuQveQVVhk9j0G8yx6K324C7nMLHgG76e0,6296
+sleap_nn/tracking/candidates/local_queues.py,sha256=Nx3R5wwEwq0gbfH-fi3oOumfkQo8_sYe5GN47pD9Be8,7305
+sleap_nn/training/__init__.py,sha256=vNTKsIJPZHJwFSKn5PmjiiRJunR_9e7y4_v0S6rdF8U,32
+sleap_nn/training/callbacks.py,sha256=TVnQ6plNC2MnlTiY2rSCRuw2WRk5cQSziek_VPUcOEg,25994
+sleap_nn/training/lightning_modules.py,sha256=G3c4xJkYWW-iSRawzkgTqkGd4lTsbPiMTcB5Nvq7jes,85512
+sleap_nn/training/losses.py,sha256=gbdinUURh4QUzjmNd2UJpt4FXwecqKy9gHr65JZ1bZk,1632
+sleap_nn/training/model_trainer.py,sha256=loCmEX0DfBtdV_pN-W8s31fn2_L-lbpWaq3OQXeSp-0,59337
+sleap_nn/training/utils.py,sha256=ivdkZEI0DkTCm6NPszsaDOh9jSfozkONZdl6TvvQUWI,20398
+sleap_nn-0.1.0a1.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
+sleap_nn-0.1.0a1.dist-info/METADATA,sha256=h3d4WPIu_JunY32jaRqJ4-fXp4KruTWT57FWb3L6dps,5637
+sleap_nn-0.1.0a1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+sleap_nn-0.1.0a1.dist-info/entry_points.txt,sha256=zfl5Y3hidZxWBvo8qXvu5piJAXJ_l6v7xVFm0gNiUoI,46
+sleap_nn-0.1.0a1.dist-info/top_level.txt,sha256=Kz68iQ55K75LWgSeqz4V4SCMGeFFYH-KGBOyhQh3xZE,9
+sleap_nn-0.1.0a1.dist-info/RECORD,,

{sleap_nn-0.1.0.dist-info → sleap_nn-0.1.0a1.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.10.2)
+Generator: setuptools (80.9.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

sleap_nn/data/skia_augmentation.py

@@ -1,414 +0,0 @@
-"""Skia-based augmentation functions that operate on uint8 tensors.
-
-This module provides augmentation functions using skia-python that:
-1. Match the exact API of sleap_nn.data.augmentation
-2. Operate on uint8 tensors throughout (avoiding float32 conversions)
-3. Provide ~1.5x faster augmentation compared to Kornia
-
-Usage:
-    from sleap_nn.data.skia_augmentation import (
-        apply_intensity_augmentation_skia,
-        apply_geometric_augmentation_skia,
-    )
-
-    # Apply augmentations (uint8 in, uint8 out)
-    image, instances = apply_intensity_augmentation_skia(image, instances, **config)
-    image, instances = apply_geometric_augmentation_skia(image, instances, **config)
-"""
-
-from typing import Optional, Tuple
-import numpy as np
-import torch
-import skia
-
-
-def apply_intensity_augmentation_skia(
-    image: torch.Tensor,
-    instances: torch.Tensor,
-    uniform_noise_min: float = 0.0,
-    uniform_noise_max: float = 0.04,
-    uniform_noise_p: float = 0.0,
-    gaussian_noise_mean: float = 0.02,
-    gaussian_noise_std: float = 0.004,
-    gaussian_noise_p: float = 0.0,
-    contrast_min: float = 0.5,
-    contrast_max: float = 2.0,
-    contrast_p: float = 0.0,
-    brightness_min: float = 1.0,
-    brightness_max: float = 1.0,
-    brightness_p: float = 0.0,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    """Apply intensity augmentations on uint8 image tensor.
-
-    Matches API of sleap_nn.data.augmentation.apply_intensity_augmentation.
-
-    Args:
-        image: Input tensor of shape (1, C, H, W) with dtype uint8 or float32.
-        instances: Keypoints tensor (not modified, just passed through).
-        uniform_noise_min: Minimum uniform noise (0-1 scale, maps to 0-255).
-        uniform_noise_max: Maximum uniform noise (0-1 scale).
-        uniform_noise_p: Probability of uniform noise.
-        gaussian_noise_mean: Gaussian noise mean (0-1 scale).
-        gaussian_noise_std: Gaussian noise std (0-1 scale).
-        gaussian_noise_p: Probability of Gaussian noise.
-        contrast_min: Minimum contrast factor.
-        contrast_max: Maximum contrast factor.
-        contrast_p: Probability of contrast adjustment.
-        brightness_min: Minimum brightness factor.
-        brightness_max: Maximum brightness factor.
-        brightness_p: Probability of brightness adjustment.
-
-    Returns:
-        Tuple of (augmented_image, instances). Image dtype matches input.
-    """
-    # Convert to numpy for Skia processing
-    is_float = image.dtype == torch.float32
-    if is_float:
-        img_np = (image[0].permute(1, 2, 0).numpy() * 255).astype(np.uint8)
-    else:
-        img_np = image[0].permute(1, 2, 0).numpy()
-
-    result = img_np.copy()
-
-    # Apply uniform noise (in uint8 space)
-    if uniform_noise_p > 0 and np.random.random() < uniform_noise_p:
-        noise = np.random.randint(
-            int(uniform_noise_min * 255),
-            int(uniform_noise_max * 255) + 1,
-            img_np.shape,
-            dtype=np.int16,
-        )
-        result = np.clip(result.astype(np.int16) + noise, 0, 255).astype(np.uint8)
-
-    # Apply Gaussian noise (in uint8 space)
-    if gaussian_noise_p > 0 and np.random.random() < gaussian_noise_p:
-        noise = np.random.normal(
-            gaussian_noise_mean * 255, gaussian_noise_std * 255, img_np.shape
-        ).astype(np.int16)
-        result = np.clip(result.astype(np.int16) + noise, 0, 255).astype(np.uint8)
-
-    # Apply contrast using lookup table (pure uint8)
-    if contrast_p > 0 and np.random.random() < contrast_p:
-        factor = np.random.uniform(contrast_min, contrast_max)
-        lut = np.arange(256, dtype=np.float32)
-        lut = np.clip((lut - 127.5) * factor + 127.5, 0, 255).astype(np.uint8)
-        result = lut[result]
-
-    # Apply brightness using lookup table (pure uint8)
-    if brightness_p > 0 and np.random.random() < brightness_p:
-        factor = np.random.uniform(brightness_min, brightness_max)
-        lut = np.arange(256, dtype=np.float32)
-        lut = np.clip(lut * factor, 0, 255).astype(np.uint8)
-        result = lut[result]
-
-    # Convert back to tensor
-    result_tensor = torch.from_numpy(result).permute(2, 0, 1).unsqueeze(0)
-    if is_float:
-        result_tensor = result_tensor.float() / 255.0
-
-    return result_tensor, instances
-
-
-def apply_geometric_augmentation_skia(
-    image: torch.Tensor,
-    instances: torch.Tensor,
-    rotation_min: float = -15.0,
-    rotation_max: float = 15.0,
-    rotation_p: Optional[float] = None,
-    scale_min: float = 0.9,
-    scale_max: float = 1.1,
-    scale_p: Optional[float] = None,
-    translate_width: float = 0.02,
-    translate_height: float = 0.02,
-    translate_p: Optional[float] = None,
-    affine_p: float = 0.0,
-    erase_scale_min: float = 0.0001,
-    erase_scale_max: float = 0.01,
-    erase_ratio_min: float = 1.0,
-    erase_ratio_max: float = 1.0,
-    erase_p: float = 0.0,
-    mixup_lambda_min: float = 0.01,
-    mixup_lambda_max: float = 0.05,
-    mixup_p: float = 0.0,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    """Apply geometric augmentations using Skia.
-
-    Matches API of sleap_nn.data.augmentation.apply_geometric_augmentation.
-
-    Args:
-        image: Input tensor of shape (1, C, H, W) with dtype uint8 or float32.
-        instances: Keypoints tensor of shape (1, n_instances, n_nodes, 2) or (1, n_nodes, 2).
-        rotation_min: Minimum rotation angle in degrees.
-        rotation_max: Maximum rotation angle in degrees.
-        rotation_p: Probability of rotation (independent). None = use affine_p.
-        scale_min: Minimum scale factor.
-        scale_max: Maximum scale factor.
-        scale_p: Probability of scaling (independent). None = use affine_p.
-        translate_width: Max horizontal translation as fraction of width.
-        translate_height: Max vertical translation as fraction of height.
-        translate_p: Probability of translation (independent). None = use affine_p.
-        affine_p: Probability of bundled affine transform.
-        erase_scale_min: Min proportion of image to erase.
-        erase_scale_max: Max proportion of image to erase.
-        erase_ratio_min: Min aspect ratio of erased area.
-        erase_ratio_max: Max aspect ratio of erased area.
-        erase_p: Probability of random erasing.
-        mixup_lambda_min: Min mixup strength (not implemented).
-        mixup_lambda_max: Max mixup strength (not implemented).
-        mixup_p: Probability of mixup (not implemented).
-
-    Returns:
-        Tuple of (augmented_image, augmented_instances). Image dtype matches input.
-    """
-    # Convert to numpy for Skia processing
-    is_float = image.dtype == torch.float32
-    if is_float:
-        img_np = (image[0].permute(1, 2, 0).numpy() * 255).astype(np.uint8)
-    else:
-        img_np = image[0].permute(1, 2, 0).numpy().copy()
-
-    h, w = img_np.shape[:2]
-    cx, cy = w / 2, h / 2
-
-    # Build transformation matrix
-    matrix = skia.Matrix()
-    has_transform = False
-
-    use_independent = (
-        rotation_p is not None or scale_p is not None or translate_p is not None
-    )
-
-    if use_independent:
-        if (
-            rotation_p is not None
-            and rotation_p > 0
-            and np.random.random() < rotation_p
-        ):
-            angle = np.random.uniform(rotation_min, rotation_max)
-            rot_matrix = skia.Matrix()
-            rot_matrix.setRotate(angle, cx, cy)
-            matrix = matrix.preConcat(rot_matrix)
-            has_transform = True
-
-        if scale_p is not None and scale_p > 0 and np.random.random() < scale_p:
-            scale = np.random.uniform(scale_min, scale_max)
-            scale_matrix = skia.Matrix()
-            scale_matrix.setScale(scale, scale, cx, cy)
-            matrix = matrix.preConcat(scale_matrix)
-            has_transform = True
-
-        if (
-            translate_p is not None
-            and translate_p > 0
-            and np.random.random() < translate_p
-        ):
-            tx = np.random.uniform(-translate_width, translate_width) * w
-            ty = np.random.uniform(-translate_height, translate_height) * h
-            trans_matrix = skia.Matrix()
-            trans_matrix.setTranslate(tx, ty)
-            matrix = matrix.preConcat(trans_matrix)
-            has_transform = True
-
-    elif affine_p > 0 and np.random.random() < affine_p:
-        angle = np.random.uniform(rotation_min, rotation_max)
-        scale = np.random.uniform(scale_min, scale_max)
-        tx = np.random.uniform(-translate_width, translate_width) * w
-        ty = np.random.uniform(-translate_height, translate_height) * h
-
-        matrix.setRotate(angle, cx, cy)
-        matrix.preScale(scale, scale, cx, cy)
-        matrix.preTranslate(tx, ty)
-        has_transform = True
-
-    # Apply geometric transform
-    if has_transform:
-        img_np = _transform_image_skia(img_np, matrix)
-        instances = _transform_keypoints_tensor(instances, matrix)
-
-    # Apply random erasing
-    if erase_p > 0 and np.random.random() < erase_p:
-        img_np = _apply_random_erase(
-            img_np, erase_scale_min, erase_scale_max, erase_ratio_min, erase_ratio_max
-        )
-
-    # Convert back to tensor
-    result_tensor = torch.from_numpy(img_np).permute(2, 0, 1).unsqueeze(0)
-    if is_float:
-        result_tensor = result_tensor.float() / 255.0
-
-    return result_tensor, instances
-
-
-def _transform_image_skia(image: np.ndarray, matrix: skia.Matrix) -> np.ndarray:
-    """Transform image using Skia matrix (uint8 in, uint8 out)."""
-    h, w = image.shape[:2]
-    channels = image.shape[2] if image.ndim == 3 else 1
-
-    # Skia needs RGBA
-    if channels == 1:
-        image_rgba = np.stack(
-            [image.squeeze()] * 3 + [np.full((h, w), 255, dtype=np.uint8)], axis=-1
-        )
-    elif channels == 3:
-        alpha = np.full((h, w, 1), 255, dtype=np.uint8)
-        image_rgba = np.concatenate([image, alpha], axis=-1)
-    else:
-        raise ValueError(f"Unsupported channels: {channels}")
-
-    image_rgba = np.ascontiguousarray(image_rgba, dtype=np.uint8)
-    skia_image = skia.Image.fromarray(
-        image_rgba, colorType=skia.ColorType.kRGBA_8888_ColorType
-    )
-
-    surface = skia.Surface(w, h)
-    canvas = surface.getCanvas()
-    canvas.clear(skia.Color4f(0, 0, 0, 1))
-    canvas.setMatrix(matrix)
-
-    paint = skia.Paint()
-    paint.setAntiAlias(True)
-    sampling = skia.SamplingOptions(skia.FilterMode.kLinear)
-    canvas.drawImage(skia_image, 0, 0, sampling, paint)
-
-    result = surface.makeImageSnapshot().toarray()
-
-    if channels == 1:
-        return result[:, :, 0:1]
-    return result[:, :, :channels]
-
-
-def _transform_keypoints_tensor(
-    keypoints: torch.Tensor, matrix: skia.Matrix
-) -> torch.Tensor:
-    """Transform keypoints tensor using Skia matrix."""
-    if keypoints.numel() == 0:
-        return keypoints
-
-    original_shape = keypoints.shape
-    flat = keypoints.reshape(-1, 2).numpy()
-
-    # Handle NaN values
-    valid_mask = ~np.isnan(flat).any(axis=1)
-    transformed = flat.copy()
-
-    if valid_mask.any():
-        valid_pts = flat[valid_mask]
-        skia_pts = [skia.Point(float(p[0]), float(p[1])) for p in valid_pts]
-        mapped = matrix.mapPoints(skia_pts)
-        transformed[valid_mask] = np.array([[p.x(), p.y()] for p in mapped])
-
-    return torch.from_numpy(transformed.reshape(original_shape).astype(np.float32))
-
-
-def _apply_random_erase(
-    image: np.ndarray,
-    scale_min: float,
-    scale_max: float,
-    ratio_min: float,
-    ratio_max: float,
-) -> np.ndarray:
-    """Apply random erasing (uint8)."""
-    h, w = image.shape[:2]
-    area = h * w
-
-    erase_area = np.random.uniform(scale_min, scale_max) * area
-    aspect_ratio = np.random.uniform(ratio_min, ratio_max)
-
-    erase_h = int(np.sqrt(erase_area * aspect_ratio))
-    erase_w = int(np.sqrt(erase_area / aspect_ratio))
-
-    if erase_h >= h or erase_w >= w:
-        return image
-
-    y = np.random.randint(0, h - erase_h)
-    x = np.random.randint(0, w - erase_w)
-
-    result = image.copy()
-    channels = image.shape[2] if image.ndim == 3 else 1
-    fill = np.random.randint(0, 256, size=(channels,), dtype=np.uint8)
-    result[y : y + erase_h, x : x + erase_w] = fill
-
-    return result
-
-
-def crop_and_resize_skia(
-    image: torch.Tensor,
-    boxes: torch.Tensor,
-    size: Tuple[int, int],
-) -> torch.Tensor:
-    """Crop and resize image regions using Skia.
-
-    Replacement for kornia.geometry.transform.crop_and_resize.
-
-    Args:
-        image: Input tensor of shape (1, C, H, W).
-        boxes: Bounding boxes tensor of shape (1, 4, 2) with corners:
-            [top-left, top-right, bottom-right, bottom-left].
-        size: Output size (height, width).
-
-    Returns:
-        Cropped and resized tensor of shape (1, C, out_h, out_w).
-    """
-    is_float = image.dtype == torch.float32
-    if is_float:
-        img_np = (image[0].permute(1, 2, 0).numpy() * 255).astype(np.uint8)
-    else:
-        img_np = image[0].permute(1, 2, 0).numpy()
-
-    h, w = img_np.shape[:2]
-    out_h, out_w = size
-    channels = img_np.shape[2] if img_np.ndim == 3 else 1
-
-    # Get box coordinates (top-left and bottom-right)
-    box = boxes[0].numpy()  # (4, 2)
-    x1, y1 = box[0]  # top-left
-    x2, y2 = box[2]  # bottom-right
-
-    crop_w = x2 - x1
-    crop_h = y2 - y1
-
-    # Create transformation matrix
-    matrix = skia.Matrix()
-    scale_x = out_w / crop_w
-    scale_y = out_h / crop_h
-    matrix.setScale(scale_x, scale_y)
-    matrix.preTranslate(-x1, -y1)
-
-    # Skia needs RGBA
-    if channels == 1:
-        image_rgba = np.stack(
-            [img_np.squeeze()] * 3 + [np.full((h, w), 255, dtype=np.uint8)], axis=-1
-        )
-    elif channels == 3:
-        alpha = np.full((h, w, 1), 255, dtype=np.uint8)
-        image_rgba = np.concatenate([img_np, alpha], axis=-1)
-    else:
-        raise ValueError(f"Unsupported channels: {channels}")
-
-    image_rgba = np.ascontiguousarray(image_rgba, dtype=np.uint8)
-    skia_image = skia.Image.fromarray(
-        image_rgba, colorType=skia.ColorType.kRGBA_8888_ColorType
-    )
-
-    surface = skia.Surface(out_w, out_h)
-    canvas = surface.getCanvas()
-    canvas.clear(skia.Color4f(0, 0, 0, 1))
-    canvas.setMatrix(matrix)
-
-    paint = skia.Paint()
-    paint.setAntiAlias(True)
-    sampling = skia.SamplingOptions(skia.FilterMode.kLinear)
-    canvas.drawImage(skia_image, 0, 0, sampling, paint)
-
-    result = surface.makeImageSnapshot().toarray()
-
-    if channels == 1:
-        result = result[:, :, 0:1]
-    else:
-        result = result[:, :, :channels]
-
-    result_tensor = torch.from_numpy(result).permute(2, 0, 1).unsqueeze(0)
-    if is_float:
-        result_tensor = result_tensor.float() / 255.0
-
-    return result_tensor

sleap_nn/export/__init__.py

@@ -1,21 +0,0 @@
-"""Export utilities for sleap-nn."""
-
-from sleap_nn.export.exporters import export_model, export_to_onnx, export_to_tensorrt
-from sleap_nn.export.metadata import ExportMetadata
-from sleap_nn.export.predictors import (
-    load_exported_model,
-    ONNXPredictor,
-    TensorRTPredictor,
-)
-from sleap_nn.export.utils import build_bottomup_candidate_template
-
-__all__ = [
-    "export_model",
-    "export_to_onnx",
-    "export_to_tensorrt",
-    "load_exported_model",
-    "ONNXPredictor",
-    "TensorRTPredictor",
-    "ExportMetadata",
-    "build_bottomup_candidate_template",
-]