""" 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) }