225 lines
8.3 KiB
Python
225 lines
8.3 KiB
Python
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()
|