import tempfile import unittest from pathlib import Path import cv2 import numpy as np import tifffile from skimage.util import img_as_ubyte from cleaning import ( INPAINT_RADIUS, _as_display_rgb_uint8, _as_rgb_uint8, clean_display_image_and_mask, correct_vignetting, create_debris_mask, create_tissue_mask, read_image_rgb_and_preview, write_image_rgb, ) class PreviewConversionTests(unittest.TestCase): def test_uint16_preview_uses_full_dtype_range(self): image = np.array([[0, 32768, 65535]], dtype=np.uint16) preview = _as_display_rgb_uint8(image) expected = np.array([0, 128, 255], dtype=np.uint8) np.testing.assert_array_equal(preview[:, :, 0], expected[None, :]) np.testing.assert_array_equal(preview[:, :, 1], expected[None, :]) np.testing.assert_array_equal(preview[:, :, 2], expected[None, :]) def test_uint8_preview_is_unchanged(self): image = np.array([[[10, 120, 240]]], dtype=np.uint8) preview = _as_display_rgb_uint8(image) np.testing.assert_array_equal(preview, image) def test_reader_returns_separate_processing_and_display_views(self): image = np.full((10, 10, 3), 32768, dtype=np.uint16) image[0, 0] = 0 image[-1, -1] = 65535 with tempfile.TemporaryDirectory() as directory: image_path = Path(directory) / "image.tiff" tifffile.imwrite(image_path, image, photometric="rgb") processing, display = read_image_rgb_and_preview(image_path) np.testing.assert_array_equal(processing, _as_rgb_uint8(image)) np.testing.assert_array_equal(display, img_as_ubyte(image)) def test_cleaned_output_uses_display_brightness_and_processing_mask(self): raw = np.full((32, 32, 3), 32768, dtype=np.uint16) raw[12:20, 12:20] = 0 processing = _as_rgb_uint8(raw) display = _as_display_rgb_uint8(raw) cleaned, debris_mask = clean_display_image_and_mask( processing, display, ) expected_mask = create_debris_mask(processing) expected_cleaned = cv2.inpaint( display, expected_mask, INPAINT_RADIUS, cv2.INPAINT_TELEA, ) np.testing.assert_array_equal(debris_mask, expected_mask) np.testing.assert_array_equal(cleaned, expected_cleaned) def test_vignetting_correction_flattens_detected_edge_falloff(self): size = 160 axis = np.linspace(-1, 1, size, dtype=np.float32) x, y = np.meshgrid(axis, axis) illumination = 1 - 0.18 * (x * x + y * y) base_color = np.array([230, 220, 210], dtype=np.float32) image = np.clip( illumination[..., None] * base_color, 0, 255, ).astype(np.uint8) corrected, was_corrected = correct_vignetting(image) center = np.s_[60:100, 60:100] border = np.zeros((size, size), dtype=bool) border[:20] = True border[-20:] = True border[:, :20] = True border[:, -20:] = True before_gap = image[center].mean() - image[border].mean() after_gap = corrected[center].mean() - corrected[border].mean() self.assertTrue(was_corrected) self.assertLess(abs(after_gap), abs(before_gap) * 0.25) def test_uniform_image_skips_vignetting_correction(self): image = np.full((128, 128, 3), 220, dtype=np.uint8) corrected, was_corrected = correct_vignetting(image) self.assertFalse(was_corrected) np.testing.assert_array_equal(corrected, image) def test_vignetting_correction_handles_steep_one_sided_falloff(self): height, width = 180, 240 x = np.linspace(0, 1, width, dtype=np.float32) illumination = 1 - 0.35 * np.clip((x - 0.72) / 0.28, 0, 1) base_color = np.array([245, 240, 235], dtype=np.float32) image = np.clip( illumination[None, :, None] * base_color, 0, 255, ).astype(np.uint8) image = np.repeat(image, height, axis=0) corrected, was_corrected = correct_vignetting(image) center_level = image[:, 80:140].mean() edge_level = image[:, -12:].mean() corrected_center_level = corrected[:, 80:140].mean() corrected_edge_level = corrected[:, -12:].mean() before_gap = center_level - edge_level after_gap = corrected_center_level - corrected_edge_level self.assertTrue(was_corrected) self.assertLess(after_gap, before_gap * 0.15) def test_vignetting_correction_preserves_stained_tissue(self): height, width = 180, 240 x = np.linspace(0, 1, width, dtype=np.float32) illumination = 1 - 0.30 * np.clip((x - 0.65) / 0.35, 0, 1) background_color = np.array([245, 240, 235], dtype=np.float32) image = np.clip( illumination[None, :, None] * background_color, 0, 255, ).astype(np.uint8) image = np.repeat(image, height, axis=0) tissue_mask = np.zeros((height, width), dtype=np.uint8) cv2.circle(tissue_mask, (190, 90), 35, 255, cv2.FILLED) tissue_color = np.array([155, 90, 110], dtype=np.uint8) image[tissue_mask > 0] = tissue_color corrected, was_corrected = correct_vignetting(image) self.assertTrue(was_corrected) np.testing.assert_array_equal( corrected[tissue_mask > 0], image[tissue_mask > 0], ) self.assertGreater( corrected[:, -12:][tissue_mask[:, -12:] == 0].mean(), image[:, -12:][tissue_mask[:, -12:] == 0].mean() + 40, ) def test_tissue_mask_fills_and_preserves_pale_spheroid_interior(self): height, width = 240, 320 x = np.linspace(0, 1, width, dtype=np.float32) illumination = 1 - 0.30 * np.clip((x - 0.65) / 0.35, 0, 1) background_color = np.array([245, 240, 235], dtype=np.float32) image = np.clip( illumination[None, :, None] * background_color, 0, 255, ).astype(np.uint8) image = np.repeat(image, height, axis=0) spheroid = np.zeros((height, width), dtype=np.uint8) cv2.circle(spheroid, (250, 120), 55, 255, cv2.FILLED) image[spheroid > 0] = [220, 205, 205] cv2.circle(image, (250, 120), 55, (150, 85, 105), 10) image[110:130, 240:260] = 0 tissue_mask = create_tissue_mask(image) corrected, was_corrected = correct_vignetting(image) cleaned, debris_mask = clean_display_image_and_mask(image, image) self.assertTrue(was_corrected) self.assertEqual(tissue_mask[120, 250], 255) np.testing.assert_array_equal( corrected[spheroid > 0], image[spheroid > 0], ) self.assertTrue(np.any(debris_mask[110:130, 240:260])) self.assertFalse( np.array_equal( cleaned[110:130, 240:260], image[110:130, 240:260], ) ) def test_tiff_output_uses_lossless_lzw_compression(self): image = np.full((128, 128, 3), 127, dtype=np.uint8) with tempfile.TemporaryDirectory() as directory: image_path = Path(directory) / "cleaned.tiff" write_image_rgb(image_path, image) with tifffile.TiffFile(image_path) as tiff: page = tiff.pages[0] self.assertEqual(page.compression.name, "LZW") self.assertEqual(page.tags["Predictor"].value, 2) np.testing.assert_array_equal(page.asarray(), image) self.assertLess(image_path.stat().st_size, image.nbytes) if __name__ == "__main__": unittest.main()