| """MRMS data loading utilities matching the official NowcastNet contract.""" |
|
|
| from __future__ import annotations |
|
|
| import struct |
| import zlib |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| from torch.utils.data import Dataset |
|
|
|
|
| def _read_png_gray16(path: Path) -> np.ndarray: |
| """Read the 16-bit grayscale PNG subset used by MRMS.""" |
| raw = path.read_bytes() |
| if raw[:8] != b"\x89PNG\r\n\x1a\n": |
| raise ValueError(f"Not a PNG file: {path}") |
| pos = 8 |
| idat: list[bytes] = [] |
| width = height = bit_depth = color_type = None |
| while pos < len(raw): |
| length = struct.unpack(">I", raw[pos:pos + 4])[0] |
| kind = raw[pos + 4:pos + 8] |
| payload = raw[pos + 8:pos + 8 + length] |
| pos += length + 12 |
| if kind == b"IHDR": |
| width, height, bit_depth, color_type = struct.unpack(">IIBB", payload[:10]) |
| elif kind == b"IDAT": |
| idat.append(payload) |
| elif kind == b"IEND": |
| break |
| if width is None or height is None or bit_depth != 16 or color_type != 0: |
| raise ValueError(f"Expected 16-bit grayscale PNG: {path}") |
|
|
| decoded = zlib.decompress(b"".join(idat)) |
| row_bytes = width * 2 |
| previous = np.zeros(row_bytes, dtype=np.uint8) |
| rows: list[np.ndarray] = [] |
| cursor = 0 |
| for _ in range(height): |
| filter_type = decoded[cursor] |
| row = np.frombuffer(decoded[cursor + 1:cursor + 1 + row_bytes], dtype=np.uint8).copy() |
| cursor += row_bytes + 1 |
| for index in range(row_bytes): |
| left = row[index - 2] if index >= 2 else 0 |
| up = previous[index] |
| upper_left = previous[index - 2] if index >= 2 else 0 |
| if filter_type == 1: |
| value = left |
| elif filter_type == 2: |
| value = up |
| elif filter_type == 3: |
| value = (int(left) + int(up)) // 2 |
| elif filter_type == 4: |
| predictor = int(left) + int(up) - int(upper_left) |
| distances = (abs(predictor - int(left)), abs(predictor - int(up)), abs(predictor - int(upper_left))) |
| value = (left, up, upper_left)[int(np.argmin(distances))] |
| elif filter_type == 0: |
| value = 0 |
| else: |
| raise ValueError(f"Unsupported PNG filter {filter_type}") |
| row[index] = (int(row[index]) + int(value)) & 255 |
| rows.append(row) |
| previous = row |
| return np.frombuffer(b"".join(row.tobytes() for row in rows), dtype=">u2").astype(np.uint16).reshape(height, width) |
|
|
|
|
| def decode_mrms_event(event_dir: str | Path, image_height: int = 512, image_width: int = 512, total_length: int = 29) -> np.ndarray: |
| event_dir = Path(event_dir) |
| event_name = event_dir.name |
| frames = [_read_png_gray16(event_dir / f"{event_name}-{index:02d}.png") for index in range(total_length)] |
| data = np.stack(frames).astype(np.float32) / 10.0 - 3.0 |
| if data.shape[1:] != (image_height, image_width): |
| raise ValueError(f"MRMS frame shape {data.shape[1:]} != {(image_height, image_width)}") |
| mask = np.ones_like(data, dtype=np.float32) |
| mask[data < 0] = 0 |
| data[data < 0] = 0 |
| data = np.clip(data, 0, 128) |
| return np.stack([data, mask], axis=-1) |
|
|
|
|
| class MRMSDataset(Dataset): |
| def __init__(self, data_dir: str | Path, image_height: int = 512, image_width: int = 512, total_length: int = 29, split: str = "train"): |
| self.data_dir = Path(data_dir) |
| self.events = sorted(path for path in self.data_dir.iterdir() if path.is_dir()) |
| if split == "train": |
| self.events = [path for path in self.events if len(list(path.glob("*.png"))) == total_length] |
| elif split != "test": |
| raise ValueError("split must be train or test") |
| self.image_height = image_height |
| self.image_width = image_width |
| self.total_length = total_length |
|
|
| def __len__(self) -> int: |
| return len(self.events) |
|
|
| def __getitem__(self, index: int) -> dict[str, torch.Tensor | str]: |
| event = self.events[index] |
| frames = decode_mrms_event(event, self.image_height, self.image_width, self.total_length) |
| return {"radar_frames": torch.from_numpy(frames), "event": event.name} |
|
|
