""" 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 sklearn.ensemble import RandomForestClassifier from sklearn.tree import DecisionTreeClassifier from sklearn.svm import SVC from xgboost import XGBClassifier import joblib from datetime import datetime import json import os import warnings import hashlib from pathlib import Path warnings.filterwarnings('ignore') # PyTorch for CNN try: import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim from torch.utils.data import TensorDataset, DataLoader PYTORCH_AVAILABLE = True except ImportError: PYTORCH_AVAILABLE = False print("Warning: PyTorch not available. CNN model will not work.") # Define CNN model class for PyTorch class CNNClassifier(nn.Module): def __init__(self, n_features, n_classes): super(CNNClassifier, self).__init__() self.n_features = n_features self.n_classes = n_classes # For small feature sets (like 3 features), use simpler architecture if n_features < 8: # Simple fully connected network for small features self.use_conv = False self.fc1 = nn.Linear(n_features, 64) self.dropout1 = nn.Dropout(0.3) self.fc2 = nn.Linear(64, 128) self.dropout2 = nn.Dropout(0.5) self.fc3 = nn.Linear(128, n_classes) else: # CNN architecture for larger feature sets self.use_conv = True self.conv1 = nn.Conv1d(in_channels=1, out_channels=32, kernel_size=3, padding=1) self.pool1 = nn.MaxPool1d(kernel_size=2) self.conv2 = nn.Conv1d(in_channels=32, out_channels=64, kernel_size=3, padding=1) self.pool2 = nn.MaxPool1d(kernel_size=2) # Calculate size after convolutions conv_output_size = (n_features // 2 // 2) * 64 # Fully connected layers self.fc1 = nn.Linear(conv_output_size, 128) self.dropout = nn.Dropout(0.5) self.fc2 = nn.Linear(128, n_classes) def forward(self, x): # x shape: (batch, n_features) or (batch, 1, n_features) if self.use_conv: # CNN path for larger feature sets if len(x.shape) == 2: x = x.unsqueeze(1) # Add channel dimension x = F.relu(self.conv1(x)) x = self.pool1(x) x = F.relu(self.conv2(x)) x = self.pool2(x) x = x.view(x.size(0), -1) # Flatten x = F.relu(self.fc1(x)) x = self.dropout(x) x = self.fc2(x) else: # Fully connected path for small feature sets if len(x.shape) == 3: x = x.squeeze(1) # Remove channel dimension if present x = F.relu(self.fc1(x)) x = self.dropout1(x) x = F.relu(self.fc2(x)) x = self.dropout2(x) x = self.fc3(x) return x def predict(self, X): """Scikit-learn style predict method""" self.eval() with torch.no_grad(): if isinstance(X, np.ndarray): X = torch.FloatTensor(X) # Handle both 2D and 3D inputs if not self.use_conv and len(X.shape) == 3: X = X.squeeze(1) elif self.use_conv and len(X.shape) == 2: X = X.unsqueeze(1) outputs = self(X) _, predicted = torch.max(outputs, 1) return predicted.cpu().numpy() def score(self, X, y): """Scikit-learn style score method""" predictions = self.predict(X) if isinstance(y, torch.Tensor): y = y.cpu().numpy() return np.mean(predictions == y) # Microsoft Planetary Computer imports import planetary_computer from pystac_client import Client from odc.stac import load as stac_load # Feature extraction from feature_extractor import get_feature_extractor 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', model_type='xgboost', n_estimators=100, max_depth=20, learning_rate=0.1, use_gpu=True, use_cache=True, test_size=0.2, feature_mode='simple', 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 test_size: Fraction of data to use for test set (0-1) feature_mode: 'simple' (3 features), 'temporal' (39 features), or 'extended' (15 features) 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_{model_type}_{timestamp}.joblib' # ============ CACHE SYSTEM ============ # Create cache directory cache_dir = Path("dataset_cache") cache_dir.mkdir(exist_ok=True) # Generate cache key from parameters cache_params = f"{bbox}_{time_range}_{max_scenes}_{cloud_cover}_{resolution}" cache_key = hashlib.md5(cache_params.encode()).hexdigest() cache_file = cache_dir / f"training_data_{cache_key}.joblib" features = None labels = None # Try to load from cache if use_cache and cache_file.exists(): update_status(f"📦 Loading cached dataset from {cache_file.name}...", 5) try: cached_data = joblib.load(cache_file) features = cached_data['features'] labels = cached_data['labels'] update_status(f"✅ Loaded {len(features)} samples from cache (skipped satellite download!)", 50) except Exception as e: update_status(f"⚠️ Cache load failed: {str(e)}, downloading fresh data...", 10) features = None # If no cache or cache failed, download data if features is None: update_status("📡 Cache not found or disabled, downloading satellite data...", 10) # Connect to Microsoft Planetary Computer update_status("Connecting to Microsoft Planetary Computer...", 12) 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] # Load different bands based on feature mode if feature_mode == 'simple': bands_to_load = ["B04", "B08", "SCL"] else: # temporal or extended bands_to_load = ["B02", "B03", "B04", "B08", "B11", "SCL"] ds_s2 = stac_load( items_s2, bands=bands_to_load, crs="EPSG:32648", resolution=resolution, bbox=bbox, patch_url=planetary_computer.sign, fail_on_error=False, ) # Rename for compatibility (simple mode) if "B04" in ds_s2 and "red" not in ds_s2: 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() # ============ FEATURE EXTRACTION ============ update_status(f"Initializing FeatureExtractor (mode={feature_mode})...", 50) extractor = get_feature_extractor(mode=feature_mode) # 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 using FeatureExtractor update_status("Extracting features from satellite data...", 60) if feature_mode == 'simple': # For simple mode: calculate NDVI first 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) # Extract features at training points 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_masked.sel(x=x_coord, y=y_coord, method='nearest').mean(dim='time').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 = [float(ndvi_val), float(vh_val), float(vv_val)] if not np.isnan(feature_vec).any(): features.append(feature_vec) labels.append(label) except Exception as e: continue features = np.array(features) labels = np.array(labels) else: # temporal or extended mode # Apply cloud mask for temporal/extended modes if 'scl' in ds_s2 or 'SCL' in ds_s2: scl_band = ds_s2['scl'] if 'scl' in ds_s2 else ds_s2['SCL'] cloud_mask = scl_band.isin([1, 3, 8, 9, 10]) for band in ds_s2.data_vars: if band != 'scl' and band != 'SCL': ds_s2[band] = ds_s2[band].where(~cloud_mask) # Extract features at training points 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: # Extract point data from S2 point_s2 = ds_s2.sel(x=x_coord, y=y_coord, method='nearest') # Extract point data from S1 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 # Create minimal dataset for feature extraction point_data = xr.Dataset({ 'B02': point_s2['B02'], 'B03': point_s2['B03'], 'B04': point_s2['B04'], 'B08': point_s2['B08'], 'B11': point_s2['B11'] }) # Create VH/VV DataArrays (without spatial dims, just time if exists) if 'time' in point_data.dims: vh_da = xr.DataArray([vh_val] * len(point_data.time), dims=['time']) vv_da = xr.DataArray([vv_val] * len(point_data.time), dims=['time']) else: vh_da = xr.DataArray([vh_val]) vv_da = xr.DataArray([vv_val]) # Extract features using FeatureExtractor # Note: extractor.extract returns (n_pixels, n_features), we take first row feature_vec = extractor.extract( s2_data=point_data, vh_data=vh_da, vv_data=vv_da ) # If feature_vec is 2D, take first row if len(feature_vec.shape) > 1: feature_vec = feature_vec[0] if not np.isnan(feature_vec).any(): features.append(feature_vec) labels.append(label) except Exception as e: continue features = np.array(features) labels = np.array(labels) check_cancellation() update_status(f"Extracted {len(features)} valid training samples", 70) # ============ SAVE TO CACHE ============ if use_cache: update_status(f"💾 Saving dataset to cache for future use...", 72) try: cache_data = { 'features': features, 'labels': labels, 'bbox': bbox, 'time_range': time_range, 'resolution': resolution, 'feature_mode': feature_mode, 'timestamp': datetime.now().isoformat() } joblib.dump(cache_data, cache_file) update_status(f"✅ Cached to {cache_file.name}", 75) except Exception as e: update_status(f"⚠️ Cache save failed: {str(e)}", 75) # 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=test_size, random_state=42, stratify=labels_encoded ) # Train model based on selected type update_status(f"Training {model_type.upper()} model...", 75) device = 'cuda:0' if use_gpu else 'cpu' if model_type == 'xgboost': model = XGBClassifier( n_estimators=n_estimators, max_depth=max_depth, learning_rate=learning_rate, device=device if use_gpu else 'cpu', tree_method='hist', random_state=42, eval_metric='mlogloss', verbosity=0 ) elif model_type == 'random_forest': model = RandomForestClassifier( n_estimators=n_estimators, max_depth=max_depth, random_state=42, n_jobs=-1, # Use all cores verbose=0 ) elif model_type == 'decision_tree': model = DecisionTreeClassifier( max_depth=max_depth, random_state=42 ) elif model_type == 'svm': model = SVC( kernel='rbf', random_state=42, verbose=False ) elif model_type == 'cnn': if not PYTORCH_AVAILABLE: raise ImportError("PyTorch is required for CNN. Install: pip install torch") # CNN requires reshaping data n_features = X_train.shape[1] n_classes = len(np.unique(y_train)) # Build PyTorch CNN model device = torch.device('cuda' if torch.cuda.is_available() and use_gpu else 'cpu') update_status(f"Building CNN model on {device}...", 75) model = CNNClassifier(n_features, n_classes).to(device) # Convert to PyTorch tensors X_train_tensor = torch.FloatTensor(X_train).unsqueeze(1) # Add channel dim: (N, 1, features) y_train_tensor = torch.LongTensor(y_train) X_test_tensor = torch.FloatTensor(X_test).unsqueeze(1) y_test_tensor = torch.LongTensor(y_test) # Create data loaders train_dataset = TensorDataset(X_train_tensor, y_train_tensor) train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True) # Loss and optimizer criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=0.001) # Train CNN update_status("Training CNN model with PyTorch...", 80) epochs = min(50, n_estimators // 2) # Use n_estimators as epochs model.train() for epoch in range(epochs): epoch_loss = 0.0 for batch_X, batch_y in train_loader: batch_X, batch_y = batch_X.to(device), batch_y.to(device) optimizer.zero_grad() outputs = model(batch_X) loss = criterion(outputs, batch_y) loss.backward() optimizer.step() epoch_loss += loss.item() if (epoch + 1) % 10 == 0: avg_loss = epoch_loss / len(train_loader) update_status(f"CNN Epoch {epoch+1}/{epochs}, Loss: {avg_loss:.4f}", 80 + (epoch / epochs) * 10) # Move model to CPU for saving (compatible with non-GPU systems) model = model.cpu() model.device_used = str(device) else: raise ValueError(f"Unknown model type: {model_type}. Choose: xgboost, random_forest, decision_tree, svm, cnn") # Fit non-CNN models if model_type != 'cnn': model.fit(X_train, y_train) # Evaluate update_status("Evaluating model...", 90) if model_type == 'cnn': # PyTorch CNN evaluation train_score = model.score(X_train, y_train) test_score = model.score(X_test, y_test) y_pred = model.predict(X_test) else: train_score = model.score(X_train, y_train) test_score = model.score(X_test, y_test) y_pred = model.predict(X_test) # Generate classification report and confusion matrix update_status("Generating classification report...", 92) class_names = label_encoder.classes_.tolist() # Classification report as dict from sklearn.metrics import classification_report, confusion_matrix cls_report = classification_report(y_test, y_pred, target_names=class_names, output_dict=True, zero_division=0) # Confusion matrix conf_matrix = confusion_matrix(y_test, y_pred).tolist() # Save model using ModelManager update_status("Saving model...", 95) os.makedirs(os.path.dirname(output_model_path), exist_ok=True) # Get feature names from extractor if feature_mode == 'temporal': # Calculate n_timesteps from data n_timesteps = len(features[0]) // 3 - 1 # (NDVI + NDWI + NDBI) * n_timesteps + 3 radar features feature_names = extractor.get_feature_names(n_timesteps=n_timesteps) else: feature_names = extractor.get_feature_names() # Prepare metadata info = { "timestamp": datetime.now().isoformat(), "data_source": "Microsoft Planetary Computer STAC", "collections": ["sentinel-2-l2a", "sentinel-1-rtc"], "features": feature_names, "feature_mode": feature_mode, "training_samples": len(X_train), "testing_samples": len(X_test), "test_size": test_size, "train_accuracy": float(train_score), "test_accuracy": float(test_score), "model_type": model_type, "device": device if model_type == 'xgboost' else 'cpu', "n_estimators": n_estimators if model_type in ['xgboost', 'random_forest', 'cnn'] else None, "max_depth": max_depth if model_type != 'cnn' else None, "learning_rate": learning_rate if model_type == 'xgboost' else None, "cnn_epochs": min(50, n_estimators // 2) if model_type == 'cnn' else None, "n_features": X_train.shape[1], "n_classes": len(np.unique(y_train)), "class_names": class_names, "classification_report": cls_report, "confusion_matrix": conf_matrix, "bbox": bbox, "time_range": time_range, "resolution": resolution } # Use ModelManager to save from model_manager import get_model_manager model_manager = get_model_manager() model_filename = os.path.basename(output_model_path) model_manager.save_model( model=model, metadata=info, model_filename=model_filename, label_encoder=label_encoder ) 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), "test_size": test_size, "classes": class_names, "classification_report": cls_report, "confusion_matrix": conf_matrix, "model_type": model_type, "bbox": bbox, "time_range": time_range, "resolution": resolution } 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) }