Bug fix

tests/test_training_dataset_prep.py

@@ -0,0 +1,117 @@
+#!/usr/bin/env python3
+"""
+Test script for training dataset preparation with 16-bit TIFFs.
+"""
+
+import numpy as np
+import tifffile
+from pathlib import Path
+import tempfile
+import sys
+import os
+import shutil
+
+# Add parent directory to path to import modules
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+from src.utils.image import Image
+
+
+def test_float32_3ch_conversion():
+    """Test conversion of 16-bit TIFF to float32 3-channel TIFF."""
+    print("\n=== Testing Float32 3-Channel Conversion ===")
+
+    # Create temporary directory structure
+    with tempfile.TemporaryDirectory() as tmpdir:
+        tmpdir = Path(tmpdir)
+        src_dir = tmpdir / "original"
+        dst_dir = tmpdir / "converted"
+        src_dir.mkdir()
+        dst_dir.mkdir()
+
+        # Create test 16-bit TIFF
+        test_data = np.zeros((100, 100), dtype=np.uint16)
+        for i in range(100):
+            for j in range(100):
+                test_data[i, j] = int((i + j) / 198 * 65535)
+
+        test_file = src_dir / "test_16bit.tif"
+        tifffile.imwrite(test_file, test_data)
+        print(f"Created test 16-bit TIFF: {test_file}")
+        print(f"  Shape: {test_data.shape}")
+        print(f"  Dtype: {test_data.dtype}")
+        print(f"  Range: [{test_data.min()}, {test_data.max()}]")
+
+        # Simulate the conversion process
+        print("\nConverting to float32 3-channel...")
+        img_obj = Image(test_file)
+
+        # Convert to float32 [0-1]
+        float_data = img_obj.to_normalized_float32()
+
+        # Replicate to 3 channels
+        if len(float_data.shape) == 2:
+            float_3ch = np.stack([float_data] * 3, axis=-1)
+        else:
+            float_3ch = float_data
+
+        # Save as float32 TIFF
+        output_file = dst_dir / "test_float32_3ch.tif"
+        tifffile.imwrite(output_file, float_3ch.astype(np.float32))
+        print(f"Saved float32 3-channel TIFF: {output_file}")
+
+        # Verify the output
+        loaded = tifffile.imread(output_file)
+        print(f"\nVerifying output:")
+        print(f"  Shape: {loaded.shape}")
+        print(f"  Dtype: {loaded.dtype}")
+        print(f"  Channels: {loaded.shape[2] if len(loaded.shape) == 3 else 1}")
+        print(f"  Range: [{loaded.min():.6f}, {loaded.max():.6f}]")
+        print(f"  Unique values: {len(np.unique(loaded[:,:,0]))}")
+
+        # Assertions
+        assert loaded.dtype == np.float32, f"Expected float32, got {loaded.dtype}"
+        assert loaded.shape[2] == 3, f"Expected 3 channels, got {loaded.shape[2]}"
+        assert (
+            0.0 <= loaded.min() <= loaded.max() <= 1.0
+        ), f"Expected [0,1] range, got [{loaded.min()}, {loaded.max()}]"
+
+        # Verify all channels are identical (replicated grayscale)
+        assert np.array_equal(
+            loaded[:, :, 0], loaded[:, :, 1]
+        ), "Channel 0 and 1 should be identical"
+        assert np.array_equal(
+            loaded[:, :, 0], loaded[:, :, 2]
+        ), "Channel 0 and 2 should be identical"
+
+        # Verify float32 precision (not quantized to uint8 steps)
+        unique_vals = len(np.unique(loaded[:, :, 0]))
+        print(f"\n  Precision check:")
+        print(f"    Unique values in channel: {unique_vals}")
+        print(f"    Source unique values: {len(np.unique(test_data))}")
+
+        # The final unique values should match source (no loss from conversion)
+        assert unique_vals == len(
+            np.unique(test_data)
+        ), f"Expected {len(np.unique(test_data))} unique values, got {unique_vals}"
+
+        print("\n✓ All conversion tests passed!")
+        print("  - Float32 dtype preserved")
+        print("  - 3 channels created")
+        print("  - Range [0-1] maintained")
+        print("  - No precision loss from conversion")
+        print("  - Channels properly replicated")
+
+        return True
+
+
+if __name__ == "__main__":
+    try:
+        success = test_float32_3ch_conversion()
+        sys.exit(0 if success else 1)
+    except Exception as e:
+        print(f"\n✗ Test failed with error: {e}")
+        import traceback
+
+        traceback.print_exc()
+        sys.exit(1)