Fix custom image subclass transforms

[siddharth10ss]

Fixes #1391

Summary

Implements the new_like() factory method to support custom Image subclasses with different constructor signatures, with comprehensive fallback mechanism and full spatial transform coverage.

Description

TorchIO spatial transforms previously reconstructed images using type(image)(tensor=..., affine=...), which failed for custom Image subclasses with different __init__ signatures (e.g., requiring additional parameters like history). This PR solves the issue by:

1. Adding new_like() method to the base Image class as a factory method for creating new instances
2. Fallback mechanism using copy.deepcopy() for subclasses that don't override new_like()
3. Updating ALL spatial transforms to use the new pattern
4. Comprehensive test suite with 12 test cases

Changes

Core Implementation (src/torchio/data/image.py)

• Added new_like(tensor, affine=None) method to base Image class
• Returns new instance with updated tensor/affine while preserving attributes
• Fallback to copy.deepcopy() for robust subclass support
• Subclasses can override to handle custom initialization

Transform Updates

Updated all spatial transforms to use new_like():

• Crop (src/torchio/transforms/preprocessing/spatial/crop.py)
• Pad (src/torchio/transforms/preprocessing/spatial/pad.py)
• CropOrPad (uses Crop internally, works automatically)
• Resample (src/torchio/transforms/preprocessing/spatial/resample.py)
• ToOrientation (src/torchio/transforms/preprocessing/spatial/to_orientation.py)
• Transpose (src/torchio/transforms/preprocessing/spatial/transpose.py)
• ToReferenceSpace (src/torchio/transforms/preprocessing/spatial/to_reference_space.py)

Added subject.update_attributes() calls for proper Subject-Image synchronization.

Test Suite (tests/transforms/preprocessing/test_image_factory.py)

12 comprehensive test cases:

• Custom subclass with required parameters
• Custom subclass with optional parameters
• Chained transforms preserving attributes
• All spatial transforms (Crop, Pad, Resample, etc.)
• Fallback mechanism verification
• Subject attribute synchronization
• Backward compatibility

Example Usage

import torchio as tio
import torch

class HistoryScalarImage(tio.ScalarImage):
    def __init__(self, tensor, affine, history, **kwargs):
        super().__init__(tensor=tensor, affine=affine, **kwargs)
        self.history = history
    
    def new_like(self, tensor, affine=None):
        """Override to preserve custom attributes."""
        return type(self)(
            tensor=tensor,
            affine=affine if affine is not None else self.affine,
            history=self.history,
            check_nans=self.check_nans,
            reader=self.reader,
        )

# Now works seamlessly with all transforms!
img = HistoryScalarImage(torch.rand(1,10,10,10), affine=torch.eye(4), history=["created"])
subject = tio.Subject(image=img)

transform = tio.Compose([
    tio.Crop(cropping=2),
    tio.Pad(padding=5),
    tio.Resample(target_shape=(64, 64, 64))
])

result = transform(subject)  #  Success! history preserved throughout
print(result.image.history)  # ["created"]

[Copilot]

Pull request overview

This PR implements a new_like() factory method in the base Image class to support custom Image subclasses with different constructor signatures. The issue arose when spatial transforms tried to reconstruct images using type(image)(tensor=..., affine=...), which failed for custom subclasses requiring additional parameters.

Changes:

• Added new_like() method to the base Image class with a fallback mechanism using copy.deepcopy() for subclasses that don't override it
• Updated all spatial transforms (Crop, Pad, Resample, ToOrientation, Transpose, ToReferenceSpace) to use new_like() instead of direct constructor calls
• Added comprehensive test coverage with two test files containing 12+ test cases covering various scenarios

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
src/torchio/data/image.py	Added new_like() factory method with try-except fallback mechanism and comprehensive documentation
src/torchio/transforms/preprocessing/spatial/crop.py	Updated to use new_like() for creating cropped images
src/torchio/transforms/preprocessing/spatial/pad.py	Updated to use new_like() and added subject.update_attributes() call
src/torchio/transforms/preprocessing/spatial/resample.py	Updated to use new_like() and added subject.update_attributes() call
src/torchio/transforms/preprocessing/spatial/to_orientation.py	Updated to use new_like() and added subject.update_attributes() call
src/torchio/transforms/preprocessing/spatial/transpose.py	Updated to use new_like() and added subject.update_attributes() call
src/torchio/transforms/preprocessing/spatial/to_reference_space.py	Updated build_image_from_reference() to use new_like() and added subject.update_attributes() call
tests/transforms/test_custom_image_subclass.py	Comprehensive test suite for custom Image subclasses with various transforms
tests/transforms/test_custom_image_subclass_fix.py	Additional test suite with similar coverage, potentially duplicating the other test file

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[siddharth10ss]

Update: All Feedback Addressed

I've completed the changes requested in the Copilot code review in commit 33903bc:

Changes Made:

1. Removed duplicate test file - Deleted tests/transforms/test_custom_image_subclass_fix.py
2. Fixed variable shadowing - In to_reference_space.py, renamed the shadowing variable:

   # Before: reference = downsample(reference)  
   # After:  downsampled_reference = downsample(reference)

[fepegar]

Hi @siddharth10ss, thanks for working on this! The new_like() factory method is a good pattern, and your PR correctly identified the problem and the solution direction.

I've opened a separate PR with a more minimal implementation:

• Add Image.new_like() for custom subclass support #1445

Here's why, and what's different:

The core insight is that only 2 call sites actually reconstruct images via type(image)(...):

1. Crop._crop_image (the copy_patch=True path)
2. build_image_from_reference in to_reference_space.py

The other transforms (Pad, Resample, ToOrientation, Transpose, ToReferenceSpace.apply_transform) use in-place mutation (image.set_data() + image.affine = ...), which already works fine for custom subclasses —
the object stays the same instance and type. Changing those from in-place to object-replacement is what caused the test_resample::test_affine CI failure in this PR, and it's also a semantic/behavioral change
that could break downstream code holding references to image objects.

What the new PR keeps from yours:

• The new_like(tensor, affine) method on Image — same API
• The override pattern for custom subclasses
• Comprehensive tests, including the exact snippet from Some transforms (e.g., Crop, Pad, CropOrPad) not compatible with custom Image subclasses with different constructors #1391

What's different:

• Only 3 source files changed (vs 8), ~60 insertions (vs 343)
• No copy.deepcopy fallback — subclasses must explicitly override new_like() (avoids surprise perf costs and the in-method import)
• Pad, Resample, ToOrientation, Transpose are untouched — no behavioral change
• Image.copy() also uses new_like() now, so copy.copy() works for subclasses too
• All tox checks pass (pytest, types, lint, format)

Thanks again for the initiative — your analysis of the problem was spot on, and the new_like idea from both you and the @tcollins-hub is the right call.

[fepegar]

Closing in favor of #1445.