Victor Phan
533 lines
20 KiB
Python
533 lines
20 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 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)
|
|
}
|