tests/data/test_get_data_chunks.py (3)

Line range hint 111-112: Add tests for different scale values

While the scale parameter is set to 1.0, we should add test cases with different scale values to verify the resizing functionality works as expected.

 config.data_config.preprocessing.scale = 1.0
 config.data_config.preprocessing.is_rgb = True
+# Test with different scale values
+scales = [0.5, 2.0]
+for scale in scales:
+    config.data_config.preprocessing.scale = scale
+    samples = []
+    for idx, lf in enumerate(labels):
+        res = centered_instance_data_chunks(
+            (lf, idx),
+            data_config=config.data_config,
+            anchor_ind=0,
+            crop_size=(160, 160),
+            max_instances=2,
+            max_hw=max_hw,
+        )
+        samples.extend(res)
+    expected_size = int(226 * scale)
+    assert transform(samples[0]["instance_image"]).shape == (3, expected_size, expected_size)

---

Line range hint 209-260: Consider reducing test setup duplication

The transformation setup and RGB testing pattern is repeated across all test functions. Consider extracting common setup into fixtures or helper functions.

@pytest.fixture
def transform():
    return T.ToTensor()

@pytest.fixture
def rgb_config(config):
    config.data_config.preprocessing.is_rgb = True
    return config

def assert_image_shapes(sample, transform, is_rgb=False):
    channels = 3 if is_rgb else 1
    assert transform(sample["image"]).shape == (channels, 384, 384)

---

Line range hint 1-260: Consider enhancing test coverage for image transformations

While the changes successfully implement PIL image transformations, consider these improvements:

1. Add property-based tests to verify transformation behavior across different image sizes and scales
2. Add edge cases for extreme scale values
3. Consider testing memory usage to verify the PR's objective of reducing memory footprint

sleap_nn/data/instance_cropping.py (1)

50-55: LGTM! Consider consolidating the NaN handling logic.

The NaN handling logic is robust and well-implemented. However, the code could be slightly more concise while maintaining readability.

Consider this alternative implementation that reduces repetition:

-            diff_x = np.nanmax(pts[:, 0]) - np.nanmin(pts[:, 0])
-            diff_x = 0 if np.isnan(diff_x) else diff_x
-            max_length = np.maximum(max_length, diff_x)
-            diff_y = np.nanmax(pts[:, 1]) - np.nanmin(pts[:, 1])
-            diff_y = 0 if np.isnan(diff_y) else diff_y
-            max_length = np.maximum(max_length, diff_y)
+            diffs = [np.nanmax(pts[:, i]) - np.nanmin(pts[:, i]) for i in range(2)]
+            diffs = [0 if np.isnan(d) else d for d in diffs]
+            max_length = np.maximum(max_length, max(diffs))

tests/data/test_streaming_datasets.py (1)

35-35: LGTM! Consider adding shape calculation comments.

The addition of scale parameter and testing multiple scales (0.5 and 1.0) aligns well with the PR objective. The cleanup is properly handled.

Consider adding comments explaining how the output shapes are calculated, e.g.:

# Image shape: (1, 1, 200, 200)  # Original 400x400 with scale=0.5
# Confidence maps: (1, 2, 100, 100)  # Image/2 due to output_stride=2
# PAFs: (50, 50, 2)  # Image/4 due to output_stride=4

Also applies to: 64-81

sleap_nn/data/resizing.py (4)

Line range hint 1-152: Document image format requirements and consider PIL conversion handling

Given the PR's objective to convert between Tensor and PIL formats, consider:

1. Documenting the expected input/output image formats in function docstrings
2. Adding PIL conversion utilities in this module since it's the central image processing location
3. Ensuring tvf.resize behavior is consistent with both Tensor and PIL inputs

Consider adding these utility functions:

def to_pil_image(image: torch.Tensor) -> "PIL.Image":
    """Convert a torch tensor to PIL Image for storage efficiency."""
    return tvf.to_pil_image(image)

