Victor Phan
303 lines
10 KiB
Python
303 lines
10 KiB
Python
"""
|
|
Training module for land classification using Sentinel-2 and Sentinel-1 data
|
|
from Microsoft Planetary Computer STAC API
|
|
"""
|
|
|
|
import numpy as np
|
|
import xarray as xr
|
|
import geopandas as gpd
|
|
from sklearn.model_selection import train_test_split
|
|
from sklearn.preprocessing import LabelEncoder
|
|
from sklearn.metrics import classification_report, confusion_matrix
|
|
from xgboost import XGBClassifier
|
|
import joblib
|
|
from datetime import datetime
|
|
import json
|
|
import os
|
|
|
|
# Microsoft Planetary Computer imports
|
|
import planetary_computer
|
|
from pystac_client import Client
|
|
from odc.stac import load as stac_load
|
|
|
|
|
|
def train_model(
|
|
bbox=[105.6, 9.3, 106.2, 9.8],
|
|
time_range='2023-03-01/2023-05-31',
|
|
max_scenes=12,
|
|
cloud_cover=30,
|
|
resolution=20,
|
|
training_shapefile='train/ST_training data_updated_1130points_new.shp',
|
|
n_estimators=100,
|
|
max_depth=20,
|
|
learning_rate=0.1,
|
|
use_gpu=True,
|
|
output_model_path=None,
|
|
status_callback=None,
|
|
cancel_check=None
|
|
):
|
|
"""
|
|
Train a land classification model using Sentinel-2 and Sentinel-1 data
|
|
|
|
Args:
|
|
bbox: [min_lon, min_lat, max_lon, max_lat]
|
|
time_range: "YYYY-MM-DD/YYYY-MM-DD"
|
|
max_scenes: maximum number of scenes to load
|
|
cloud_cover: maximum cloud cover percentage
|
|
resolution: resolution in meters (e.g., 20)
|
|
training_shapefile: path to training shapefile
|
|
n_estimators: number of trees for XGBoost
|
|
max_depth: maximum tree depth
|
|
learning_rate: learning rate for XGBoost
|
|
use_gpu: whether to use GPU for training
|
|
output_model_path: path to save trained model (auto-generated if None)
|
|
status_callback: Optional callback function to report progress
|
|
cancel_check: Optional function that returns True if training should be cancelled
|
|
|
|
Returns:
|
|
Dictionary containing training results
|
|
"""
|
|
|
|
def update_status(message, progress=None):
|
|
"""Helper to update status"""
|
|
if status_callback:
|
|
# Try calling with both arguments, fallback to just message
|
|
try:
|
|
status_callback(message, progress)
|
|
except TypeError:
|
|
status_callback(message)
|
|
print(message)
|
|
|
|
def check_cancellation():
|
|
"""Check if training should be cancelled"""
|
|
if cancel_check and cancel_check():
|
|
raise InterruptedError("Training cancelled by user")
|
|
|
|
try:
|
|
# Auto-generate output path if not provided
|
|
if output_model_path is None:
|
|
timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
|
|
output_model_path = f'model_train/model_xgboost_gpu_{timestamp}.joblib'
|
|
|
|
# Connect to Microsoft Planetary Computer
|
|
update_status("Connecting to Microsoft Planetary Computer...", 0)
|
|
catalog = Client.open("https://planetarycomputer.microsoft.com/api/stac/v1")
|
|
check_cancellation()
|
|
|
|
# Search for Sentinel-2 scenes
|
|
update_status("Searching for Sentinel-2 scenes...", 10)
|
|
query_s2 = catalog.search(
|
|
collections=["sentinel-2-l2a"],
|
|
bbox=bbox,
|
|
datetime=time_range,
|
|
query={"eo:cloud_cover": {"lt": cloud_cover}}
|
|
)
|
|
items_s2 = list(query_s2.item_collection())
|
|
|
|
check_cancellation()
|
|
|
|
# Limit scenes
|
|
if len(items_s2) > max_scenes:
|
|
step = len(items_s2) // max_scenes
|
|
items_s2 = items_s2[::step][:max_scenes]
|
|
|
|
update_status(f"Found {len(items_s2)} Sentinel-2 scenes", 20)
|
|
|
|
# Sign and load Sentinel-2 data
|
|
update_status("Loading Sentinel-2 data...", 25)
|
|
items_s2 = [planetary_computer.sign(item) for item in items_s2]
|
|
ds_s2 = stac_load(
|
|
items_s2,
|
|
bands=["B04", "B08", "SCL"],
|
|
crs="EPSG:32648",
|
|
resolution=resolution,
|
|
bbox=bbox,
|
|
patch_url=planetary_computer.sign,
|
|
fail_on_error=False,
|
|
)
|
|
ds_s2 = ds_s2.rename({"B04": "red", "B08": "nir", "SCL": "scl"})
|
|
|
|
check_cancellation()
|
|
|
|
# Search for Sentinel-1 scenes
|
|
update_status("Searching for Sentinel-1 scenes...", 35)
|
|
query_s1 = catalog.search(
|
|
collections=["sentinel-1-rtc"],
|
|
bbox=bbox,
|
|
datetime=time_range,
|
|
)
|
|
items_s1 = list(query_s1.item_collection())
|
|
|
|
# Limit scenes
|
|
if len(items_s1) > max_scenes:
|
|
step = len(items_s1) // max_scenes
|
|
items_s1 = items_s1[::step][:max_scenes]
|
|
|
|
update_status(f"Found {len(items_s1)} Sentinel-1 scenes", 40)
|
|
|
|
# Sign and load Sentinel-1 data
|
|
update_status("Loading Sentinel-1 data...", 45)
|
|
items_s1 = [planetary_computer.sign(item) for item in items_s1]
|
|
ds_s1 = stac_load(
|
|
items_s1,
|
|
bands=["vv", "vh"],
|
|
crs="EPSG:32648",
|
|
resolution=resolution,
|
|
bbox=bbox,
|
|
patch_url=planetary_computer.sign,
|
|
fail_on_error=False,
|
|
)
|
|
|
|
# Convert to dB
|
|
ds_s1['vv_db'] = 10 * np.log10(ds_s1['vv'].where(ds_s1['vv'] > 0))
|
|
ds_s1['vh_db'] = 10 * np.log10(ds_s1['vh'].where(ds_s1['vh'] > 0))
|
|
|
|
check_cancellation()
|
|
|
|
# Calculate NDVI
|
|
update_status("Calculating NDVI...", 50)
|
|
ndvi = (ds_s2['nir'] - ds_s2['red']) / (ds_s2['nir'] + ds_s2['red'] + 1e-8)
|
|
|
|
# Apply cloud mask
|
|
cloud_mask = ds_s2['scl'].isin([1, 3, 8, 9, 10])
|
|
ndvi_masked = ndvi.where(~cloud_mask)
|
|
ndvi_mean = ndvi_masked.mean(dim='time')
|
|
|
|
# Load training data
|
|
update_status("Loading training data...", 55)
|
|
train_gdf = gpd.read_file(training_shapefile)
|
|
|
|
if train_gdf.crs != 'EPSG:32648':
|
|
train_gdf = train_gdf.to_crs('EPSG:32648')
|
|
|
|
# Auto-detect label column
|
|
label_column = None
|
|
for col in ['HT_code', 'Ma_LU', 'LU2022', 'Hientrang', 'class', 'Class', 'CLASS']:
|
|
if col in train_gdf.columns:
|
|
label_column = col
|
|
break
|
|
|
|
if label_column is None:
|
|
raise ValueError(f"Cannot find label column in shapefile. Available: {list(train_gdf.columns)}")
|
|
|
|
# Extract features
|
|
update_status("Extracting features from training points...", 60)
|
|
features = []
|
|
labels = []
|
|
|
|
for idx, row in train_gdf.iterrows():
|
|
point = row.geometry
|
|
x_coord = point.x
|
|
y_coord = point.y
|
|
label = row[label_column]
|
|
|
|
try:
|
|
ndvi_val = ndvi_mean.sel(x=x_coord, y=y_coord, method='nearest').values
|
|
vh_val = ds_s1['vh_db'].sel(x=x_coord, y=y_coord, method='nearest').mean(dim='time').values
|
|
vv_val = ds_s1['vv_db'].sel(x=x_coord, y=y_coord, method='nearest').mean(dim='time').values
|
|
|
|
feature_vec = [ndvi_val, vh_val, vv_val]
|
|
|
|
if not np.isnan(feature_vec).any():
|
|
features.append(feature_vec)
|
|
labels.append(label)
|
|
except:
|
|
continue
|
|
|
|
features = np.array(features)
|
|
labels = np.array(labels)
|
|
|
|
check_cancellation()
|
|
|
|
update_status(f"Extracted {len(features)} valid training samples", 70)
|
|
|
|
# Encode labels
|
|
label_encoder = LabelEncoder()
|
|
labels_encoded = label_encoder.fit_transform(labels)
|
|
|
|
# Split data
|
|
X_train, X_test, y_train, y_test = train_test_split(
|
|
features, labels_encoded, test_size=0.2, random_state=42, stratify=labels_encoded
|
|
)
|
|
|
|
# Train XGBoost model
|
|
update_status("Training XGBoost model on GPU...", 75)
|
|
|
|
device = 'cuda:0' if use_gpu else 'cpu'
|
|
|
|
xgb_model = XGBClassifier(
|
|
n_estimators=n_estimators,
|
|
max_depth=max_depth,
|
|
learning_rate=learning_rate,
|
|
device=device,
|
|
tree_method='hist',
|
|
random_state=42,
|
|
eval_metric='mlogloss',
|
|
verbosity=0
|
|
)
|
|
|
|
xgb_model.fit(X_train, y_train)
|
|
|
|
# Evaluate
|
|
update_status("Evaluating model...", 90)
|
|
train_score = xgb_model.score(X_train, y_train)
|
|
test_score = xgb_model.score(X_test, y_test)
|
|
|
|
# Save model
|
|
update_status("Saving model...", 95)
|
|
os.makedirs(os.path.dirname(output_model_path), exist_ok=True)
|
|
joblib.dump({'model': xgb_model, 'label_encoder': label_encoder}, output_model_path)
|
|
|
|
# Save model info
|
|
info = {
|
|
"timestamp": datetime.now().isoformat(),
|
|
"data_source": "Microsoft Planetary Computer STAC",
|
|
"collections": ["sentinel-2-l2a", "sentinel-1-rtc"],
|
|
"features": ["NDVI_mean", "VH_dB_mean", "VV_dB_mean"],
|
|
"training_samples": len(X_train),
|
|
"testing_samples": len(X_test),
|
|
"train_accuracy": float(train_score),
|
|
"test_accuracy": float(test_score),
|
|
"model_type": "XGBClassifier",
|
|
"device": device,
|
|
"tree_method": "hist",
|
|
"n_estimators": n_estimators,
|
|
"max_depth": max_depth,
|
|
"learning_rate": learning_rate,
|
|
"bbox": bbox,
|
|
"time_range": time_range,
|
|
"resolution": resolution
|
|
}
|
|
|
|
info_path = output_model_path.replace('.joblib', '_info.json')
|
|
with open(info_path, 'w') as f:
|
|
json.dump(info, f, indent=2)
|
|
|
|
update_status("Training complete!", 100)
|
|
|
|
return {
|
|
"success": True,
|
|
"model_path": output_model_path,
|
|
"info_path": info_path,
|
|
"train_accuracy": train_score,
|
|
"test_accuracy": test_score,
|
|
"training_samples": len(X_train),
|
|
"testing_samples": len(X_test),
|
|
"classes": label_encoder.classes_.tolist()
|
|
}
|
|
|
|
except InterruptedError as e:
|
|
update_status(f"Cancelled: {str(e)}", -1)
|
|
return {
|
|
"success": False,
|
|
"error": str(e),
|
|
"cancelled": True
|
|
}
|
|
|
|
except Exception as e:
|
|
update_status(f"Error: {str(e)}", -1)
|
|
return {
|
|
"success": False,
|
|
"error": str(e)
|
|
}
|