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 # Add GDAL optimizations for fast HTTP access os.environ["GDAL_HTTP_MAX_RETRY"] = "5" os.environ["GDAL_HTTP_RETRY_DELAY"] = "2" os.environ["GDAL_HTTP_CONNECTION_TIMEOUT"] = "10" os.environ["GDAL_HTTP_TIMEOUT"] = "30" os.environ["CPL_VSIL_CURL_ALLOWED_EXTENSIONS"] = ".tif,.tiff" os.environ["GDAL_DISABLE_READDIR_ON_OPEN"] = "YES" 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, ) return list(search.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, ) return list(search.items()) def process_point_fast(idx, x_coord, y_coord, label, 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 # Add GDAL config per worker just in case import os os.environ["GDAL_HTTP_MAX_RETRY"] = "5" os.environ["GDAL_HTTP_CONNECTION_TIMEOUT"] = "5" os.environ["GDAL_HTTP_TIMEOUT"] = "10" 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] 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][:8] # Less temporal depth to speed up 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() 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)) # Get up to 4 valid 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: return None # --- 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][:6] 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() 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)) 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 patch_s1 = patch_s1.isel(time=slice(0, min(4, len(patch_s1.time)))) return { 'patch_s2': patch_s2, 'patch_s1': patch_s1, 'label': label } except Exception as e: return None def extract_fusion_fast(s2_items, s1_items, gdf, patch_size=16): print(f"Extracting S1+S2 Fusion patches using 24 parallel jobs (FAST MODE)...") s2_items_dicts = [item.to_dict() for item in s2_items] s1_items_dicts = [item.to_dict() for item in s1_items] # We pass individual scalar values to avoid pickling the whole row object jobs = [] for idx, row in gdf.iterrows(): jobs.append((idx, row.geometry.x, row.geometry.y, row['HT_code'] - 1)) results = Parallel(n_jobs=24, backend="loky", pre_dispatch='1.5*n_jobs')( delayed(process_point_fast)(idx, x, y, lbl, s2_items_dicts, s1_items_dicts, patch_size) for idx, x, y, lbl in tqdm(jobs, total=len(jobs), desc="Downloading S1+S2 Patches") ) X = [] y = [] cloud_remover = DeepInpaintingStrategy(model_path="cloud_removal_model/cloud_removal_unet_best.pth") 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'] patch_cloud_mask = patch_s2["SCL"].isin([3, 8, 9, 10]) clean_patch, _ = cloud_remover.remove_clouds(patch_s2, patch_cloud_mask) b4 = np.clip(clean_patch["B04"].values / 10000.0, 0, 1) b8 = np.clip(clean_patch["B08"].values / 10000.0, 0, 1) b3 = np.clip(clean_patch["B03"].values / 10000.0, 0, 1) b2 = np.clip(clean_patch["B02"].values / 10000.0, 0, 1) ndvi = (b8 - b4) / (b8 + b4 + 1e-6) ndwi = (b3 - b8) / (b3 + b8 + 1e-6) 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) vv = np.clip(patch_s1["vv"].values, 0, 1.0) vh = np.clip(patch_s1["vh"].values, 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) 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 (FAST MODE)") 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_fast(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()