Files
object-segmentation/tests/test_training_dataset_prep.py
2025-12-13 01:18:16 +02:00

118 lines
4.0 KiB
Python

#!/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)