LordZeee/ndvi-convlstm

1# prepare_training_data_with_custom_lcu.py2import os3import tensorflow as tf4import numpy as np5import yaml6import argparse7from datetime import datetime, timedelta8import glob9import re10import rasterio  # For reading .tif files11 12CONFIG = None13 14 15def get_base_prefix_for_mixer(tfrecord_filename):16    name_part = os.path.basename(tfrecord_filename)17    if name_part.endswith(".tfrecord"):18        name_part = name_part[: -len(".tfrecord")]19    name_part = re.sub(r"-\d{5}-of-\d{5}$", "", name_part)20    name_part = re.sub(r"-\d{5}$", "", name_part)21    return name_part22 23 24def read_tfrecord_image_data_flexible(25    tfrecord_path, patch_height, patch_width, primary_band_names, record_type_for_log=""26):27    if not primary_band_names:28        return None, None29 30    def _try_read_with_names(current_tfrecord_path, names_to_try_list):31        if not names_to_try_list:32            return None33        valid_names_to_try = [34            name for name in names_to_try_list if name and isinstance(name, str)35        ]36        if not valid_names_to_try:37            return None38        try:39            raw_dataset = tf.data.TFRecordDataset([current_tfrecord_path])40            feature_description = {41                b: tf.io.FixedLenFeature([patch_height * patch_width], tf.float32)42                for b in valid_names_to_try43            }44 45            def _parse_function(example_proto):46                parsed = tf.io.parse_single_example(example_proto, feature_description)47                bands_list = [48                    tf.reshape(parsed[b], [patch_height, patch_width])49                    for b in valid_names_to_try50                ]51                return tf.stack(bands_list, axis=-1)52 53            parsed_dataset = raw_dataset.map(_parse_function)54            for image_tensor in parsed_dataset.take(1):55                return image_tensor.numpy()56            return None57        except tf.errors.InvalidArgumentError as e_tf_invalid_arg:58            err_msg_lower = str(e_tf_invalid_arg).lower()59            if (60                "is required but could not be found" in err_msg_lower61                or "feature is not in schema" in err_msg_lower62                or "could not find feature" in err_msg_lower63            ):64                return "KEY_MISMATCH_RETRY"65            return None66        except Exception:67            return None68 69    used_band_names_list, image_data = [], None70    if primary_band_names and len(primary_band_names) > 0:71        image_data = _try_read_with_names(tfrecord_path, primary_band_names)72        if image_data is not None and not isinstance(image_data, str):73            used_band_names_list = primary_band_names74            return image_data, used_band_names_list75    print("Error: tfrecord is None!")76    return None, None77 78 79def read_custom_lcu_tif(lcu_tif_path, expected_height, expected_width):80    """Reads the custom LCU .tif file and returns it as a (H, W, 1) numpy array."""81    if not lcu_tif_path or not os.path.exists(lcu_tif_path):82        # print(f"      Warn: Custom LCU TIF not found: {lcu_tif_path}")83        return None84    try:85        with rasterio.open(lcu_tif_path) as src:86            if src.count != 1:87                print(88                    f"      Warn: Custom LCU TIF {lcu_tif_path} has {src.count} bands, expected 1. Using first band."89                )90 91            data = src.read(1).astype(np.float32)  # Read first band, ensure float3292 93            if data.shape[0] != expected_height or data.shape[1] != expected_width:94                print(95                    f"      Warn: Custom LCU TIF {lcu_tif_path} has shape {data.shape}, expected ({expected_height},{expected_width}). Resizing/Cropping might be needed if critical."96                )97                # Simple center crop or resize if necessary - for now, just warning98                # For robust solution, implement proper resize/crop if shapes can vary99                return None  # Or attempt resize: cv2.resize(data, (expected_width, expected_height), interpolation=cv2.INTER_NEAREST)100 101            return data[:, :, np.newaxis]  # Add channel dimension -> (H, W, 1)102    except Exception as e:103        print(f"      Error reading custom LCU TIF {lcu_tif_path}: {e}")104        return None105 106 107def get_doy_features(date_obj):108    doy = date_obj.timetuple().tm_yday109    sin_doy = np.sin(2 * np.pi * doy / 365.25)110    cos_doy = np.cos(2 * np.pi * doy / 365.25)111    return sin_doy, cos_doy112 113 114def find_additional_layers_tfrecord(115    base_add_dir, mc_id, week_date_obj, template, suffix116):117    # ... (This function remains the same as the last version) ...118    week_ymd = week_date_obj.strftime("%Y%m%d")119    try:120        base_part = template.format(minicube_id=mc_id, week_date_nodash=week_ymd)121    except KeyError as e:122        print(f"    KeyError in template: {e}")123        return None124    if suffix.lower() != ".tfrecord":125        fp = os.path.join(base_add_dir, base_part + "-00000" + suffix)126        return fp if os.path.exists(fp) else None127    glob_p = os.path.join(base_add_dir, f"{base_part}*.tfrecord")128    files = [129        f for f in glob.glob(glob_p) if not os.path.basename(f).endswith("-mixer.json")130    ]131    if not files:132        return None133    exact_match = os.path.join(base_add_dir, base_part + ".tfrecord")134    return exact_match if exact_match in files else sorted(files)[0]135 136 137def process_minicubes_with_additional_layers(config_data):138    global CONFIG139    CONFIG = config_data140 141    # --- Directory Setup ---142    original_tfrecord_dir = CONFIG["directories"].get("local_tfrecord_download_dir")143    additional_tfrecord_dir = CONFIG["directories"].get(144        "additional_layers_tfrecord_download_dir"145    )146    custom_lcu_dir = CONFIG["directories"].get("custom_land_cover_dir")  # New directory147    processed_output_base_dir = CONFIG["directories"].get("processed_training_data_dir")148 149    req_dirs_to_check = [150        (original_tfrecord_dir, "local_tfrecord_download_dir"),151        (processed_output_base_dir, "processed_training_data_dir"),152    ]153    # Additional layers and custom LCU are optional sources, check if configured later154    if CONFIG.get("training_options", {}).get("additional_bands_names"):155        req_dirs_to_check.append(156            (additional_tfrecord_dir, "additional_layers_tfrecord_download_dir")157        )158 159    # Check if custom_LCU is in final_npz_order to see if custom_lcu_dir is needed160    final_npz_order_check = CONFIG.get("training_options", {}).get(161        "final_input_band_names_for_npz", []162    )163    if any(164        "custom_lcu" in b.lower() for b in final_npz_order_check165    ):  # Check if custom LCU is expected166        if not custom_lcu_dir:167            print(168                "Error: 'custom_LCU' in final bands, but 'custom_land_cover_dir' not in config."169            )170            return171        req_dirs_to_check.append((custom_lcu_dir, "custom_land_cover_dir"))172 173    for dir_path, dir_name in req_dirs_to_check:174        if not dir_path:175            print(f"Error: '{dir_name}' not specified in config.")176            return177        if dir_name != "processed_training_data_dir" and not os.path.exists(178            dir_path179        ):  # Source dirs must exist180            print(f"Error: Source directory '{dir_path}' ({dir_name}) not found.")181            return182    os.makedirs(processed_output_base_dir, exist_ok=True)183    print(f"Outputting combined .npz to: {processed_output_base_dir}")184 185    # --- Config Parameters ---186    mcp = CONFIG.get("minicube_parameters", {})187    patch_h, patch_w = mcp.get("minicube_size_pixels", 64), mcp.get(188        "minicube_size_pixels", 64189    )190    tp = CONFIG.get("time_parameters", {})191    in_len, pred_len = tp.get("input_len", 6), tp.get("pred_len", 4)192    total_seq = in_len + pred_len193 194    train_opts = CONFIG.get("training_options", {})195    o_dyn_b = train_opts.get("original_dynamic_bands_exported", [])196    o_stat_b = train_opts.get("original_static_bands_exported", [])197    add_b_primary = train_opts.get("additional_bands_names", [])198    final_npz_order = train_opts.get("final_input_band_names_for_npz", [])199 200    # --- Config Validation (Simplified) ---201    if not all([o_dyn_b, o_stat_b, final_npz_order]):202        print("ERR: Base band configs missing.")203        return204    # Further validation as in previous script for add_b lengths vs conceptual can be kept or adapted205 206    add_naming = CONFIG.get("additional_layers_naming", {})207    add_prefix_template = add_naming.get(208        "file_prefix_template", "minicube_{minicube_id}_alllayers_{week_date_nodash}"209    )210    add_suffix = add_naming.get("file_suffix", ".tfrecord")211 212    try:213        ndvi_idx_orig_dyn = o_dyn_b.index("NDVI")214        clear_px_idx_orig_dyn = o_dyn_b.index("NDVI_clear_pixel_count")215    except ValueError:216        print("ERR: NDVI/NDVI_clear_pixel_count missing in original_dynamic_bands.")217        return218 219    num_seq_total = 0220    nodata_f = CONFIG.get("nodata_value_in_export", -9999.0)221    custom_lcu_band_name_in_final_npz = (222        "mmda_LCU"  # The name you use in final_npz_order223    )224 225    for eco_info in CONFIG.get("ecoregions", []):226        eco_n = eco_info["name"]227        print(f"\nProcessing Ecoregion: {eco_n}")228        orig_eco_d = os.path.join(original_tfrecord_dir, eco_n)229        add_eco_d = (230            os.path.join(additional_tfrecord_dir, eco_n)231            if additional_tfrecord_dir232            else None233        )234        eco_out_d = os.path.join(processed_output_base_dir, eco_n)235        os.makedirs(eco_out_d, exist_ok=True)236 237        if not os.path.exists(orig_eco_d):238            print(f"  Orig dir missing for {eco_n}. Skip.")239            continue240 241        has_add_dir = bool(add_eco_d and os.path.exists(add_eco_d))242        can_try_add = bool(add_b_primary)243        if can_try_add and not has_add_dir:244            print(f"  Warn: Add. layers configured, but dir for {eco_n} missing.")245 246        has_custom_lcu_dir = bool(custom_lcu_dir and os.path.exists(custom_lcu_dir))247        if (248            custom_lcu_band_name_in_final_npz in final_npz_order249            and not has_custom_lcu_dir250        ):251            print(252                f"  ERR: Custom LCU specified in final bands, but 'custom_land_cover_dir' missing or invalid. Skipping Ecoregion {eco_n}."253            )254            continue255 256        s_tf_fs = [257            f258            for f in glob.glob(os.path.join(orig_eco_d, "mc_*_static*.tfrecord"))259            if "-mixer.json" not in f260        ]261        mc_ids = set(get_base_prefix_for_mixer(s)[: -len("_static")] for s in s_tf_fs)262        print(f"  Found {len(mc_ids)} unique MC IDs in {eco_n}.")263 264        for mc_id in mc_ids:265            # print(f"    Processing MC: {mc_id}") # Verbose266            f_o_s_tf_paths = glob.glob(267                os.path.join(orig_eco_d, f"{mc_id}_static*.tfrecord")268            )269            f_o_s_tf = [f for f in f_o_s_tf_paths if "-mixer.json" not in f]270            if not f_o_s_tf:271                print(f"      Warn: No orig static for {mc_id}.")272                continue273 274            o_s_arr, _ = read_tfrecord_image_data_flexible(275                f_o_s_tf[0], patch_h, patch_w, o_stat_b, "OrigStatic"276            )277            if o_s_arr is None or o_s_arr.shape[-1] != len(o_stat_b):278                print(f"      Warn: Fail read orig static {mc_id}.")279                continue280 281            # --- Load Custom LCU TIF for this minicube (once per minicube) ---282            custom_lcu_data_arr = None283            if (284                custom_lcu_band_name_in_final_npz in final_npz_order285                and has_custom_lcu_dir286            ):287                # Expected filename: mc_Cherkasy_Croplands_Initial_001.tif (minicube_id directly)288                custom_lcu_tif_path = os.path.join(custom_lcu_dir, f"{mc_id}.tif")289                custom_lcu_data_arr = read_custom_lcu_tif(290                    custom_lcu_tif_path, patch_h, patch_w291                )292                if custom_lcu_data_arr is None:293                    print(294                        f"      Warn: Custom LCU for {mc_id} not found or failed to load from {custom_lcu_tif_path}. Will fill with NoData if required."295                    )296            # --- End Custom LCU Loading ---297 298            o_d_paths = sorted(299                [300                    f301                    for f in glob.glob(302                        os.path.join(orig_eco_d, f"{mc_id}_dynamic_*.tfrecord")303                    )304                    if "-mixer.json" not in f305                ]306            )307            wk_data_mc = []308            for o_d_p in o_d_paths:309                o_d_b = get_base_prefix_for_mixer(o_d_p)310                try:311                    d_s = o_d_b.split("_dynamic_")[1]312                    d_o = datetime.strptime(d_s, "%Y%m%d")313                except (IndexError, ValueError):314                    print(f"        Warn: Can't parse date from {o_d_b} for MC {mc_id}")315                    continue316 317                o_w_d_arr, _ = read_tfrecord_image_data_flexible(318                    o_d_p, patch_h, patch_w, o_dyn_b, "OrigDynamic"319                )320                if o_w_d_arr is None or o_w_d_arr.shape[-1] != len(o_dyn_b):321                    continue322 323                add_w_arr, add_w_bands_read = None, []324                if can_try_add and has_add_dir:325                    add_tf_p = find_additional_layers_tfrecord(326                        add_eco_d, mc_id, d_o, add_prefix_template, add_suffix327                    )328                    if add_tf_p:329                        add_w_arr, add_w_bands_read = read_tfrecord_image_data_flexible(330                            add_tf_p,331                            patch_h,332                            patch_w,333                            add_b_primary,334                            "AdditionalLayers",335                        )336                wk_data_mc.append(337                    {338                        "date": d_o,339                        "orig_dyn": o_w_d_arr,340                        "add_layers_data": add_w_arr,341                        "add_layers_bands_read_ordered": add_w_bands_read,342                    }343                )344 345            if not wk_data_mc:346                print(f"      No valid weekly data loaded for {mc_id}.")347                continue348            wk_data_mc.sort(key=lambda x: x["date"])349 350            n_seq_mc = 0351            if len(wk_data_mc) >= total_seq:352                for i in range(len(wk_data_mc) - total_seq + 1):353                    seq_raw = wk_data_mc[i : i + total_seq]354                    in_t_w_list, avg_c_px_list = [], []355 356                    for t in range(in_len):357                        w_dat = seq_raw[t]358                        curr_w_b_srcs = {}  # {clean_conceptual_name: numpy_array_HxWx1}359 360                        for idx, n in enumerate(o_dyn_b):361                            curr_w_b_srcs[n] = w_dat["orig_dyn"][:, :, idx : idx + 1]362                        for idx, n in enumerate(o_stat_b):363                            curr_w_b_srcs[n] = o_s_arr[:, :, idx : idx + 1]364 365                        if (366                            w_dat["add_layers_data"] is not None367                            and w_dat["add_layers_bands_read_ordered"]368                            and add_b_primary369                        ):370                            if len(w_dat["add_layers_bands_read_ordered"]) == len(371                                add_b_primary372                            ):373                                for k_idx, conceptual_clean_name in enumerate(374                                    add_b_primary375                                ):376                                    if k_idx < w_dat["add_layers_data"].shape[-1]:377                                        curr_w_b_srcs[conceptual_clean_name] = w_dat[378                                            "add_layers_data"379                                        ][:, :, k_idx : k_idx + 1]380 381                        # Add custom LCU data if available (it's static for the minicube)382                        if custom_lcu_data_arr is not None:383                            curr_w_b_srcs[custom_lcu_band_name_in_final_npz] = (384                                custom_lcu_data_arr  # Already (H,W,1)385                            )386 387                        s_d, c_d = get_doy_features(w_dat["date"])388                        curr_w_b_srcs["sin_DOY"] = np.full(389                            (patch_h, patch_w, 1), s_d, dtype=np.float32390                        )391                        curr_w_b_srcs["cos_DOY"] = np.full(392                            (patch_h, patch_w, 1), c_d, dtype=np.float32393                        )394 395                        w_final_ch_list = []396                        for fb_name in final_npz_order:397                            if fb_name in curr_w_b_srcs:398                                w_final_ch_list.append(curr_w_b_srcs[fb_name])399                            else:400                                # This band is in final_npz_order but not found in any source for this week/minicube401                                # print(f"          Warn: Final band '{fb_name}' not sourced for {mc_id} wk {w_dat['date']:%Y%m%d}. Filling NoData.")402                                w_final_ch_list.append(403                                    np.full(404                                        (patch_h, patch_w, 1),405                                        nodata_f,406                                        dtype=np.float32,407                                    )408                                )409 410                        in_t_w_list.append(np.concatenate(w_final_ch_list, axis=-1))411 412                        clear_px_d = w_dat["orig_dyn"][:, :, clear_px_idx_orig_dyn]413                        v_px = clear_px_d[clear_px_d != nodata_f]414                        avg_c_px_list.append(np.mean(v_px) if len(v_px) > 0 else 0.0)415 416                    if len(in_t_w_list) != in_len:417                        continue418 419                    final_in_t = np.stack(in_t_w_list, axis=0)420 421                    tgt_arrs = [info["orig_dyn"] for info in seq_raw[in_len:]]422                    tgt_t_f = np.stack(tgt_arrs, axis=0)423                    final_tgt_t = tgt_t_f[424                        :, :, :, ndvi_idx_orig_dyn : ndvi_idx_orig_dyn + 1425                    ]426 427                    avg_c = (428                        np.nanmean([c for c in avg_c_px_list if not np.isnan(c)])429                        if any(not np.isnan(c) for c in avg_c_px_list)430                        else 0.0431                    )432                    seq_sds = seq_raw[0]["date"].strftime("%Y%m%d")433 434                    out_fname = f"{mc_id}_sequence_{seq_sds}.npz"435                    out_fpath = os.path.join(eco_out_d, out_fname)436 437                    np.savez_compressed(438                        out_fpath,439                        input_data=final_in_t,440                        target_data=final_tgt_t,441                        input_bands=final_npz_order,442                        target_bands=[o_dyn_b[ndvi_idx_orig_dyn]],443                        minicube_id=mc_id,444                        sequence_start_date=seq_sds,445                        avg_input_clear_pixel_count=avg_c,446                    )447                    n_seq_mc += 1448                    num_seq_total += 1449            # print(f"      Processed {n_seq_mc} sequences for {mc_id}.") # Verbose450    print(f"\n--- FINISHED. Total sequences saved: {num_seq_total} ---")451 452 453if __name__ == "__main__":454    parser = argparse.ArgumentParser(455        description="Combines GEE TFRecords and custom LCU into .npz sequences."456    )457    parser.add_argument("config_file", type=str, help="Path to YAML config.")458    args = parser.parse_args()459    print(f"Loading config: {args.config_file}")460    try:461        with open(args.config_file, "r") as f:462            CONFIG = yaml.safe_load(f)463    except FileNotFoundError:464        print(f"Error: Config file '{args.config_file}' not found.")465        exit(1)466    except Exception as e:467        print(f"Error loading config: {e}")468        exit(1)469 470    dirs_c = CONFIG.get("directories", {})471    tr_opts_c = CONFIG.get("training_options", {})472    req_dirs = [473        "local_tfrecord_download_dir",474        "processed_training_data_dir",475    ]  # Additional and custom LCU dirs are optional based on usage476 477    req_band_cfg_keys = [478        "original_dynamic_bands_exported",479        "original_static_bands_exported",480        "final_input_band_names_for_npz",481    ]482 483    if not all(k in dirs_c for k in req_dirs):484        print(485            f"ERR: Required dir paths missing. Need at least: {req_dirs}. Check 'directories'."486        )487        exit(1)488    if not all(k in tr_opts_c for k in req_band_cfg_keys):489        print(490            f"ERR: Required band lists missing. Need at least: {req_band_cfg_keys}. Check 'training_options'."491        )492        exit(1)493    if not tr_opts_c.get("final_input_band_names_for_npz"):494        print("ERR: 'final_input_band_names_for_npz' cannot be empty.")495        exit(1)496 497    # Validate additional bands config if present498    add_b_pri = tr_opts_c.get("additional_bands_names", [])499    if add_b_pri:  # If we intend to use additional bands500        if not dirs_c.get("additional_layers_tfrecord_download_dir"):501            print(502                "ERR: 'additional_bands_names' defined, but 'additional_layers_tfrecord_download_dir' missing."503            )504            exit(1)505 506    # Validate custom LCU config if used507    final_bands_check_lcu = tr_opts_c.get("final_input_band_names_for_npz", [])508    custom_lcu_name_example = "custom_LCU"  # Or whatever you decide to name it509    if custom_lcu_name_example in final_bands_check_lcu and not dirs_c.get(510        "custom_land_cover_dir"511    ):512        print(513            f"ERR: '{custom_lcu_name_example}' in final bands, but 'custom_land_cover_dir' missing in directories config."514        )515        exit(1)516 517    process_minicubes_with_additional_layers(CONFIG)518