basketballcantho/remote-sensing
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

hoàn thành model swing-unet

Browse Source
This commit is contained in:
Victor Phan
2026-01-05 16:20:58 +07:00
parent f401765996
commit ae3e5fffdd
4 changed files with 509 additions and 95 deletions
Download Patch File Download Diff File Expand all files Collapse all files
train_module.py
+150 -12
View File
@@ -277,7 +277,7 @@ def train_model(
use_gpu=True,
use_cache=True,
test_size=0.2,
feature_mode='simple', # Changed from 'odc' - simple mode works with B04, B08, SCL only
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
@@ -406,9 +406,11 @@ def train_model(
# Load different bands based on feature mode
if feature_mode == 'simple':
bands_to_load = ["B04", "B08", "SCL"]
else: # temporal or extended
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,
@@ -428,9 +430,11 @@ def train_model(
if 'time' in ds_s2.dims:
print(f"[DEBUG S2] Time range: {ds_s2.time.min().values} to {ds_s2.time.max().values}")
# Rename for compatibility (simple mode)
if "B04" in ds_s2 and "red" not in ds_s2:
# 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()
@@ -575,6 +579,7 @@ def train_model(
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}")
@@ -637,7 +642,85 @@ def train_model(
features = np.array(features)
labels = np.array(labels)
else: # temporal or extended mode
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']
@@ -882,19 +965,39 @@ def train_model(
train_dataset = TensorDataset(X_train_tensor, y_train_tensor)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
# Loss and optimizer with weight decay
criterion = nn.CrossEntropyLoss()
# 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...", 80)
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)
@@ -903,16 +1006,51 @@ def train_model(
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()
if (epoch + 1) % 10 == 0:
avg_loss = epoch_loss / len(train_loader)
lr = optimizer.param_groups[0]['lr']
update_status(f"Swin-UNet Epoch {epoch+1}/{epochs}, Loss: {avg_loss:.4f}, LR: {lr:.6f}", 80 + (epoch / epochs) * 10)
# 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)