1399 lines
60 KiB
Python
1399 lines
60 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
|
|
from lightgbm import LGBMClassifier
|
|
import joblib
|
|
from datetime import datetime
|
|
import json
|
|
import os
|
|
import warnings
|
|
import hashlib
|
|
from pathlib import Path
|
|
warnings.filterwarnings('ignore')
|
|
|
|
# PyTorch for CNN and advanced models
|
|
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
|
|
import torchvision.models as models
|
|
PYTORCH_AVAILABLE = True
|
|
except ImportError:
|
|
PYTORCH_AVAILABLE = False
|
|
print("Warning: PyTorch not available. CNN and advanced models 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)
|
|
|
|
|
|
# Swin-UNet Classifier for feature vectors
|
|
class SwinUNetClassifier(nn.Module):
|
|
"""
|
|
Swin Transformer U-Net style architecture adapted for feature vector classification.
|
|
Combines hierarchical Swin Transformer blocks with skip connections.
|
|
"""
|
|
def __init__(self, n_features, n_classes, embed_dim=128, depths=(2, 2, 6, 2), num_heads=(4, 8, 16, 32)):
|
|
super(SwinUNetClassifier, self).__init__()
|
|
self.n_features = n_features
|
|
self.n_classes = n_classes
|
|
self.embed_dim = embed_dim
|
|
|
|
# Feature adapter - convert input features to embedding
|
|
self.adapter = nn.Sequential(
|
|
nn.Linear(n_features, embed_dim * 2),
|
|
nn.ReLU(),
|
|
nn.Dropout(0.1),
|
|
nn.Linear(embed_dim * 2, embed_dim)
|
|
)
|
|
|
|
# Encoder path with hierarchical structure
|
|
# Stage 1 - 1/4 resolution
|
|
self.encoder1 = nn.Sequential(
|
|
nn.Linear(embed_dim, embed_dim),
|
|
nn.LayerNorm(embed_dim),
|
|
nn.GELU(),
|
|
nn.Dropout(0.1)
|
|
)
|
|
self.down1 = nn.Linear(embed_dim, embed_dim * 2)
|
|
|
|
# Stage 2 - 1/8 resolution
|
|
self.encoder2 = nn.Sequential(
|
|
nn.Linear(embed_dim * 2, embed_dim * 2),
|
|
nn.LayerNorm(embed_dim * 2),
|
|
nn.GELU(),
|
|
nn.Dropout(0.1)
|
|
)
|
|
self.down2 = nn.Linear(embed_dim * 2, embed_dim * 4)
|
|
|
|
# Stage 3 - 1/16 resolution (bottleneck)
|
|
self.encoder3 = nn.Sequential(
|
|
nn.Linear(embed_dim * 4, embed_dim * 4),
|
|
nn.LayerNorm(embed_dim * 4),
|
|
nn.GELU(),
|
|
nn.Dropout(0.1)
|
|
)
|
|
|
|
# Decoder path with skip connections
|
|
self.up2 = nn.Linear(embed_dim * 4, embed_dim * 2)
|
|
self.decoder2 = nn.Sequential(
|
|
nn.Linear(embed_dim * 4, embed_dim * 2), # Concatenated with skip
|
|
nn.LayerNorm(embed_dim * 2),
|
|
nn.GELU(),
|
|
nn.Dropout(0.1)
|
|
)
|
|
|
|
self.up1 = nn.Linear(embed_dim * 2, embed_dim)
|
|
self.decoder1 = nn.Sequential(
|
|
nn.Linear(embed_dim * 2, embed_dim), # Concatenated with skip
|
|
nn.LayerNorm(embed_dim),
|
|
nn.GELU(),
|
|
nn.Dropout(0.1)
|
|
)
|
|
|
|
# Classification head
|
|
self.classifier = nn.Sequential(
|
|
nn.Linear(embed_dim, embed_dim // 2),
|
|
nn.GELU(),
|
|
nn.Dropout(0.3),
|
|
nn.Linear(embed_dim // 2, n_classes)
|
|
)
|
|
|
|
# Attention mechanism for better feature aggregation
|
|
self.attention = nn.MultiheadAttention(embed_dim, num_heads=4, batch_first=True)
|
|
|
|
def forward(self, x):
|
|
# x shape: (batch, n_features)
|
|
if len(x.shape) == 3:
|
|
x = x.squeeze(1)
|
|
|
|
batch_size = x.shape[0]
|
|
|
|
# Feature adaptation
|
|
x = self.adapter(x) # (batch, embed_dim)
|
|
|
|
# Add sequence dimension for attention (treat as sequence of length 1)
|
|
x_seq = x.unsqueeze(1) # (batch, 1, embed_dim)
|
|
|
|
# Encoder path
|
|
# Stage 1
|
|
x1 = self.encoder1(x_seq) # (batch, 1, embed_dim)
|
|
x_down1 = self.down1(x1.squeeze(1)) # (batch, embed_dim*2)
|
|
|
|
# Stage 2
|
|
x2 = self.encoder2(x_down1.unsqueeze(1)) # (batch, 1, embed_dim*2)
|
|
x_down2 = self.down2(x2.squeeze(1)) # (batch, embed_dim*4)
|
|
|
|
# Stage 3 (bottleneck)
|
|
x3 = self.encoder3(x_down2.unsqueeze(1)) # (batch, 1, embed_dim*4)
|
|
|
|
# Decoder path with skip connections
|
|
# Up2
|
|
x_up2 = self.up2(x3.squeeze(1)) # (batch, embed_dim*2)
|
|
x_cat2 = torch.cat([x_up2, x_down1], dim=1) # (batch, embed_dim*4) - concatenate skip
|
|
# Create proper 3D tensor for decoder
|
|
x_cat2_seq = x_cat2.unsqueeze(1) # (batch, 1, embed_dim*4)
|
|
x_dec2 = self.decoder2(x_cat2) # (batch, embed_dim*2)
|
|
|
|
# Up1
|
|
x_up1 = self.up1(x_dec2) # (batch, embed_dim)
|
|
x_cat1 = torch.cat([x_up1, x.squeeze(1)], dim=1) # (batch, embed_dim*2) - concatenate skip
|
|
x_dec1 = self.decoder1(x_cat1) # (batch, embed_dim)
|
|
|
|
# Apply attention mechanism for better aggregation
|
|
x_dec1_seq = x_dec1.unsqueeze(1) # (batch, 1, embed_dim)
|
|
attn_out, _ = self.attention(x_dec1_seq, x_dec1_seq, x_dec1_seq)
|
|
|
|
# Classification
|
|
output = self.classifier(attn_out.squeeze(1))
|
|
return output
|
|
|
|
def predict(self, X):
|
|
"""Scikit-learn style predict"""
|
|
self.eval()
|
|
with torch.no_grad():
|
|
if isinstance(X, np.ndarray):
|
|
X = torch.FloatTensor(X)
|
|
outputs = self(X)
|
|
_, predicted = torch.max(outputs, 1)
|
|
return predicted.cpu().numpy()
|
|
|
|
def score(self, X, y):
|
|
"""Scikit-learn style score"""
|
|
predictions = self.predict(X)
|
|
if isinstance(y, torch.Tensor):
|
|
y = y.cpu().numpy()
|
|
return np.mean(predictions == y)
|
|
|
|
|
|
# MobileNetV3 + LR-ASPP Classifier
|
|
class MobileNetLRASPPClassifier(nn.Module):
|
|
"""
|
|
MobileNetV3 backbone with LR-ASPP (Lite Reduced Atrous Spatial Pyramid Pooling) for semantic segmentation
|
|
Lightweight architecture optimized for efficiency and speed
|
|
"""
|
|
def __init__(self, n_features, n_classes):
|
|
super(MobileNetLRASPPClassifier, self).__init__()
|
|
self.n_features = n_features
|
|
self.n_classes = n_classes
|
|
|
|
# Feature extraction layers (MobileNetV3-inspired)
|
|
self.feature_extractor = nn.Sequential(
|
|
nn.Linear(n_features, 128),
|
|
nn.BatchNorm1d(128),
|
|
nn.ReLU(inplace=True),
|
|
nn.Dropout(0.2),
|
|
|
|
nn.Linear(128, 256),
|
|
nn.BatchNorm1d(256),
|
|
nn.ReLU(inplace=True),
|
|
nn.Dropout(0.3),
|
|
|
|
nn.Linear(256, 512),
|
|
nn.BatchNorm1d(512),
|
|
nn.ReLU(inplace=True),
|
|
nn.Dropout(0.3),
|
|
)
|
|
|
|
# LR-ASPP head (simplified for feature vectors)
|
|
# Branch 1: Global average pooling
|
|
self.global_pool = nn.AdaptiveAvgPool1d(1)
|
|
self.global_conv = nn.Sequential(
|
|
nn.Linear(512, 128),
|
|
nn.ReLU(inplace=True)
|
|
)
|
|
|
|
# Branch 2: 1x1 convolution equivalent
|
|
self.branch_conv = nn.Sequential(
|
|
nn.Linear(512, 128),
|
|
nn.BatchNorm1d(128),
|
|
nn.ReLU(inplace=True)
|
|
)
|
|
|
|
# Fusion and classification
|
|
self.classifier = nn.Sequential(
|
|
nn.Linear(256, 128), # 128 from global + 128 from branch
|
|
nn.BatchNorm1d(128),
|
|
nn.ReLU(inplace=True),
|
|
nn.Dropout(0.4),
|
|
nn.Linear(128, n_classes)
|
|
)
|
|
|
|
def forward(self, x):
|
|
# x shape: (batch, n_features)
|
|
features = self.feature_extractor(x)
|
|
|
|
# LR-ASPP head
|
|
# Branch 1: Global pooling
|
|
global_feat = self.global_pool(features.unsqueeze(-1)).squeeze(-1)
|
|
global_feat = self.global_conv(global_feat)
|
|
|
|
# Branch 2: Direct features
|
|
branch_feat = self.branch_conv(features)
|
|
|
|
# Concatenate branches
|
|
fused = torch.cat([global_feat, branch_feat], dim=1)
|
|
|
|
# Classification
|
|
output = self.classifier(fused)
|
|
return output
|
|
|
|
def predict(self, X):
|
|
"""Scikit-learn style predict"""
|
|
self.eval()
|
|
with torch.no_grad():
|
|
if isinstance(X, np.ndarray):
|
|
X = torch.FloatTensor(X)
|
|
outputs = self(X)
|
|
_, predicted = torch.max(outputs, 1)
|
|
return predicted.cpu().numpy()
|
|
|
|
def score(self, X, y):
|
|
"""Scikit-learn style score"""
|
|
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='odc', # ODC mode: 8 features (NDVI stats + NDWI/NDBI/EVI) for better accuracy
|
|
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), 'extended' (15 features), or 'odc' (8 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
|
|
|
|
# Initialize FeatureExtractor early (will be used for temporal/extended modes)
|
|
update_status(f"Initializing FeatureExtractor (mode={feature_mode})...", 5)
|
|
extractor = get_feature_extractor(mode=feature_mode)
|
|
|
|
# Try to load from cache
|
|
if use_cache and cache_file.exists():
|
|
update_status(f"📦 Đang load cache: {cache_file.name}...", 5)
|
|
try:
|
|
cached_data = joblib.load(cache_file)
|
|
features = cached_data['features']
|
|
labels = cached_data['labels']
|
|
|
|
# Validate cached data
|
|
if len(features) == 0:
|
|
update_status(
|
|
f"❌ Cache rỗng (0 samples)! Đây là cache từ lần training thất bại trước.\n"
|
|
f" Nguyên nhân: Bbox không overlap với shapefile HOẶC tất cả điểm bị NaN.\n"
|
|
f" Đang xóa cache lỗi và tải lại dữ liệu...", 10
|
|
)
|
|
cache_file.unlink() # Delete empty cache
|
|
features = None
|
|
else:
|
|
update_status(
|
|
f"✅ Loaded {len(features)} samples từ cache!\n"
|
|
f" ⚡ Đã bỏ qua download위성 data (tiết kiệm thời gian)", 50
|
|
)
|
|
print(f"[CACHE HIT] Using cached dataset with {len(features)} samples")
|
|
except Exception as e:
|
|
update_status(f"⚠️ Cache bị lỗi: {str(e)}\n Đang tải lại dữ liệu mới...", 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: # odc, temporal, or extended - all need full spectral bands
|
|
bands_to_load = ["B02", "B03", "B04", "B08", "B11", "SCL"]
|
|
|
|
update_status(f"Loading bands: {bands_to_load} for mode={feature_mode}", 26)
|
|
|
|
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,
|
|
)
|
|
|
|
# Debug: Print S2 data info
|
|
print(f"[DEBUG S2] Loaded S2 data")
|
|
print(f"[DEBUG S2] Dimensions: {dict(ds_s2.dims)}")
|
|
print(f"[DEBUG S2] Bands: {list(ds_s2.data_vars)}")
|
|
print(f"[DEBUG S2] CRS: {ds_s2.rio.crs if hasattr(ds_s2, 'rio') else 'No CRS'}")
|
|
print(f"[DEBUG S2] Spatial bounds: x=[{float(ds_s2.x.min())}, {float(ds_s2.x.max())}], y=[{float(ds_s2.y.min())}, {float(ds_s2.y.max())}]")
|
|
if 'time' in ds_s2.dims:
|
|
print(f"[DEBUG S2] Time range: {ds_s2.time.min().values} to {ds_s2.time.max().values}")
|
|
|
|
# Rename bands ONLY for simple mode (simple mode uses 'red', 'nir', 'scl' names)
|
|
# Other modes (odc, extended, temporal) use original band names (B02, B03, B04, B08, B11, SCL)
|
|
if feature_mode == 'simple' and "B04" in ds_s2 and "red" not in ds_s2:
|
|
ds_s2 = ds_s2.rename({"B04": "red", "B08": "nir", "SCL": "scl"})
|
|
print(f"[DEBUG S2] Renamed bands for simple mode: 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))
|
|
|
|
# Debug: Print S1 data info
|
|
print(f"[DEBUG S1] Loaded S1 data")
|
|
print(f"[DEBUG S1] Dimensions: {dict(ds_s1.dims)}")
|
|
print(f"[DEBUG S1] Bands: {list(ds_s1.data_vars)}")
|
|
print(f"[DEBUG S1] Spatial bounds: x=[{float(ds_s1.x.min())}, {float(ds_s1.x.max())}], y=[{float(ds_s1.y.min())}, {float(ds_s1.y.max())}]")
|
|
|
|
check_cancellation()
|
|
|
|
# Load training data
|
|
update_status("Loading training data...", 55)
|
|
|
|
# Normalize training shapefile path
|
|
# If path doesn't start with 'train/', add it
|
|
if not training_shapefile.startswith('train/'):
|
|
training_shapefile = f'train/{training_shapefile}'
|
|
|
|
print(f"[DEBUG] Original training shapefile: {training_shapefile}")
|
|
print(f"[DEBUG] Current working directory: {os.getcwd()}")
|
|
|
|
# Try to find the file with exact name first
|
|
if not os.path.exists(training_shapefile):
|
|
# File not found, try to find similar files in train directory
|
|
train_dir = Path('train')
|
|
if train_dir.exists():
|
|
# List all .shp files
|
|
shp_files = list(train_dir.glob('*.shp'))
|
|
print(f"[DEBUG] Available shapefile files in train/:")
|
|
for f in shp_files:
|
|
print(f" - {f.name}")
|
|
|
|
# Try to find a matching file (case-insensitive, ignore underscores vs spaces)
|
|
filename_normalized = os.path.basename(training_shapefile).lower().replace('_', ' ')
|
|
for shp_file in shp_files:
|
|
if shp_file.name.lower().replace('_', ' ') == filename_normalized:
|
|
print(f"[DEBUG] Found matching file: {shp_file}")
|
|
training_shapefile = str(shp_file)
|
|
break
|
|
|
|
if not os.path.exists(training_shapefile):
|
|
raise FileNotFoundError(
|
|
f"Training shapefile not found: {training_shapefile}\n"
|
|
f"Available files: {[f.name for f in shp_files]}"
|
|
)
|
|
else:
|
|
raise FileNotFoundError(f"Train directory not found: {train_dir}")
|
|
|
|
print(f"[DEBUG] Final training shapefile path: {training_shapefile}")
|
|
print(f"[DEBUG] File exists: {os.path.exists(training_shapefile)}")
|
|
|
|
train_gdf = gpd.read_file(training_shapefile)
|
|
|
|
# Print initial shapefile info
|
|
update_status(f"📍 Loaded {len(train_gdf)} points from shapefile", 56)
|
|
print(f"[DEBUG] Shapefile CRS: {train_gdf.crs}")
|
|
print(f"[DEBUG] Shapefile bounds: {train_gdf.total_bounds}")
|
|
|
|
# Convert to WGS84 first (if not already) to match bbox coordinates
|
|
original_crs = train_gdf.crs
|
|
if train_gdf.crs and train_gdf.crs.to_epsg() != 4326:
|
|
print(f"📍 Converting training shapefile from {train_gdf.crs} to WGS84")
|
|
train_gdf = train_gdf.to_crs("EPSG:4326")
|
|
print(f"[DEBUG] WGS84 bounds: {train_gdf.total_bounds}")
|
|
|
|
# Check bbox overlap in WGS84
|
|
shp_bounds = train_gdf.total_bounds # [minx, miny, maxx, maxy]
|
|
bbox_wgs84 = bbox # [min_lon, min_lat, max_lon, max_lat]
|
|
|
|
# Check if there's overlap
|
|
overlap_x = not (shp_bounds[2] < bbox_wgs84[0] or shp_bounds[0] > bbox_wgs84[2])
|
|
overlap_y = not (shp_bounds[3] < bbox_wgs84[1] or shp_bounds[1] > bbox_wgs84[3])
|
|
|
|
if not (overlap_x and overlap_y):
|
|
update_status(f"⚠️ WARNING: Shapefile and bbox may not overlap!", 57)
|
|
print(f"[WARNING] Shapefile bounds (WGS84): {shp_bounds}")
|
|
print(f"[WARNING] Requested bbox (WGS84): {bbox_wgs84}")
|
|
print(f"[WARNING] This may result in 0 training samples!")
|
|
else:
|
|
# Crop to bbox to see how many points are actually in the region
|
|
train_gdf_cropped = train_gdf.cx[bbox_wgs84[0]:bbox_wgs84[2], bbox_wgs84[1]:bbox_wgs84[3]]
|
|
update_status(f"📍 {len(train_gdf_cropped)} points within bbox", 57)
|
|
if len(train_gdf_cropped) == 0:
|
|
raise ValueError(
|
|
f"No training points found within bbox!\n"
|
|
f"Shapefile bounds: {shp_bounds}\n"
|
|
f"Requested bbox: {bbox_wgs84}\n"
|
|
f"Please adjust bbox to cover your training data."
|
|
)
|
|
|
|
# Then convert to UTM Zone 48N (EPSG:32648) for extraction
|
|
if train_gdf.crs.to_epsg() != 32648:
|
|
print(f"📍 Converting training shapefile from WGS84 to UTM Zone 48N (EPSG:32648)")
|
|
train_gdf = train_gdf.to_crs('EPSG:32648')
|
|
print(f"[DEBUG] UTM bounds: {train_gdf.total_bounds}")
|
|
|
|
# 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)
|
|
|
|
print(f"[DEBUG] Starting feature extraction...")
|
|
print(f"[DEBUG] Feature mode: {feature_mode}")
|
|
print(f"[DEBUG] Training GDF has {len(train_gdf)} points")
|
|
print(f"[DEBUG] Training GDF CRS: {train_gdf.crs}")
|
|
print(f"[DEBUG] Training GDF bounds (UTM): {train_gdf.total_bounds}")
|
|
print(f"[DEBUG] Label column: {label_column}")
|
|
|
|
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)
|
|
|
|
print(f"[DEBUG] NDVI shape: {ndvi_masked.shape}")
|
|
print(f"[DEBUG] NDVI range: [{float(ndvi_masked.min())}, {float(ndvi_masked.max())}]")
|
|
|
|
# Extract features at training points
|
|
features = []
|
|
labels = []
|
|
failed_extractions = 0
|
|
|
|
# Test first point to see what's happening
|
|
first_point = train_gdf.iloc[0]
|
|
print(f"[DEBUG] Testing first point:")
|
|
print(f" Coords: ({first_point.geometry.x}, {first_point.geometry.y})")
|
|
print(f" Label: {first_point[label_column]}")
|
|
|
|
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)]
|
|
|
|
# Debug first few points
|
|
if idx < 3:
|
|
print(f"[DEBUG] Point {idx}: coords=({x_coord:.2f}, {y_coord:.2f}), ndvi={ndvi_val:.3f}, vh={vh_val:.3f}, vv={vv_val:.3f}")
|
|
|
|
if not np.isnan(feature_vec).any():
|
|
features.append(feature_vec)
|
|
labels.append(label)
|
|
else:
|
|
failed_extractions += 1
|
|
if idx < 3:
|
|
print(f"[DEBUG] Point {idx} has NaN: {feature_vec}")
|
|
except Exception as e:
|
|
failed_extractions += 1
|
|
if idx < 3:
|
|
print(f"[DEBUG] Point {idx} extraction failed: {e}")
|
|
continue
|
|
|
|
if failed_extractions > 0:
|
|
update_status(f"⚠️ {failed_extractions}/{len(train_gdf)} points had NaN/missing data", 65)
|
|
|
|
features = np.array(features)
|
|
labels = np.array(labels)
|
|
|
|
elif feature_mode in ['odc', 'extended']:
|
|
# For odc/extended: Extract features for full raster first, then sample at points
|
|
update_status(f"Extracting {feature_mode} features from full raster...", 62)
|
|
|
|
# Apply cloud mask first
|
|
if 'SCL' in ds_s2:
|
|
scl_band = ds_s2['SCL']
|
|
cloud_mask = scl_band.isin([1, 3, 8, 9, 10])
|
|
for band in ds_s2.data_vars:
|
|
if band != 'SCL':
|
|
ds_s2[band] = ds_s2[band].where(~cloud_mask)
|
|
|
|
# Extract features using FeatureExtractor for entire raster
|
|
raster_features = extractor.extract(
|
|
s2_data=ds_s2,
|
|
vh_data=None, # ODC/extended don't use radar in aggregate
|
|
vv_data=None
|
|
)
|
|
|
|
print(f"[DEBUG] Extracted raster features: shape={raster_features.shape}")
|
|
print(f"[DEBUG] Feature range: [{raster_features.min()}, {raster_features.max()}]")
|
|
|
|
# Now sample at each training point
|
|
features = []
|
|
labels = []
|
|
failed_extractions = 0
|
|
|
|
# Get spatial dimensions
|
|
y_coords = ds_s2.y.values
|
|
x_coords = ds_s2.x.values
|
|
|
|
print(f"[DEBUG] S2 spatial grid: x=[{x_coords.min()}, {x_coords.max()}], y=[{y_coords.min()}, {y_coords.max()}]")
|
|
|
|
for idx, row in train_gdf.iterrows():
|
|
point = row.geometry
|
|
x_coord = point.x
|
|
y_coord = point.y
|
|
label = row[label_column]
|
|
|
|
try:
|
|
# Find nearest pixel indices
|
|
x_idx = np.argmin(np.abs(x_coords - x_coord))
|
|
y_idx = np.argmin(np.abs(y_coords - y_coord))
|
|
|
|
# Get features at this pixel
|
|
# raster_features shape: (n_pixels, n_features)
|
|
# Need to convert 2D (y, x) index to 1D pixel index
|
|
pixel_idx = y_idx * len(x_coords) + x_idx
|
|
|
|
if pixel_idx < len(raster_features):
|
|
feature_vec = raster_features[pixel_idx]
|
|
|
|
if idx < 3:
|
|
print(f"[DEBUG] Point {idx}: coords=({x_coord:.2f}, {y_coord:.2f}) -> pixel[{y_idx},{x_idx}] -> idx={pixel_idx}, features={feature_vec[:3]}...")
|
|
|
|
if not np.isnan(feature_vec).any():
|
|
features.append(feature_vec)
|
|
labels.append(label)
|
|
else:
|
|
failed_extractions += 1
|
|
if idx < 3:
|
|
print(f"[DEBUG] Point {idx} has NaN features")
|
|
else:
|
|
failed_extractions += 1
|
|
if idx < 3:
|
|
print(f"[DEBUG] Point {idx} pixel_idx {pixel_idx} out of range (max={len(raster_features)})")
|
|
except Exception as e:
|
|
failed_extractions += 1
|
|
if idx < 3:
|
|
print(f"[DEBUG] Point {idx} extraction failed: {e}")
|
|
continue
|
|
|
|
if failed_extractions > 0:
|
|
update_status(f"⚠️ {failed_extractions}/{len(train_gdf)} points had NaN/missing data", 65)
|
|
|
|
features = np.array(features)
|
|
labels = np.array(labels)
|
|
|
|
else: # temporal 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 = []
|
|
failed_extractions = 0
|
|
|
|
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)
|
|
else:
|
|
failed_extractions += 1
|
|
except Exception as e:
|
|
failed_extractions += 1
|
|
continue
|
|
|
|
if failed_extractions > 0:
|
|
update_status(f"⚠️ {failed_extractions}/{len(train_gdf)} points had NaN/missing data", 65)
|
|
|
|
features = np.array(features)
|
|
labels = np.array(labels)
|
|
|
|
check_cancellation()
|
|
|
|
update_status(f"Extracted {len(features)} valid training samples", 70)
|
|
|
|
# ============ VALIDATE SAMPLES ============
|
|
if len(features) == 0:
|
|
error_msg = (
|
|
f"❌ No valid training samples extracted!\n"
|
|
f"Possible reasons:\n"
|
|
f"1. Training shapefile points don't overlap with bbox: {bbox}\n"
|
|
f"2. All points have NaN values (cloud cover, missing data)\n"
|
|
f"3. Coordinate system mismatch\n"
|
|
f"Suggestions:\n"
|
|
f"- Check if bbox matches your region\n"
|
|
f"- Try a different time range with less cloud cover\n"
|
|
f"- Verify training shapefile coordinates are correct"
|
|
)
|
|
raise ValueError(error_msg)
|
|
|
|
# Warn if very few samples
|
|
if len(features) < 20:
|
|
update_status(f"⚠️ Warning: Only {len(features)} samples extracted. Results may be unreliable.", 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)
|
|
|
|
# Validate samples after cache loading
|
|
if len(features) == 0:
|
|
error_msg = (
|
|
f"❌ No training samples available!\n"
|
|
f"The cached or loaded dataset is empty.\n"
|
|
f"Please try:\n"
|
|
f"1. Clear cache and reload data\n"
|
|
f"2. Check training shapefile and bbox overlap\n"
|
|
f"3. Adjust time range and cloud cover settings"
|
|
)
|
|
raise ValueError(error_msg)
|
|
|
|
# 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':
|
|
if use_gpu:
|
|
model = XGBClassifier(
|
|
n_estimators=n_estimators,
|
|
max_depth=max_depth,
|
|
tree_method='hist',
|
|
device='cuda:0',
|
|
random_state=42,
|
|
n_jobs=-1,
|
|
verbosity=0
|
|
)
|
|
else:
|
|
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 == 'lightgbm':
|
|
model = LGBMClassifier(
|
|
n_estimators=n_estimators if n_estimators else 300,
|
|
max_depth=max_depth if max_depth else -1,
|
|
learning_rate=learning_rate,
|
|
class_weight='balanced',
|
|
random_state=42,
|
|
device='gpu' if use_gpu else 'cpu'
|
|
)
|
|
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)
|
|
|
|
elif model_type == 'swin-unet':
|
|
if not PYTORCH_AVAILABLE:
|
|
raise ImportError("PyTorch is required for Swin-UNet. Install: pip install torch torchvision")
|
|
|
|
n_features = X_train.shape[1]
|
|
n_classes = len(np.unique(y_train))
|
|
|
|
device = torch.device('cuda' if torch.cuda.is_available() and use_gpu else 'cpu')
|
|
update_status(f"Building Swin-UNet model on {device}...", 75)
|
|
|
|
model = SwinUNetClassifier(n_features, n_classes, embed_dim=128).to(device)
|
|
|
|
# Convert to PyTorch tensors (no unsqueeze needed for Swin-UNet)
|
|
X_train_tensor = torch.FloatTensor(X_train)
|
|
y_train_tensor = torch.LongTensor(y_train)
|
|
X_test_tensor = torch.FloatTensor(X_test)
|
|
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)
|
|
|
|
# Calculate class weights for imbalanced data
|
|
class_counts = np.bincount(y_train)
|
|
class_weights = 1.0 / (class_counts + 1e-6) # Avoid division by zero
|
|
class_weights = class_weights / class_weights.sum() * len(class_counts) # Normalize
|
|
class_weights_tensor = torch.FloatTensor(class_weights).to(device)
|
|
|
|
print(f"[SWIN-UNET] Class distribution: {class_counts}")
|
|
print(f"[SWIN-UNET] Class weights: {class_weights}")
|
|
|
|
# Loss with class weights and optimizer with weight decay
|
|
criterion = nn.CrossEntropyLoss(weight=class_weights_tensor)
|
|
optimizer = optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=0.01)
|
|
|
|
# LR scheduler for better convergence
|
|
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=50)
|
|
|
|
# Early stopping to prevent overfitting
|
|
best_val_loss = float('inf')
|
|
patience = 10
|
|
patience_counter = 0
|
|
|
|
# Train Swin-UNet
|
|
update_status("Training Swin-UNet model with PyTorch (with class weights)...", 80)
|
|
epochs = min(60, n_estimators // 2) # Swin-UNet benefits from more epochs
|
|
|
|
# Validation dataset
|
|
val_dataset = TensorDataset(X_test_tensor, y_test_tensor)
|
|
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)
|
|
|
|
model.train()
|
|
for epoch in range(epochs):
|
|
# Training phase
|
|
model.train()
|
|
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()
|
|
|
|
# Gradient clipping to prevent exploding gradients
|
|
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
|
|
|
|
optimizer.step()
|
|
|
|
epoch_loss += loss.item()
|
|
|
|
scheduler.step()
|
|
|
|
# Validation phase
|
|
model.eval()
|
|
val_loss = 0.0
|
|
correct = 0
|
|
total = 0
|
|
with torch.no_grad():
|
|
for batch_X, batch_y in val_loader:
|
|
batch_X, batch_y = batch_X.to(device), batch_y.to(device)
|
|
outputs = model(batch_X)
|
|
loss = criterion(outputs, batch_y)
|
|
val_loss += loss.item()
|
|
|
|
_, predicted = torch.max(outputs, 1)
|
|
total += batch_y.size(0)
|
|
correct += (predicted == batch_y).sum().item()
|
|
|
|
avg_train_loss = epoch_loss / len(train_loader)
|
|
avg_val_loss = val_loss / len(val_loader)
|
|
val_acc = 100 * correct / total
|
|
lr = optimizer.param_groups[0]['lr']
|
|
|
|
if (epoch + 1) % 5 == 0:
|
|
update_status(f"Swin-UNet Epoch {epoch+1}/{epochs}, Train Loss: {avg_train_loss:.4f}, Val Loss: {avg_val_loss:.4f}, Val Acc: {val_acc:.2f}%, LR: {lr:.6f}", 80 + (epoch / epochs) * 10)
|
|
print(f"[SWIN-UNET] Epoch {epoch+1}/{epochs} - Train Loss: {avg_train_loss:.4f}, Val Loss: {avg_val_loss:.4f}, Val Acc: {val_acc:.2f}%")
|
|
|
|
# Early stopping check
|
|
if avg_val_loss < best_val_loss:
|
|
best_val_loss = avg_val_loss
|
|
patience_counter = 0
|
|
else:
|
|
patience_counter += 1
|
|
if patience_counter >= patience:
|
|
print(f"[SWIN-UNET] Early stopping at epoch {epoch+1} (best val loss: {best_val_loss:.4f})")
|
|
update_status(f"Swin-UNet early stopped at epoch {epoch+1}", 90)
|
|
break
|
|
|
|
model = model.cpu()
|
|
model.device_used = str(device)
|
|
|
|
elif model_type == 'mobilenet-lraspp':
|
|
if not PYTORCH_AVAILABLE:
|
|
raise ImportError("PyTorch is required for MobileNetV3 + LR-ASPP. Install: pip install torch torchvision")
|
|
|
|
n_features = X_train.shape[1]
|
|
n_classes = len(np.unique(y_train))
|
|
|
|
device = torch.device('cuda' if torch.cuda.is_available() and use_gpu else 'cpu')
|
|
update_status(f"Building MobileNetV3 + LR-ASPP model on {device}...", 75)
|
|
|
|
model = MobileNetLRASPPClassifier(n_features, n_classes).to(device)
|
|
|
|
# Convert to PyTorch tensors
|
|
X_train_tensor = torch.FloatTensor(X_train)
|
|
y_train_tensor = torch.LongTensor(y_train)
|
|
X_test_tensor = torch.FloatTensor(X_test)
|
|
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=64, shuffle=True) # Larger batch for efficiency
|
|
|
|
# Calculate class weights for imbalanced data
|
|
class_counts = np.bincount(y_train)
|
|
class_weights = 1.0 / (class_counts + 1e-6)
|
|
class_weights = class_weights / class_weights.sum() * len(class_counts)
|
|
class_weights_tensor = torch.FloatTensor(class_weights).to(device)
|
|
|
|
print(f"[MOBILENET] Class distribution: {class_counts}")
|
|
print(f"[MOBILENET] Class weights: {class_weights}")
|
|
|
|
# Loss with class weights
|
|
criterion = nn.CrossEntropyLoss(weight=class_weights_tensor)
|
|
optimizer = optim.Adam(model.parameters(), lr=learning_rate, weight_decay=0.0001)
|
|
|
|
# LR scheduler
|
|
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=5)
|
|
|
|
# Early stopping
|
|
best_val_loss = float('inf')
|
|
patience = 10
|
|
patience_counter = 0
|
|
|
|
# Train MobileNetV3 + LR-ASPP
|
|
update_status("Training MobileNetV3 + LR-ASPP model with PyTorch...", 80)
|
|
epochs = min(60, n_estimators // 2)
|
|
|
|
# Validation dataset
|
|
val_dataset = TensorDataset(X_test_tensor, y_test_tensor)
|
|
val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False)
|
|
|
|
model.train()
|
|
for epoch in range(epochs):
|
|
# Training phase
|
|
model.train()
|
|
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()
|
|
|
|
# Gradient clipping
|
|
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
|
|
|
|
optimizer.step()
|
|
epoch_loss += loss.item()
|
|
|
|
# Validation phase
|
|
model.eval()
|
|
val_loss = 0.0
|
|
correct = 0
|
|
total = 0
|
|
with torch.no_grad():
|
|
for batch_X, batch_y in val_loader:
|
|
batch_X, batch_y = batch_X.to(device), batch_y.to(device)
|
|
outputs = model(batch_X)
|
|
loss = criterion(outputs, batch_y)
|
|
val_loss += loss.item()
|
|
|
|
_, predicted = torch.max(outputs, 1)
|
|
total += batch_y.size(0)
|
|
correct += (predicted == batch_y).sum().item()
|
|
|
|
avg_train_loss = epoch_loss / len(train_loader)
|
|
avg_val_loss = val_loss / len(val_loader)
|
|
val_acc = 100 * correct / total
|
|
|
|
# Update learning rate
|
|
scheduler.step(avg_val_loss)
|
|
lr = optimizer.param_groups[0]['lr']
|
|
|
|
if (epoch + 1) % 5 == 0:
|
|
update_status(f"MobileNet Epoch {epoch+1}/{epochs}, Train Loss: {avg_train_loss:.4f}, Val Loss: {avg_val_loss:.4f}, Val Acc: {val_acc:.2f}%, LR: {lr:.6f}", 80 + (epoch / epochs) * 10)
|
|
print(f"[MOBILENET] Epoch {epoch+1}/{epochs} - Train Loss: {avg_train_loss:.4f}, Val Loss: {avg_val_loss:.4f}, Val Acc: {val_acc:.2f}%")
|
|
|
|
# Early stopping
|
|
if avg_val_loss < best_val_loss:
|
|
best_val_loss = avg_val_loss
|
|
patience_counter = 0
|
|
else:
|
|
patience_counter += 1
|
|
if patience_counter >= patience:
|
|
print(f"[MOBILENET] Early stopping at epoch {epoch+1} (best val loss: {best_val_loss:.4f})")
|
|
update_status(f"MobileNet early stopped at epoch {epoch+1}", 90)
|
|
break
|
|
|
|
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, swin-unet, mobilenet-lraspp")
|
|
|
|
# Fit non-neural-network models
|
|
if model_type not in ['cnn', 'swin-unet', 'mobilenet-lraspp']:
|
|
model.fit(X_train, y_train)
|
|
|
|
# Evaluate
|
|
update_status("Evaluating model...", 90)
|
|
if model_type in ['cnn', 'swin-unet', 'mobilenet-lraspp']:
|
|
# PyTorch models 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, accuracy_score
|
|
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', 'lightgbm', 'cnn', 'swin-unet', 'mobilenet-lraspp'] else None,
|
|
"max_depth": max_depth if model_type not in ['cnn', 'swin-unet', 'mobilenet-lraspp'] else None,
|
|
"learning_rate": learning_rate if model_type in ['xgboost', 'lightgbm', 'swin-unet', 'mobilenet-lraspp'] else None,
|
|
"epochs": min(50, n_estimators // 2) if model_type == 'cnn' else (min(60, n_estimators // 2) if model_type in ['swin-unet', 'mobilenet-lraspp'] 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
|
|
)
|
|
|
|
# Construct info path (model manager saves it in model_train/)
|
|
info_path = os.path.join('model_train', model_filename.replace('.joblib', '_info.json'))
|
|
|
|
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)
|
|
}
|