Replace T.Normalize with T.ToDtype for improved data type handling in transformations

Author: rhoadesScholar

README.md

@@ -89,7 +89,7 @@ spatial_transforms = {
 
 # Define value transformations
 raw_value_transforms = T.Compose([
-    T.Normalize(mean=[0.0], std=[255.0]),           # Normalize to [0,1]
+    T.ToDtype(torch.float, scale=True),           # Normalize to [0,1] and convert to float
     GaussianNoise(std=0.05),          # Add noise for augmentation
     RandomContrast((0.8, 1.2)),       # Vary contrast
 ])
@@ -293,7 +293,7 @@ from cellmap_data.transforms import (
 
 # Input preprocessing
 raw_transforms = T.Compose([
-    T.Normalize(mean=[0.0], std=[255.0]),      # Normalize to [0,1]
+    T.ToDtype(torch.float, scale=True),      # Normalize to [0,1]
     GaussianNoise(std=0.1),      # Add noise
     RandomContrast((0.8, 1.2)),  # Vary contrast
     NaNtoNum({"nan": 0})         # Handle NaN values

docs/source/cellmap_data.transforms.augment.normalize.rst

@@ -108,14 +108,12 @@ def __init__(
         spatial_transforms: Optional[Mapping[str, Any]] = None,
         train_raw_value_transforms: Optional[T.Transform] = T.Compose(
             [
-                T.Normalize(mean=[0.0], std=[255.0]),
                 T.ToDtype(torch.float, scale=True),
                 NaNtoNum({"nan": 0, "posinf": None, "neginf": None}),
             ],
         ),
         val_raw_value_transforms: Optional[T.Transform] = T.Compose(
             [
-                T.Normalize(mean=[0.0], std=[255.0]),
                 T.ToDtype(torch.float, scale=True),
                 NaNtoNum({"nan": 0, "posinf": None, "neginf": None}),
             ],

src/cellmap_data/datasplit.py

@@ -192,7 +192,7 @@ def test_value_transforms_configuration(self, minimal_dataset_config):
 
         raw_transforms = T.Compose(
             [
-                T.Normalize(mean=[0.0], std=[255.0]),
+                T.ToDtype(torch.float, scale=True),
             ]
         )
 

tests/test_cellmap_dataset.py

@@ -286,7 +286,7 @@ def test_loader_with_transforms(self, tmp_path):
         input_arrays = {"raw": {"shape": (8, 8, 8), "scale": (4.0, 4.0, 4.0)}}
         target_arrays = {"gt": {"shape": (8, 8, 8), "scale": (4.0, 4.0, 4.0)}}
 
-        raw_transforms = T.Compose([T.Normalize(mean=[0.0], std=[255.0])])
+        raw_transforms = T.Compose([T.ToDtype(torch.float, scale=True)])
         target_transforms = T.Compose([Binarize(threshold=0.5)])
 
         dataset = CellMapDataset(

tests/test_dataloader.py

@@ -6,6 +6,7 @@
 
 import pytest
 import torchvision.transforms.v2 as T
+import torch
 
 from cellmap_data import CellMapDatasetWriter
 
@@ -290,7 +291,7 @@ def test_writer_with_value_transforms(self, tmp_path):
 
         output_path = tmp_path / "output.zarr"
 
-        raw_transform = T.Normalize(mean=[0.0], std=[255.0])
+        raw_transform = T.ToDtype(torch.float, scale=True)
 
         target_bounds = {
             "pred": {

tests/test_dataset_writer.py

@@ -43,7 +43,7 @@ def test_basic_training_setup(self, tmp_path):
 
         raw_transforms = T.Compose(
             [
-                T.Normalize(mean=[0.0], std=[255.0]),
+                T.ToDtype(torch.float, scale=True),
                 GaussianNoise(std=0.05),
             ]
         )
@@ -220,7 +220,7 @@ def test_complete_augmentation_pipeline(self, tmp_path):
         raw_transforms = T.Compose(
             [
                 NaNtoNum({"nan": 0.0}),
-                T.Normalize(mean=[0.0], std=[255.0]),
+                T.ToDtype(torch.float, scale=True),
                 GaussianNoise(std=0.05),
                 RandomContrast(contrast_range=(0.8, 1.2)),
                 RandomGamma(gamma_range=(0.8, 1.2)),
@@ -271,7 +271,7 @@ def test_per_target_transforms(self, tmp_path):
         # Different transforms for different targets
         target_transforms = {
             "labels": T.Compose([Binarize(threshold=0.5)]),
-            "distances": T.Compose([T.Normalize(mean=[0.0], std=[100.0])]),
+            "distances": T.Compose([T.ToDtype(torch.float, scale=True)]),
         }
 
         target_arrays = {

tests/test_integration.py

@@ -336,7 +336,7 @@ def test_sequential_transforms(self):
 
         transforms = T.Compose(
             [
-                T.Normalize(mean=[0.0], std=[255.0]),
+                T.ToDtype(torch.float32, scale=True),
                 GaussianNoise(std=0.01),
                 RandomContrast(contrast_range=(0.9, 1.1)),
             ]
@@ -356,7 +356,7 @@ def test_transform_pipeline(self):
         # Realistic preprocessing pipeline
         raw_transforms = T.Compose(
             [
-                T.Normalize(mean=[-128.0], std=[128.0]),  # Normalize around 0
+                T.ToDtype(torch.float32, scale=True),
                 GaussianNoise(std=0.05),
                 RandomContrast(contrast_range=(0.8, 1.2)),
             ]