File size: 4,246 Bytes
439c523 | 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 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | """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}
|