245 lines
8.1 KiB
Python
245 lines
8.1 KiB
Python
import asyncio
|
|
import json
|
|
import logging
|
|
import os
|
|
import re
|
|
import tarfile
|
|
import tempfile
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
|
|
import asyncpg
|
|
import h5py
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
|
logger = logging.getLogger("radar-renderer")
|
|
|
|
DATA_DIR = Path(os.getenv("DATA_DIR", "/data"))
|
|
RAW_RADAR_DIR = DATA_DIR / "raw" / "radar"
|
|
RENDERED_RADAR_DIR = DATA_DIR / "processed" / "radar"
|
|
RADAR_PRODUCT = os.getenv("RADAR_PRODUCT", "rv")
|
|
RENDER_INTERVAL_SECONDS = int(os.getenv("RADAR_RENDER_INTERVAL_SECONDS", "300"))
|
|
DATABASE_URL = os.getenv("DATABASE_URL", "postgresql://hexawetter:hexawetter@postgres:5432/hexawetter")
|
|
RADAR_FILE_RE = re.compile(r"^[A-Za-z0-9_.-]+\.(?:tar|tar\.bz2|bz2|gz)$", re.IGNORECASE)
|
|
|
|
# DWD RV composite extent (EPSG:4326 approximate)
|
|
DEFAULT_BOUNDS = {
|
|
"south": 47.0,
|
|
"west": 3.0,
|
|
"north": 56.0,
|
|
"east": 16.0,
|
|
}
|
|
|
|
# Precipitation colormap (mm/h)
|
|
COLOR_STOPS = [
|
|
(0.0, (0, 0, 0, 0)),
|
|
(0.1, (100, 180, 255, 80)),
|
|
(0.5, (0, 120, 255, 140)),
|
|
(1.0, (0, 200, 80, 170)),
|
|
(2.0, (255, 255, 0, 200)),
|
|
(5.0, (255, 140, 0, 220)),
|
|
(10.0, (255, 0, 0, 230)),
|
|
(20.0, (180, 0, 120, 240)),
|
|
(50.0, (120, 0, 80, 250)),
|
|
]
|
|
|
|
|
|
def apply_colormap(values: np.ndarray) -> np.ndarray:
|
|
rgba = np.zeros((*values.shape, 4), dtype=np.uint8)
|
|
mask = np.isfinite(values) & (values > 0)
|
|
if not np.any(mask):
|
|
return rgba
|
|
vmax = max(float(np.nanpercentile(values[mask], 99)), 1.0)
|
|
normalized = np.clip(values / vmax, 0, 1)
|
|
stops = COLOR_STOPS
|
|
for idx in range(len(stops) - 1):
|
|
low_val, low_color = stops[idx]
|
|
high_val, high_color = stops[idx + 1]
|
|
if high_val <= low_val:
|
|
continue
|
|
band = mask & (values >= low_val) & (values < high_val)
|
|
if not np.any(band):
|
|
continue
|
|
ratio = (values[band] - low_val) / (high_val - low_val)
|
|
for channel in range(4):
|
|
rgba[..., channel][band] = (
|
|
low_color[channel] + ratio * (high_color[channel] - low_color[channel])
|
|
).astype(np.uint8)
|
|
top_band = mask & (values >= stops[-1][0])
|
|
if np.any(top_band):
|
|
for channel in range(4):
|
|
rgba[..., channel][top_band] = stops[-1][1][channel]
|
|
return rgba
|
|
|
|
|
|
def read_hdf5_dataset(path: Path) -> np.ndarray | None:
|
|
try:
|
|
with h5py.File(path, "r") as handle:
|
|
if "dataset1/data1/data" in handle:
|
|
data = handle["dataset1/data1/data"][:]
|
|
what = handle["dataset1/what"].attrs
|
|
nodata = what.get("nodata", -9999)
|
|
data = np.where(data == nodata, np.nan, data)
|
|
return data.astype(np.float32)
|
|
for key in handle.keys():
|
|
if isinstance(handle[key], h5py.Group) and "data" in handle[key]:
|
|
subgroup = handle[key]
|
|
for subkey in subgroup.keys():
|
|
if "data" in subgroup[subkey]:
|
|
return np.array(subgroup[subkey]["data"], dtype=np.float32)
|
|
except Exception as exc:
|
|
logger.warning("HDF5 read failed for %s: %s", path.name, exc)
|
|
return None
|
|
|
|
|
|
def extract_archive(archive_path: Path, temp_dir: Path) -> list[Path]:
|
|
extracted: list[Path] = []
|
|
if archive_path.suffix == ".bz2" and archive_path.name.endswith(".tar.bz2"):
|
|
mode = "r:bz2"
|
|
elif archive_path.suffix == ".gz":
|
|
mode = "r:gz"
|
|
else:
|
|
mode = "r"
|
|
with tarfile.open(archive_path, mode) as tar:
|
|
tar.extractall(temp_dir, filter="data")
|
|
for path in temp_dir.rglob("*"):
|
|
if path.is_file():
|
|
extracted.append(path)
|
|
return extracted
|
|
|
|
|
|
def render_png(data: np.ndarray, output_path: Path) -> dict:
|
|
rgba = apply_colormap(data)
|
|
image = Image.fromarray(rgba, mode="RGBA")
|
|
image = image.transpose(Image.Transpose.FLIP_TOP_BOTTOM)
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
image.save(output_path, format="PNG", optimize=True)
|
|
return DEFAULT_BOUNDS
|
|
|
|
|
|
async def upsert_meta(
|
|
pool: asyncpg.Pool,
|
|
filename: str,
|
|
file_size: int,
|
|
modified_at: datetime,
|
|
render_status: str,
|
|
render_path: str | None = None,
|
|
render_bounds: dict | None = None,
|
|
render_error: str | None = None,
|
|
) -> None:
|
|
await pool.execute(
|
|
"""
|
|
INSERT INTO radar_file_meta (
|
|
product, filename, file_size, modified_at, render_status,
|
|
render_path, render_bounds, render_error, updated_at
|
|
) VALUES ($1, $2, $3, $4, $5, $6, $7::jsonb, $8, NOW())
|
|
ON CONFLICT (filename) DO UPDATE SET
|
|
file_size = EXCLUDED.file_size,
|
|
modified_at = EXCLUDED.modified_at,
|
|
render_status = EXCLUDED.render_status,
|
|
render_path = EXCLUDED.render_path,
|
|
render_bounds = EXCLUDED.render_bounds,
|
|
render_error = EXCLUDED.render_error,
|
|
updated_at = NOW()
|
|
""",
|
|
RADAR_PRODUCT,
|
|
filename,
|
|
file_size,
|
|
modified_at,
|
|
render_status,
|
|
render_path,
|
|
json.dumps(render_bounds) if render_bounds else None,
|
|
render_error,
|
|
)
|
|
|
|
|
|
async def process_file(pool: asyncpg.Pool, archive_path: Path) -> None:
|
|
filename = archive_path.name
|
|
modified_at = datetime.fromtimestamp(archive_path.stat().st_mtime, timezone.utc)
|
|
output_path = RENDERED_RADAR_DIR / RADAR_PRODUCT / f"{filename}.png"
|
|
if output_path.exists() and output_path.stat().st_size > 0:
|
|
await upsert_meta(
|
|
pool,
|
|
filename,
|
|
archive_path.stat().st_size,
|
|
modified_at,
|
|
"done",
|
|
str(output_path),
|
|
DEFAULT_BOUNDS,
|
|
)
|
|
return
|
|
|
|
await upsert_meta(pool, filename, archive_path.stat().st_size, modified_at, "processing")
|
|
try:
|
|
with tempfile.TemporaryDirectory() as temp_name:
|
|
temp_dir = Path(temp_name)
|
|
extracted = extract_archive(archive_path, temp_dir)
|
|
data = None
|
|
for candidate in extracted:
|
|
if candidate.suffix.lower() in {".h5", ".hdf5"} or "RV" in candidate.name.upper():
|
|
data = read_hdf5_dataset(candidate)
|
|
if data is not None:
|
|
break
|
|
if data is None:
|
|
for candidate in extracted:
|
|
data = read_hdf5_dataset(candidate)
|
|
if data is not None:
|
|
break
|
|
if data is None:
|
|
raise RuntimeError("No renderable HDF5 dataset found in archive")
|
|
|
|
bounds = render_png(data, output_path)
|
|
await upsert_meta(
|
|
pool,
|
|
filename,
|
|
archive_path.stat().st_size,
|
|
modified_at,
|
|
"done",
|
|
str(output_path),
|
|
bounds,
|
|
)
|
|
logger.info("rendered %s -> %s", filename, output_path)
|
|
except Exception as exc:
|
|
await upsert_meta(
|
|
pool,
|
|
filename,
|
|
archive_path.stat().st_size,
|
|
modified_at,
|
|
"error",
|
|
render_error=str(exc),
|
|
)
|
|
logger.error("render failed for %s: %s", filename, exc)
|
|
|
|
|
|
async def render_pending() -> None:
|
|
target_dir = RAW_RADAR_DIR / "composite" / RADAR_PRODUCT
|
|
target_dir.mkdir(parents=True, exist_ok=True)
|
|
RENDERED_RADAR_DIR.mkdir(parents=True, exist_ok=True)
|
|
|
|
pool = await asyncpg.create_pool(DATABASE_URL, min_size=1, max_size=3)
|
|
try:
|
|
archives = sorted(
|
|
[p for p in target_dir.iterdir() if p.is_file() and RADAR_FILE_RE.fullmatch(p.name)],
|
|
key=lambda p: p.stat().st_mtime,
|
|
)[-12:]
|
|
for archive_path in archives:
|
|
await process_file(pool, archive_path)
|
|
finally:
|
|
await pool.close()
|
|
|
|
|
|
async def main() -> None:
|
|
logger.info("radar renderer started (product=%s)", RADAR_PRODUCT)
|
|
while True:
|
|
try:
|
|
await render_pending()
|
|
except Exception as exc:
|
|
logger.exception("render cycle failed: %s", exc)
|
|
await asyncio.sleep(RENDER_INTERVAL_SECONDS)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|