import os import gc import json import joblib import numpy as np import pandas as pd import geopandas as gpd import xarray as xr from tqdm import tqdm from joblib import Parallel, delayed import pystac_client import planetary_computer import odc.stac from shapely.geometry import Point, shape from pyproj import Transformer import warnings warnings.filterwarnings('ignore') from core.cloud_removal import DeepInpaintingStrategy def get_s2_items(bbox, time_range): catalog = pystac_client.Client.open( "https://planetarycomputer.microsoft.com/api/stac/v1", modifier=planetary_computer.sign_inplace, ) search = catalog.search( collections=["sentinel-2-l2a"], bbox=bbox, datetime=time_range, ) items = list(search.items()) print(f"Found {len(items)} Sentinel-2 scenes") return items def get_s1_items(bbox, time_range): catalog = pystac_client.Client.open( "https://planetarycomputer.microsoft.com/api/stac/v1", modifier=planetary_computer.sign_inplace, ) search = catalog.search( collections=["sentinel-1-rtc"], bbox=bbox, datetime=time_range, ) items = list(search.items()) print(f"Found {len(items)} Sentinel-1 scenes") return items def process_point_s1_s2(idx, row, s2_items_dicts, s1_items_dicts, patch_size=16): try: import pystac import odc.stac import planetary_computer from shapely.geometry import Point, shape from pyproj import Transformer s2_items = [pystac.Item.from_dict(d) for d in s2_items_dicts] s1_items = [pystac.Item.from_dict(d) for d in s1_items_dicts] x_coord = row['geometry'].x y_coord = row['geometry'].y transformer = Transformer.from_crs("epsg:32648", "epsg:4326", always_xy=True) lon, lat = transformer.transform(x_coord, y_coord) point = Point(lon, lat) # --- SENTINEL-2 --- filtered_s2 = [item for item in s2_items if shape(item.geometry).contains(point)] if not filtered_s2: return None filtered_s2 = [planetary_computer.sign(item) for item in filtered_s2][:10] patch_s2 = odc.stac.load( filtered_s2, bands=["B02", "B03", "B04", "B08", "SCL"], x=(x_coord - 100, x_coord + 100), y=(y_coord - 100, y_coord + 100), crs="EPSG:32648", resolution=10, patch_url=planetary_computer.sign, fail_on_error=False ).compute() b2_sums = patch_s2["B02"].sum(dim=["x", "y"]) valid_times = b2_sums > 0 patch_s2 = patch_s2.isel(time=valid_times) if len(patch_s2.time) == 0: return None patch_s2 = patch_s2.isel(time=slice(0, min(4, len(patch_s2.time)))) if "SCL" not in patch_s2 or "B02" not in patch_s2: return None if patch_s2.dims['x'] < patch_size or patch_s2.dims['y'] < patch_size: return None patch_s2 = patch_s2.isel(x=slice(0, patch_size), y=slice(0, patch_size)) # --- SENTINEL-1 --- filtered_s1 = [item for item in s1_items if shape(item.geometry).contains(point)] if not filtered_s1: return None filtered_s1 = [planetary_computer.sign(item) for item in filtered_s1][:10] patch_s1 = odc.stac.load( filtered_s1, bands=["vv", "vh"], x=(x_coord - 100, x_coord + 100), y=(y_coord - 100, y_coord + 100), crs="EPSG:32648", resolution=10, patch_url=planetary_computer.sign, fail_on_error=False ).compute() vv_sums = patch_s1["vv"].sum(dim=["x", "y"]) valid_s1_times = vv_sums > 0 patch_s1 = patch_s1.isel(time=valid_s1_times) if len(patch_s1.time) == 0: return None # take up to 4 timesteps to match S2 patch_s1 = patch_s1.isel(time=slice(0, min(4, len(patch_s1.time)))) if patch_s1.dims['x'] < patch_size or patch_s1.dims['y'] < patch_size: return None patch_s1 = patch_s1.isel(x=slice(0, patch_size), y=slice(0, patch_size)) return { 'patch_s2': patch_s2, 'patch_s1': patch_s1, 'label': row['HT_code'] - 1 } except Exception as e: return None def extract_fusion_patches(s2_items, s1_items, gdf, patch_size=16): print(f"Extracting S1+S2 Fusion patches for {len(gdf)} points using 8 parallel jobs...") s2_items_dicts = [item.to_dict() for item in s2_items] s1_items_dicts = [item.to_dict() for item in s1_items] results = Parallel(n_jobs=8, backend="loky")( delayed(process_point_s1_s2)(idx, row, s2_items_dicts, s1_items_dicts, patch_size) for idx, row in tqdm(gdf.iterrows(), total=len(gdf), desc="Downloading S1+S2 Patches") ) X = [] y = [] cloud_remover = DeepInpaintingStrategy(model_path="cloud_removal_model/cloud_removal_unet_best.pth") if cloud_remover.model is None: print("Warning: Could not load DeepInpainting model.") print("Applying Cloud Removal & Merging Sentinel-1...") valid_results = [r for r in results if r is not None] print(f"Valid points extracted: {len(valid_results)}/{len(gdf)}") for res in tqdm(valid_results, desc="Processing Fusion Features"): try: patch_s2 = res['patch_s2'] patch_s1 = res['patch_s1'] label = res['label'] # --- PROCESS S2 --- patch_cloud_mask = patch_s2["SCL"].isin([3, 8, 9, 10]) clean_patch, _ = cloud_remover.remove_clouds(patch_s2, patch_cloud_mask) b4 = clean_patch["B04"].values b8 = clean_patch["B08"].values b3 = clean_patch["B03"].values b2 = clean_patch["B02"].values ndvi = (b8 - b4) / (b8 + b4 + 1e-6) ndwi = (b3 - b8) / (b3 + b8 + 1e-6) b2 = np.clip(b2 / 10000.0, 0, 1) b3 = np.clip(b3 / 10000.0, 0, 1) b4 = np.clip(b4 / 10000.0, 0, 1) b8 = np.clip(b8 / 10000.0, 0, 1) features_t_s2 = np.stack([b2, b3, b4, b8, ndvi, ndwi], axis=1) # (time, 6, 16, 16) t_len = features_t_s2.shape[0] if t_len < 4: pad = np.zeros((4 - t_len, 6, 16, 16)) features_t_s2 = np.concatenate([features_t_s2, pad], axis=0) # --- PROCESS S1 --- vv = patch_s1["vv"].values vh = patch_s1["vh"].values vv = np.clip(vv, 0, 1.0) vh = np.clip(vh, 0, 1.0) features_t_s1 = np.stack([vv, vh], axis=1) # (time, 2, 16, 16) t_len_s1 = features_t_s1.shape[0] if t_len_s1 < 4: pad_s1 = np.zeros((4 - t_len_s1, 2, 16, 16)) features_t_s1 = np.concatenate([features_t_s1, pad_s1], axis=0) # --- MERGE S1 and S2 --- features_t = np.concatenate([features_t_s2, features_t_s1], axis=1) # (4, 8, 16, 16) features = features_t.reshape(32, 16, 16) features = np.nan_to_num(features, nan=0.0) X.append(features) y.append(label) except Exception as e: pass return np.array(X), np.array(y) def main(): print("🚀 BẮT ĐẦU TRÍCH XUẤT FUSION S1 + S2 (32-CHANNELS)") cache_file = "dataset_cache/training_data_fusion_32ch.joblib" bbox = [105.5, 9.2, 106.3, 10.0] time_range = "2023-01-01/2023-04-30" s2_items = get_s2_items(bbox, time_range) s1_items = get_s1_items(bbox, time_range) gdf = gpd.read_file("train/ST_training_data_updated_1130points_new.shp") gdf = gdf.to_crs("EPSG:32648") X, y = extract_fusion_patches(s2_items, s1_items, gdf, patch_size=16) print(f"Final extracted shape: X={X.shape}, y={y.shape}") os.makedirs('dataset_cache', exist_ok=True) joblib.dump({'X': X, 'y': y}, cache_file) print(f"Saved 32-channel Fusion cache to {cache_file}") if __name__ == "__main__": main()