""" 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 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, 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_{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] 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) # ============ 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, '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=0.2, 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) else: train_score = model.score(X_train, y_train) test_score = 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': 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": 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)), "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) }