|
| 1 | +"""MinIsWhite photometric handling must be backend-consistent (#1797).""" |
| 2 | +from __future__ import annotations |
| 3 | + |
| 4 | +import http.server |
| 5 | +import importlib.util |
| 6 | +import socketserver |
| 7 | +import threading |
| 8 | + |
| 9 | +import numpy as np |
| 10 | +import pytest |
| 11 | + |
| 12 | +from xrspatial.geotiff import open_geotiff |
| 13 | + |
| 14 | +tifffile = pytest.importorskip("tifffile") |
| 15 | + |
| 16 | + |
| 17 | +def _gpu_available() -> bool: |
| 18 | + """True if cupy is importable and CUDA is initialised.""" |
| 19 | + if importlib.util.find_spec("cupy") is None: |
| 20 | + return False |
| 21 | + try: |
| 22 | + import cupy |
| 23 | + return bool(cupy.cuda.is_available()) |
| 24 | + except Exception: |
| 25 | + return False |
| 26 | + |
| 27 | + |
| 28 | +_HAS_GPU = _gpu_available() |
| 29 | +_gpu_only = pytest.mark.skipif( |
| 30 | + not _HAS_GPU, |
| 31 | + reason="cupy + CUDA required", |
| 32 | +) |
| 33 | + |
| 34 | + |
| 35 | +class _RangeHandler(http.server.BaseHTTPRequestHandler): |
| 36 | + payload: bytes = b'' |
| 37 | + |
| 38 | + def do_GET(self): # noqa: N802 |
| 39 | + rng = self.headers.get('Range') |
| 40 | + if rng and rng.startswith('bytes='): |
| 41 | + spec = rng[len('bytes='):] |
| 42 | + start_s, _, end_s = spec.partition('-') |
| 43 | + start = int(start_s) |
| 44 | + end = int(end_s) if end_s else len(self.payload) - 1 |
| 45 | + chunk = self.payload[start:end + 1] |
| 46 | + self.send_response(206) |
| 47 | + self.send_header('Content-Type', 'application/octet-stream') |
| 48 | + self.send_header( |
| 49 | + 'Content-Range', |
| 50 | + f'bytes {start}-{start + len(chunk) - 1}/{len(self.payload)}', |
| 51 | + ) |
| 52 | + self.send_header('Content-Length', str(len(chunk))) |
| 53 | + self.end_headers() |
| 54 | + self.wfile.write(chunk) |
| 55 | + return |
| 56 | + self.send_response(200) |
| 57 | + self.send_header('Content-Type', 'application/octet-stream') |
| 58 | + self.send_header('Content-Length', str(len(self.payload))) |
| 59 | + self.end_headers() |
| 60 | + self.wfile.write(self.payload) |
| 61 | + |
| 62 | + def log_message(self, *_args, **_kwargs): |
| 63 | + pass |
| 64 | + |
| 65 | + |
| 66 | +def _serve(payload: bytes): |
| 67 | + handler_cls = type( |
| 68 | + 'RangeHandler1797', (_RangeHandler,), {'payload': payload} |
| 69 | + ) |
| 70 | + httpd = socketserver.TCPServer(('127.0.0.1', 0), handler_cls) |
| 71 | + port = httpd.server_address[1] |
| 72 | + thread = threading.Thread(target=httpd.serve_forever, daemon=True) |
| 73 | + thread.start() |
| 74 | + return httpd, port |
| 75 | + |
| 76 | + |
| 77 | +@pytest.fixture |
| 78 | +def miniswhite_http_url(tmp_path, monkeypatch): |
| 79 | + monkeypatch.setenv('XRSPATIAL_GEOTIFF_ALLOW_PRIVATE_HOSTS', '1') |
| 80 | + stored = np.array([[0, 1, 2], [10, 128, 255]], dtype=np.uint8) |
| 81 | + path = tmp_path / "tmp_1797_miniswhite.tif" |
| 82 | + tifffile.imwrite(str(path), stored, photometric='miniswhite') |
| 83 | + httpd, port = _serve(path.read_bytes()) |
| 84 | + try: |
| 85 | + yield f'http://127.0.0.1:{port}/tmp_1797_miniswhite.tif', stored |
| 86 | + finally: |
| 87 | + httpd.shutdown() |
| 88 | + httpd.server_close() |
| 89 | + |
| 90 | + |
| 91 | +def test_http_miniswhite_matches_local_reader(miniswhite_http_url): |
| 92 | + url, stored = miniswhite_http_url |
| 93 | + |
| 94 | + got = open_geotiff(url) |
| 95 | + |
| 96 | + np.testing.assert_array_equal(got.values, np.iinfo(stored.dtype).max - stored) |
| 97 | + |
| 98 | + |
| 99 | +def test_http_dask_miniswhite_matches_local_reader(miniswhite_http_url): |
| 100 | + url, stored = miniswhite_http_url |
| 101 | + |
| 102 | + got = open_geotiff(url, chunks=2).compute() |
| 103 | + |
| 104 | + np.testing.assert_array_equal(got.values, np.iinfo(stored.dtype).max - stored) |
| 105 | + |
| 106 | + |
| 107 | +@_gpu_only |
| 108 | +def test_gpu_miniswhite_matches_cpu_reader(tmp_path): |
| 109 | + from xrspatial.geotiff._writer import write |
| 110 | + |
| 111 | + stored = np.array([[0, 1, 2], [10, 128, 255]], dtype=np.uint8) |
| 112 | + path = str(tmp_path / "tmp_1797_miniswhite_gpu.tif") |
| 113 | + write(stored, path, compression='deflate', tiled=True, tile_size=16, |
| 114 | + photometric='miniswhite') |
| 115 | + |
| 116 | + cpu = open_geotiff(path) |
| 117 | + gpu = open_geotiff(path, gpu=True) |
| 118 | + |
| 119 | + np.testing.assert_array_equal(cpu.values, np.iinfo(stored.dtype).max - stored) |
| 120 | + np.testing.assert_array_equal(gpu.data.get(), cpu.values) |
0 commit comments