initial commit
This commit is contained in:
TheOnlyMace
244
renderer/renderer.py
Normal file
244
renderer/renderer.py
Normal file
@@ -0,0 +1,244 @@
|
||||
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())
|
||||
Reference in New Issue
Block a user