127 lines
4.0 KiB
Python
127 lines
4.0 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Test script for YOLO preprocessing of 16-bit TIFF images.
|
|
"""
|
|
|
|
import numpy as np
|
|
import tifffile
|
|
from pathlib import Path
|
|
import tempfile
|
|
import sys
|
|
import os
|
|
|
|
# Add parent directory to path to import modules
|
|
sys.path.insert(0, str(Path(__file__).parent.parent))
|
|
|
|
from src.model.yolo_wrapper import YOLOWrapper
|
|
from src.utils.image import Image
|
|
from PIL import Image as PILImage
|
|
|
|
|
|
def create_test_16bit_tiff(output_path: str) -> str:
|
|
"""Create a test 16-bit grayscale TIFF file.
|
|
|
|
Args:
|
|
output_path: Path where to save the test TIFF
|
|
|
|
Returns:
|
|
Path to the created TIFF file
|
|
"""
|
|
# Create a 16-bit grayscale test image (200x200)
|
|
# With values ranging from 0 to 65535 (full 16-bit range)
|
|
height, width = 200, 200
|
|
|
|
# Create a gradient pattern
|
|
test_data = np.zeros((height, width), dtype=np.uint16)
|
|
for i in range(height):
|
|
for j in range(width):
|
|
# Create a diagonal gradient
|
|
test_data[i, j] = int((i + j) / (height + width - 2) * 65535)
|
|
|
|
# Save as TIFF
|
|
tifffile.imwrite(output_path, test_data)
|
|
print(f"Created test 16-bit TIFF: {output_path}")
|
|
print(f" Shape: {test_data.shape}")
|
|
print(f" Dtype: {test_data.dtype}")
|
|
print(f" Min value: {test_data.min()}")
|
|
print(f" Max value: {test_data.max()}")
|
|
|
|
return output_path
|
|
|
|
|
|
def test_yolo_preprocessing():
|
|
"""Test YOLO preprocessing of 16-bit TIFF images."""
|
|
print("\n=== Testing YOLO Preprocessing of 16-bit TIFF ===")
|
|
|
|
# Create temporary test file
|
|
with tempfile.NamedTemporaryFile(suffix=".tif", delete=False) as tmp:
|
|
test_path = tmp.name
|
|
|
|
try:
|
|
# Create test image
|
|
create_test_16bit_tiff(test_path)
|
|
|
|
# Create YOLOWrapper instance (no actual model loading needed for this test)
|
|
print("\nTesting YOLOWrapper._prepare_source()...")
|
|
wrapper = YOLOWrapper()
|
|
|
|
# Call _prepare_source to preprocess the image
|
|
prepared_path, cleanup_path = wrapper._prepare_source(test_path)
|
|
|
|
print(f"\nPreprocessing complete:")
|
|
print(f" Original path: {test_path}")
|
|
print(f" Prepared path: {prepared_path}")
|
|
print(f" Cleanup path: {cleanup_path}")
|
|
|
|
# Verify the prepared image exists
|
|
assert os.path.exists(prepared_path), "Prepared image should exist"
|
|
|
|
# Load the prepared image and verify it's uint8 RGB
|
|
prepared_img = PILImage.open(prepared_path)
|
|
print(f"\nPrepared image properties:")
|
|
print(f" Mode: {prepared_img.mode}")
|
|
print(f" Size: {prepared_img.size}")
|
|
print(f" Format: {prepared_img.format}")
|
|
|
|
# Convert to numpy to check values
|
|
img_array = np.array(prepared_img)
|
|
print(f" Shape: {img_array.shape}")
|
|
print(f" Dtype: {img_array.dtype}")
|
|
print(f" Min value: {img_array.min()}")
|
|
print(f" Max value: {img_array.max()}")
|
|
print(f" Mean value: {img_array.mean():.2f}")
|
|
|
|
# Verify it's RGB uint8
|
|
assert prepared_img.mode == "RGB", "Prepared image should be RGB"
|
|
assert img_array.dtype == np.uint8, "Prepared image should be uint8"
|
|
assert img_array.shape[2] == 3, "Prepared image should have 3 channels"
|
|
assert (
|
|
0 <= img_array.min() <= img_array.max() <= 255
|
|
), "Values should be in [0, 255]"
|
|
|
|
# Cleanup prepared file if needed
|
|
if cleanup_path and os.path.exists(cleanup_path):
|
|
os.remove(cleanup_path)
|
|
print(f"\nCleaned up prepared image: {cleanup_path}")
|
|
|
|
print("\n✓ All YOLO preprocessing tests passed!")
|
|
return True
|
|
|
|
except Exception as e:
|
|
print(f"\n✗ Test failed with error: {e}")
|
|
import traceback
|
|
|
|
traceback.print_exc()
|
|
return False
|
|
|
|
finally:
|
|
# Cleanup
|
|
if os.path.exists(test_path):
|
|
os.remove(test_path)
|
|
print(f"Cleaned up test file: {test_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
success = test_yolo_preprocessing()
|
|
sys.exit(0 if success else 1)
|