File size: 2,017 Bytes
574aa8f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 | from __future__ import annotations
import sys
import tempfile
import unittest
from pathlib import Path
import numpy as np
import torch
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from denoise import denoise_array, rgb_to_ycbcr, ycbcr_to_rgb # noqa: E402
from image_io import load_image, save_image # noqa: E402
from model import CosmicClarityDenoiseCNN, model_statistics # noqa: E402
class CoreTests(unittest.TestCase):
def test_model_shape_and_parameter_count(self) -> None:
model = CosmicClarityDenoiseCNN()
output = model(torch.rand(1, 3, 16, 16))
self.assertEqual(tuple(output.shape), (1, 3, 16, 16))
self.assertEqual(model_statistics(model)["parameters"], 2_945_283)
def test_ycbcr_roundtrip(self) -> None:
image = np.random.default_rng(3).random((20, 21, 3), dtype=np.float32)
reconstructed = ycbcr_to_rgb(*rgb_to_ycbcr(image))
np.testing.assert_allclose(image, reconstructed, atol=3e-5)
def test_zero_strength_is_identity(self) -> None:
image = np.random.default_rng(4).random((9, 10, 3), dtype=np.float32)
result = denoise_array(image, torch.nn.Identity(), torch.device("cpu"), strength=0.0)
np.testing.assert_array_equal(image, result)
def test_fits_roundtrip_preserves_channel_layout(self) -> None:
from astropy.io import fits
original = np.arange(3 * 8 * 9, dtype=np.uint16).reshape(3, 8, 9)
with tempfile.TemporaryDirectory() as directory:
source = Path(directory) / "source.fits"
output = Path(directory) / "output.fits"
fits.PrimaryHDU(original).writeto(source)
record = load_image(source)
save_image(output, record, record.pixels)
with fits.open(output) as hdul:
restored = hdul[0].data
self.assertEqual(restored.shape, original.shape)
np.testing.assert_array_equal(restored, original)
if __name__ == "__main__":
unittest.main()
|