extract_himawari_dataset.py

#!/usr/bin/env python3
"""
Extract IR Band (B13) + Cloud Mask + Cloud Phase from Himawari data.

Creates a 3-channel training image for diffusion models:
channel 0 = IR B13 brightness
channel 1 = Cloud mask
channel 2 = Cloud phase

Outputs:
  ./dataset/images/*.png
  ./dataset/metadata/*.json
"""

import bz2
import json
import tempfile
import numpy as np
from pathlib import Path
from tqdm import tqdm
from PIL import Image
import xarray as xr

# ---------------------------------------------
# Paths (no CLI args needed)
# ---------------------------------------------
CONFIG = {
    "L1_DIR": "./L1",
    "L2_MASK_DIR": "./L2_mask",
    "L2_PHASE_DIR": "./L2_phase",
    "OUT_DIR": "./dataset",
    "IMAGE_SIZE": 256
}


def decompress_bz2(path):
    """Decompress .DAT.bz2 and return temp file path."""
    if not path.endswith(".bz2"):
        return path
    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".DAT").name
    with bz2.open(path, "rb") as f_in, open(tmp, "wb") as f_out:
        f_out.write(f_in.read())
    return tmp


def extract_b13_array(dat_path):
    """Extract B13 brightness temperature from raw DAT (HDF5 disguised)."""
    try:
        data = xr.open_dataset(dat_path, engine="h5netcdf")
        if "tbb" in data:
            arr = data["tbb"].values
        else:
            arr = list(data.data_vars.values())[0].values  # fallback
        arr = np.nan_to_num(arr)
        return arr
    except Exception:
        print("nothing to print")
        return None


def extract_nc_array(path):
    ds = xr.open_dataset(path)
    arr = list(ds.data_vars.values())[0].values
    arr = np.nan_to_num(arr)
    return arr


def resize(arr, size):
    img = Image.fromarray(arr.astype(np.float32))
    img = img.resize((size, size), Image.BILINEAR)
    return np.array(img)


def main():
    L1 = Path(CONFIG["L1_DIR"])
    MASK = Path(CONFIG["L2_MASK_DIR"])
    PHASE = Path(CONFIG["L2_PHASE_DIR"])
    OUT = Path(CONFIG["OUT_DIR"])

    (OUT / "images").mkdir(parents=True, exist_ok=True)
    (OUT / "metadata").mkdir(parents=True, exist_ok=True)

    # Collect timestamps
    timestamps = sorted([p.name for p in L1.iterdir() if p.is_dir()])

    print(f"Found {len(timestamps)} timestamps")

    for ts in tqdm(timestamps):
        # -----------------------------------
        # 1. Find matching files
        # -----------------------------------
        ts_dir = L1 / ts

        dat_files = sorted(ts_dir.glob("*B13*.DAT.bz2"))
        if len(dat_files) == 0:
            continue

        mask_file = MASK / ts / next((p.name for p in (MASK / ts).glob("*.nc")), None)
        phase_file = PHASE / ts / next((p.name for p in (PHASE / ts).glob("*.nc")), None)

        if not mask_file or not phase_file:
            continue

        # -----------------------------------
        # 2. Merge DAT segments into single array
        # -----------------------------------
        merged = None
        for part in dat_files:
            tmp = decompress_bz2(str(part))
            arr = extract_b13_array(tmp)
            if arr is None:
                continue
            if merged is None:
                merged = arr
            else:
                merged = np.maximum(merged, arr)

        if merged is None:
            print(f"Failed to read B13 for {ts}")
            continue

        # -----------------------------------
        # 3. Read cloud mask & cloud phase
        # -----------------------------------
        mask_arr = extract_nc_array(mask_file)
        phase_arr = extract_nc_array(phase_file)

        # -----------------------------------
        # 4. Normalize/resize
        # -----------------------------------
        merged = resize(merged, CONFIG["IMAGE_SIZE"])
        mask_arr = resize(mask_arr, CONFIG["IMAGE_SIZE"])
        phase_arr = resize(phase_arr, CONFIG["IMAGE_SIZE"])

        # Normalize IR
        merged = merged.astype(np.float32)
        merged = (merged - merged.min()) / (merged.max() - merged.min() + 1e-6)

        # Clip mask + phase
        mask_arr = np.clip(mask_arr, 0, 3)
        phase_arr = np.clip(phase_arr, 0, 3)

        # -----------------------------------
        # 5. Stack into 3-channel training tensor
        # -----------------------------------
        stacked = np.stack([merged, mask_arr / 3.0, phase_arr / 3.0], axis=-1)

        # Convert to uint8 image for storage
        img_uint8 = (stacked * 255).astype(np.uint8)

        img = Image.fromarray(img_uint8, mode="RGB")

        out_img = OUT / "images" / f"{ts}.png"
        out_meta = OUT / "metadata" / f"{ts}.json"

        img.save(out_img)

        # -----------------------------------
        # 6. Metadata
        # -----------------------------------
        metadata = {
            "timestamp": ts,
            "image": str(out_img),
            "mask": str(mask_file),
            "phase": str(phase_file),
            "channels": ["IR_B13", "CloudMask", "CloudPhase"]
        }

        with open(out_meta, "w") as f:
            json.dump(metadata, f, indent=2)

    print("Extraction complete.")


if __name__ == "__main__":
    main()