def to_tensor(image: "PIL.Image") -> torch.Tensor:
    """Convert a PIL Image back to torch tensor for processing."""
    return tvf.to_tensor(image)

---

Line range hint 249-286: Standardize error handling in SizeMatcher

The error handling in SizeMatcher.__iter__ uses Exception while similar checks in apply_sizematcher use ValueError. Consider standardizing to use ValueError consistently.

-                raise Exception(
+                raise ValueError(
                     f"Max height {self.max_height} should be greater than the current image height: {img_height}"
                 )
             if pad_width < 0:
-                raise Exception(
+                raise ValueError(
                     f"Max width {self.max_width} should be greater than the current image width: {img_width}"
                 )

---

Line range hint 31-63: Make padding more configurable

The padding functions use hardcoded parameters that could be made configurable:

1. Padding mode is always 'constant'
2. Padding value is always 0

Consider updating the function signature:

 def apply_pad_to_stride(
     image: torch.Tensor,
     max_stride: int,
+    pad_mode: str = "constant",
+    pad_value: float = 0.0,
 ) -> torch.Tensor:

---

Line range hint 89-106: Consider performance optimizations

The apply_resizer function could be optimized for batch processing:

1. Add batch support to process multiple images at once
2. Consider in-place operations where possible to reduce memory usage

Example batch implementation:

def apply_resizer_batch(
    images: torch.Tensor,
    instances: torch.Tensor,
    scale: float = 1.0
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Batch version of apply_resizer."""
    if scale != 1.0:
        # images shape: (batch_size, channels, height, width)
        images = resize_image(images, scale)
        instances = instances * scale
    return images, instances

sleap_nn/data/get_data_chunks.py (2)

7-7: Consider cleaning up unused imports.

Since normalization steps have been removed, consider cleaning up any unused imports from the normalization module to maintain code cleanliness.

Also applies to: 17-17

---

74-82: Document tensor dimension requirements.

The code assumes specific tensor dimensions for the image (presence of batch dimension that gets squeezed). Consider documenting the expected input tensor format in the function's docstring to prevent potential runtime errors.

tests/training/test_model_trainer.py (2)

Line range hint 296-301: Fix inconsistent checkpoint file references

There are still references to "best.ckpt" in the test cleanup code while the loading code has been updated to use "last.ckpt". This inconsistency should be resolved to maintain test reliability.

Apply this change to all cleanup sections:

-if (Path(centroid_config.trainer_config.save_ckpt_path) / "best.ckpt").exists():
+if (Path(centroid_config.trainer_config.save_ckpt_path) / "last.ckpt").exists():
     os.remove(
         (
-            Path(centroid_config.trainer_config.save_ckpt_path) / "best.ckpt"
+            Path(centroid_config.trainer_config.save_ckpt_path) / "last.ckpt"
         ).as_posix()
     )

Also applies to: 449-454

---

Line range hint 1-24: Add test coverage for new image conversion features

The PR introduces Tensor-to-PIL image conversion and relocates the resizer function, but there's no test coverage for these changes. Consider adding test cases to verify:

1. Correct conversion between Tensor and PIL formats
2. Proper functioning of the relocated resizer

Add new test cases like:

def test_image_format_conversion():
    """Test conversion between Tensor and PIL image formats."""
    # Test tensor to PIL conversion
    tensor_image = torch.rand(1, 3, 224, 224)
    pil_image = convert_tensor_to_pil(tensor_image)
    assert isinstance(pil_image, Image.Image)
    
    # Test PIL to tensor conversion
    tensor_result = convert_pil_to_tensor(pil_image)
    assert torch.allclose(tensor_image, tensor_result)

def test_resizer_in_get_chunks():
    """Test the relocated resizer function in get_chunks."""
    # Test resizing behavior
    original_size = (224, 224)
    target_size = (112, 112)
    chunks = get_chunks(...)  # Add appropriate test data
    assert chunks[0].size == target_size

Configuration used: CodeRabbit UI
Review profile: CHILL

#!/bin/bash
# Description: Verify PIL conversion integration

# Look for PIL/image conversion logic in the codebase
echo "Searching for PIL conversion logic..."
rg -l "from PIL import|Image.fromarray|torchvision.*transforms" --type python

# Look for image processing in the data loading pipeline
echo "Searching for image processing in data pipeline..."
rg -A 10 "def (load|process|transform).*image" --type python

#!/bin/bash
# Description: Check how scale parameter is used in data chunk functions

# Look for the data chunk function implementations
echo "Searching for data chunk function implementations..."
rg -A 10 "def (bottomup|centered_instance|centroid|single_instance)_data_chunks" --type python

# Look for image scaling/resizing logic
echo "Searching for image scaling implementation..."
rg -A 5 "scale.*image|resize.*image" --type python

#!/bin/bash
# Let's try a different approach to find scale parameter usage

# Search for scale parameter in function definitions and their implementations
echo "Searching for scale parameter usage..."
rg -A 10 "def.*scale.*:|scale\s*=\s*" --type-add 'python:*.py'

# Search specifically in the data chunks related files
echo "Searching in data chunks related files..."
fd -e py chunks | xargs rg -A 10 "scale"

# Look for image processing related functions
echo "Searching for image processing..."
rg -A 10 "(process|transform).*image" --type-add 'python:*.py'

#!/bin/bash
# Search for tensor shape expectations in the codebase
rg -A 2 "\.shape\s*==.*384.*384" --glob '!tests/**'

#!/bin/bash
# Search for shape-related assertions or checks in the main codebase
rg -A 2 "shape|size.*384" --glob '!tests/**'

# Search for tensor operations or model forward passes that might expect specific shapes
ast-grep --pattern 'def forward($$$) { $$$ }'

# Look for any configuration files that might specify expected dimensions
fd -e yaml -e json -e ini -e cfg -x cat {}

#!/bin/bash
# Search for tensor shape expectations in model forward passes and data processing
rg -A 2 "ToTensor|transform.*tensor" --glob '!tests/**'

# Look for any configuration files that specify input/output tensor shapes
rg -A 2 "shape.*384|384.*shape" --glob '!tests/**'

#!/bin/bash
# Search for PIL image conversion code that might interact with this function
rg -l "PIL" . | xargs rg "find_instance_crop_size"

# Search for test files that verify NaN handling
rg -l "test.*instance.*crop" . | xargs rg "nan"

#!/bin/bash
# Let's first find the function definition and its usage
ast-grep --pattern 'def find_instance_crop_size'

# Search for any PIL-related image processing in the same file
rg "PIL|Image" sleap_nn/data/instance_cropping.py

# Find tests related to instance cropping
fd "test.*instance.*crop.*\.py$"

#!/bin/bash
# Let's examine the test file to understand how NaN handling is tested
cat tests/data/test_instance_cropping.py

# Let's also look at the full implementation of instance_cropping.py
cat sleap_nn/data/instance_cropping.py

#!/bin/bash
# Check for PIL conversion tests
rg -l "PIL|Image\.fromarray|ToTensor" tests/
# Check for aspect ratio tests
rg -l "aspect_ratio|different.*width.*height" tests/
# Check for invalid scale tests
rg -l "invalid.*scale|negative.*scale|zero.*scale" tests/

#!/bin/bash
# Check for existing image-related tests
rg -l "test.*image" tests/

#!/bin/bash
# Search for other dtype conversions in the codebase
rg "to\(torch\.float32\)" --type py

#!/bin/bash
# Search for dtype conversions with different patterns
rg "\.to\(.*float" --type py -A 2

# Search for tensor type conversions
rg "\.type\(torch\." --type py -A 2

# Search for float dtype mentions in function signatures and docstrings
rg "dtype.*float|float.*dtype" --type py

#!/bin/bash
# Check for PIL usage in generate_bin_files.py
rg "PIL" --type python "generate_bin_files.py"

# Check for resize implementation in get_chunks
rg "get_chunks.*scale" --type python