extract_fldk_rgb.py

#!/usr/bin/env python3
"""
Extract True Color RGB images from Himawari AHI-L1b-FLDK data.

FLDK (Full Disk) data contains all 16 bands of the Advanced Himawari Imager.
This script extracts Band 1 (Blue), Band 2 (Green), and Band 3 (Red) to create
true color RGB composites.

Usage:
    python extract_fldk_rgb.py --input-dir /path/to/fldk/data --output-dir ./fldk_rgb_output
"""

import os
import argparse
import bz2
from pathlib import Path
from typing import Dict, List, Tuple, Optional
import json

import numpy as np
from satpy import Scene
from PIL import Image
from tqdm import tqdm


def decompress_bz2_files(input_dir: str, temp_dir: str) -> List[str]:
    """
    Decompress .DAT.bz2 files to temporary directory.

    Args:
        input_dir: Directory containing .DAT.bz2 files
        temp_dir: Temporary directory for decompressed files

    Returns:
        List of decompressed file paths
    """
    os.makedirs(temp_dir, exist_ok=True)

    bz2_files = list(Path(input_dir).glob("*.bz2"))
    decompressed_files = []

    print(f"Decompressing {len(bz2_files)} .bz2 files...")
    for bz2_file in tqdm(bz2_files):
        output_file = Path(temp_dir) / bz2_file.stem

        if output_file.exists():
            decompressed_files.append(str(output_file))
            continue

        try:
            with bz2.open(bz2_file, 'rb') as f_in:
                with open(output_file, 'wb') as f_out:
                    f_out.write(f_in.read())
            decompressed_files.append(str(output_file))
        except Exception as e:
            print(f"Error decompressing {bz2_file}: {e}")

    return decompressed_files


def group_files_by_timestamp(file_paths: List[str]) -> Dict[str, List[str]]:
    """
    Group files by observation timestamp.

    FLDK filenames format: HS_H08_20241012_0000_B01_FLDK_R20_S0101.DAT
    """
    groups = {}

    for file_path in file_paths:
        filename = Path(file_path).name
        parts = filename.split('_')

        if len(parts) < 4:
            continue

        # Extract date and time
        date_str = parts[2]  # YYYYMMDD
        time_str = parts[3]  # HHMM

        timestamp = f"{date_str}_{time_str}"

        if timestamp not in groups:
            groups[timestamp] = []
        groups[timestamp].append(file_path)

    return groups


def extract_rgb_bands(file_group: List[str]) -> Tuple[Optional[np.ndarray], Optional[Dict]]:
    """
    Extract RGB bands from a group of FLDK files.

    Args:
        file_group: List of file paths for the same timestamp

    Returns:
        RGB array (H, W, 3) and metadata dict, or (None, None) if failed
    """
    try:
        # Load scene with satpy
        scn = Scene(filelist=file_group, reader='ahi_hsd')

        # Load required bands for true color
        # Band 1 (0.47 µm) - Blue
        # Band 2 (0.51 µm) - Green
        # Band 3 (0.64 µm) - Red
        scn.load(['B01', 'B02', 'B03'])

        # Check if all bands loaded
        if 'B01' not in scn or 'B02' not in scn or 'B03' not in scn:
            return None, None

        # Get band arrays
        blue = scn['B01'].values
        green = scn['B02'].values
        red = scn['B03'].values

        # Stack into RGB
        rgb = np.stack([red, green, blue], axis=-1)

        # Normalize to 0-1 range
        # Reflectance bands typically range 0-1 or 0-100%
        rgb_min = np.nanpercentile(rgb, 1)
        rgb_max = np.nanpercentile(rgb, 99)
        rgb = np.clip((rgb - rgb_min) / (rgb_max - rgb_min), 0, 1)

        # Handle NaN values (outside Earth disk)
        rgb = np.nan_to_num(rgb, nan=0.0)

        # Extract metadata
        metadata = {
            'observation_time_utc': scn.start_time.isoformat(),
            'satellite': scn.attrs.get('platform_name', 'Himawari-8'),
            'sensor': 'AHI',
            'area_id': (scn['B01'].attrs.get('area', {}).id
                        if hasattr(scn['B01'].attrs.get('area', {}), 'id')
                        else 'FLDK'),
            'bands': ['B01', 'B02', 'B03'],
            'band_wavelengths': ['0.47µm', '0.51µm', '0.64µm'],
            'image_shape': rgb.shape[:2],
            'resolution': f"{scn['B01'].attrs.get('resolution', 'unknown')} km",
        }

        # Add geographic bounds if available
        try:
            area_def = scn['B01'].attrs.get('area')
            if hasattr(area_def, 'area_extent'):
                lons, lats = area_def.get_lonlats()
                metadata['min_lat'] = float(np.nanmin(lats))
                metadata['max_lat'] = float(np.nanmax(lats))
                metadata['min_lon'] = float(np.nanmin(lons))
                metadata['max_lon'] = float(np.nanmax(lons))
        except Exception:
            pass

        return rgb, metadata

    except Exception as e:
        print(f"Error processing file group: {e}")
        return None, None


