Spaces:
Sleeping
Sleeping
| # tests/test_fsrcnn.py | |
| import os | |
| import importlib | |
| import numpy as np | |
| import torch | |
| import pytest | |
| # Adjust this import if your file isn't named fsrcnn_app.py | |
| import app as app | |
| def _reset_cache_between_tests(): | |
| # Ensure cache isolation between tests | |
| app.MODEL_CACHE.clear() | |
| yield | |
| app.MODEL_CACHE.clear() | |
| def test_fsrcnn_forward_output_shape_cpu_only(): | |
| """FSRCNN forward should upscale 1-channel input by its scale factor.""" | |
| model = app.FSRCNN(scale_factor=3).eval() | |
| x = torch.randn(1, 1, 10, 12) # (N, C, H, W) | |
| with torch.inference_mode(): | |
| y = model(x) | |
| assert y.shape == (1, 1, 30, 36), "Output shape must be (H*scale, W*scale)" | |
| def test_run_fsrcnn_on_y_shape_and_dtype(): | |
| """run_fsrcnn_on_y should return uint8 image with upscaled spatial dims.""" | |
| y = np.random.randint(0, 256, (9, 7), dtype=np.uint8) | |
| model = app.FSRCNN(scale_factor=2).eval() | |
| out = app.run_fsrcnn_on_y(y, model) | |
| assert out.dtype == np.uint8 | |
| assert out.shape == (9 * 2, 7 * 2) | |
| def test_bicubic_upscale_rgb_shape_and_dtype(): | |
| rgb = np.random.randint(0, 256, (16, 24, 3), dtype=np.uint8) | |
| out = app.bicubic_upscale_rgb(rgb, scale=4) | |
| assert out.dtype == np.uint8 | |
| assert out.shape == (16 * 4, 24 * 4, 3) | |
| def test_rgb_ycbcr_roundtrip_close(): | |
| """RGB -> YCrCb -> RGB roundtrip should be close (small max diff).""" | |
| rgb = np.random.randint(0, 256, (32, 32, 3), dtype=np.uint8) | |
| ycrcb = app.rgb_to_ycbcr(rgb) | |
| back = app.ycbcr_to_rgb(ycrcb) | |
| # Allow small numerical differences from color conversion | |
| assert np.max(np.abs(back.astype(int) - rgb.astype(int))) <= 2 | |
| def test_fsrcnn_upscale_falls_back_to_bicubic_when_no_weights(tmp_path): | |
| """When no valid weights are provided, FSRCNN code must return bicubic result.""" | |
| rgb = np.random.randint(0, 256, (12, 10, 3), dtype=np.uint8) | |
| scale = 3 | |
| # Ensure a fresh cache so the "no-weights" path is exercised | |
| app.MODEL_CACHE.clear() | |
| out_fallback = app.fsrcnn_upscale_rgb(rgb, scale=scale, weights=None) | |
| out_bicubic = app.bicubic_upscale_rgb(rgb, scale=scale) | |
| assert out_fallback.shape == out_bicubic.shape | |
| # Code path returns bicubic directly; should be byte-identical | |
| assert np.array_equal(out_fallback, out_bicubic) | |
| def test_ui_accepts_grayscale_and_rgba_and_clips(): | |
| """The UI helper should handle grayscale, RGBA, and non-uint8 inputs.""" | |
| # Grayscale -> stacked to RGB | |
| gray = np.random.randint(0, 256, (8, 8), dtype=np.uint8) | |
| out_gray = app.upscale_ui(gray, 2, "Bicubic", "", "", "") | |
| assert out_gray.shape == (16, 16, 3) | |
| assert out_gray.dtype == np.uint8 | |
| # RGBA -> drop alpha | |
| rgba = np.random.randint(0, 256, (8, 8, 4), dtype=np.uint8) | |
| out_rgba = app.upscale_ui(rgba, 2, "Bicubic", "", "", "") | |
| assert out_rgba.shape == (16, 16, 3) | |
| assert out_rgba.dtype == np.uint8 | |
| # Float input -> should clip/convert to uint8 internally | |
| f_rgb = np.random.randn(8, 8, 3).astype(np.float32) * 1000.0 # intentionally wild | |
| out_float = app.upscale_ui(f_rgb, 2, "Bicubic", "", "", "") | |
| assert out_float.dtype == np.uint8 | |
| assert out_float.shape == (16, 16, 3) | |
| def test_maybe_downscale_for_memory_respects_limit(): | |
| big = np.random.randint(0, 256, (4000, 4000, 3), dtype=np.uint8) # 16M px | |
| capped = app.maybe_downscale_for_memory(big, max_pixels=1_000_000) | |
| assert capped.shape[0] * capped.shape[1] <= 1_000_000 | |
| def test_get_model_cache_per_scale(): | |
| m2, w2 = app.get_model(2, weights_path=None) | |
| m3, w3 = app.get_model(3, weights_path=None) | |
| # Cache populated for both scales | |
| assert 2 in app.MODEL_CACHE and 3 in app.MODEL_CACHE | |
| assert isinstance(m2, app.FSRCNN) and isinstance(m3, app.FSRCNN) | |
| assert m2 is not m3, "Different scales should use different model instances" | |