Spaces:
Running
Running
| 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() | |