def save_rgb_image(rgb: np.ndarray, output_path: str):
    """Save RGB array as PNG image."""
    # Convert to 8-bit
    rgb_8bit = (rgb * 255).astype(np.uint8)

    # Create PIL image
    img = Image.fromarray(rgb_8bit, mode='RGB')

    # Save
    img.save(output_path, 'PNG')


def save_metadata(metadata: Dict, output_path: str):
    """Save metadata as JSON."""
    with open(output_path, 'w') as f:
        json.dump(metadata, f, indent=2)


def process_fldk_directory(
    input_dir: str,
    output_dir: str,
    temp_dir: str = "./temp_fldk",
    cleanup_temp: bool = True
):
    """
    Process all FLDK files in directory and extract RGB images.

    Args:
        input_dir: Directory containing .DAT.bz2 files
        output_dir: Output directory for RGB images and metadata
        temp_dir: Temporary directory for decompressed files
        cleanup_temp: Whether to delete temporary files after processing
    """
    # Create output directories
    image_dir = Path(output_dir) / "images"
    metadata_dir = Path(output_dir) / "metadata"
    image_dir.mkdir(parents=True, exist_ok=True)
    metadata_dir.mkdir(parents=True, exist_ok=True)

    # Step 1: Decompress files
    decompressed_files = decompress_bz2_files(input_dir, temp_dir)

    if not decompressed_files:
        print("No files to process!")
        return

    # Step 2: Group by timestamp
    print("Grouping files by timestamp...")
    file_groups = group_files_by_timestamp(decompressed_files)
    print(f"Found {len(file_groups)} timestamps")

    # Step 3: Process each timestamp
    print("Extracting RGB images...")
    success_count = 0

    for timestamp, files in tqdm(file_groups.items()):
        # Extract RGB
        rgb, metadata = extract_rgb_bands(files)

        if rgb is None:
            continue

        # Save image
        image_filename = f"fldk_rgb_{timestamp}.png"
        image_path = image_dir / image_filename
        save_rgb_image(rgb, str(image_path))

        # Save metadata
        metadata['filename'] = image_filename
        metadata_filename = f"fldk_rgb_{timestamp}.json"
        metadata_path = metadata_dir / metadata_filename
        save_metadata(metadata, str(metadata_path))

        success_count += 1

    print(f"\nSuccessfully processed {success_count}/{len(file_groups)} timestamps")
    print(f"Images saved to: {image_dir}")
    print(f"Metadata saved to: {metadata_dir}")

    # Step 4: Cleanup
    if cleanup_temp:
        print(f"Cleaning up temporary files in {temp_dir}...")
        import shutil
        shutil.rmtree(temp_dir, ignore_errors=True)


def main():
    parser = argparse.ArgumentParser(
        description="Extract True Color RGB images from Himawari AHI-L1b-FLDK data"
    )
    parser.add_argument(
        "--input-dir",
        required=True,
        help="Directory containing .DAT.bz2 files"
    )
    parser.add_argument(
        "--output-dir",
        default="./fldk_rgb_output",
        help="Output directory for RGB images and metadata"
    )
    parser.add_argument(
        "--temp-dir",
        default="./temp_fldk",
        help="Temporary directory for decompressed files"
    )
    parser.add_argument(
        "--keep-temp",
        action="store_true",
        help="Keep temporary decompressed files (don't delete)"
    )

    args = parser.parse_args()

    # Verify input directory exists
    if not os.path.isdir(args.input_dir):
        print(f"Error: Input directory does not exist: {args.input_dir}")
        return

    # Check for .bz2 files
    bz2_count = len(list(Path(args.input_dir).glob("*.bz2")))
    if bz2_count == 0:
        print(f"Error: No .bz2 files found in {args.input_dir}")
        return

    print(f"Found {bz2_count} .bz2 files in {args.input_dir}")

    # Process directory
    process_fldk_directory(
        args.input_dir,
        args.output_dir,
        args.temp_dir,
        cleanup_temp=not args.keep_temp
    )

    print("\nDone! You can now use this data for training:")
    print(f"  python train_diffusion_truecolor.py --data-dir {args.output_dir}")


if __name__ == "__main__":
    main()