hoàn thành chức năng remove cloud train
This commit is contained in:
Victor Phan
@@ -0,0 +1,512 @@
|
||||
"""
|
||||
Train Cloud Removal Model using SEN12MS-CR Dataset
|
||||
Huấn luyện model Deep Learning để khử mây từ ảnh Sentinel-2
|
||||
|
||||
Dataset: SEN12MS-CR (Sentinel-12 Multi-Seasonal Cloud Removal)
|
||||
- Input: S2 cloudy images (ảnh Sentinel-2 bị mây)
|
||||
- Target: S2 clean images (ảnh Sentinel-2 sạch)
|
||||
- Optional: S1 SAR data (radar data không bị ảnh hưởng bởi mây)
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.optim as optim
|
||||
from torch.utils.data import Dataset, DataLoader
|
||||
from pathlib import Path
|
||||
import matplotlib.pyplot as plt
|
||||
from tqdm import tqdm
|
||||
|
||||
# Add winter_dataset to path
|
||||
sys.path.insert(0, str(Path(__file__).parent / "winter_dataset"))
|
||||
from sen12ms_cr_dataLoader import SEN12MSCRDataset, Seasons, S1Bands, S2Bands
|
||||
|
||||
|
||||
# ============ DATASET WRAPPER ============
|
||||
|
||||
class CloudRemovalDataset(Dataset):
|
||||
"""
|
||||
PyTorch Dataset wrapper cho SEN12MS-CR
|
||||
Input: S2 cloudy + S1 (optional)
|
||||
Target: S2 clean
|
||||
"""
|
||||
|
||||
def __init__(self, base_dir, season=Seasons.WINTER, use_s1=True,
|
||||
s2_bands=S2Bands.ALL, normalize=True):
|
||||
"""
|
||||
Args:
|
||||
base_dir: Đường dẫn đến thư mục chứa dữ liệu
|
||||
season: Mùa (SPRING, SUMMER, FALL, WINTER)
|
||||
use_s1: Có sử dụng dữ liệu S1 (radar) không
|
||||
s2_bands: Các band S2 cần dùng
|
||||
normalize: Normalize dữ liệu về [0, 1]
|
||||
"""
|
||||
self.dataset = SEN12MSCRDataset(base_dir)
|
||||
self.season = season
|
||||
self.use_s1 = use_s1
|
||||
self.s2_bands = s2_bands
|
||||
self.normalize = normalize
|
||||
|
||||
# Lấy tất cả scene và patch IDs
|
||||
season_ids = self.dataset.get_season_ids(season)
|
||||
|
||||
# Tạo list of (scene_id, patch_id) pairs
|
||||
self.samples = []
|
||||
for scene_id, patch_ids in season_ids.items():
|
||||
for patch_id in patch_ids:
|
||||
self.samples.append((scene_id, patch_id))
|
||||
|
||||
# Get band count
|
||||
n_s2_bands = len(s2_bands.value) if hasattr(s2_bands, 'value') else len(s2_bands)
|
||||
|
||||
print(f"[DATASET] Loaded {len(self.samples)} samples from {season.value}")
|
||||
print(f"[DATASET] Use S1: {use_s1}, S2 bands: {n_s2_bands}")
|
||||
|
||||
def __len__(self):
|
||||
return len(self.samples)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
scene_id, patch_id = self.samples[idx]
|
||||
|
||||
# Load triplet: S1, S2 clean, S2 cloudy
|
||||
s1, s2_clean, s2_cloudy, bounds = self.dataset.get_s1s2s2cloudy_triplet(
|
||||
self.season,
|
||||
scene_id,
|
||||
patch_id,
|
||||
s1_bands=S1Bands.ALL if self.use_s1 else S1Bands.NONE,
|
||||
s2_bands=self.s2_bands,
|
||||
s2cloudy_bands=self.s2_bands
|
||||
)
|
||||
|
||||
# Normalize to [0, 1] if needed
|
||||
if self.normalize:
|
||||
s2_clean = s2_clean.astype(np.float32) / 10000.0 # S2 values are in [0, 10000]
|
||||
s2_cloudy = s2_cloudy.astype(np.float32) / 10000.0
|
||||
if self.use_s1:
|
||||
# S1 values need different normalization (dB scale)
|
||||
s1 = (s1.astype(np.float32) + 30) / 50.0 # Normalize from [-30, 20] to [0, 1]
|
||||
s1 = np.clip(s1, 0, 1)
|
||||
|
||||
# Convert to torch tensors
|
||||
s2_clean = torch.from_numpy(s2_clean).float()
|
||||
s2_cloudy = torch.from_numpy(s2_cloudy).float()
|
||||
|
||||
# Input: S2 cloudy + S1 (if enabled)
|
||||
if self.use_s1:
|
||||
s1 = torch.from_numpy(s1).float()
|
||||
input_data = torch.cat([s2_cloudy, s1], dim=0)
|
||||
else:
|
||||
input_data = s2_cloudy
|
||||
|
||||
return input_data, s2_clean
|
||||
|
||||
|
||||
# ============ U-NET ARCHITECTURE ============
|
||||
|
||||
class DoubleConv(nn.Module):
|
||||
"""(Conv2d -> BatchNorm -> ReLU) x 2"""
|
||||
|
||||
def __init__(self, in_channels, out_channels):
|
||||
super().__init__()
|
||||
self.double_conv = nn.Sequential(
|
||||
nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
|
||||
nn.BatchNorm2d(out_channels),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
|
||||
nn.BatchNorm2d(out_channels),
|
||||
nn.ReLU(inplace=True)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return self.double_conv(x)
|
||||
|
||||
|
||||
class UNet(nn.Module):
|
||||
"""
|
||||
U-Net architecture cho cloud removal
|
||||
Input: S2 cloudy (+ S1 optional) [B, C_in, H, W]
|
||||
Output: S2 clean [B, C_out, H, W]
|
||||
"""
|
||||
|
||||
def __init__(self, in_channels, out_channels, features=[64, 128, 256, 512]):
|
||||
super().__init__()
|
||||
self.encoder = nn.ModuleList()
|
||||
self.decoder = nn.ModuleList()
|
||||
self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
|
||||
|
||||
# Encoder (downsampling)
|
||||
for feature in features:
|
||||
self.encoder.append(DoubleConv(in_channels, feature))
|
||||
in_channels = feature
|
||||
|
||||
# Bottleneck
|
||||
self.bottleneck = DoubleConv(features[-1], features[-1] * 2)
|
||||
|
||||
# Decoder (upsampling)
|
||||
for feature in reversed(features):
|
||||
self.decoder.append(
|
||||
nn.ConvTranspose2d(feature * 2, feature, kernel_size=2, stride=2)
|
||||
)
|
||||
self.decoder.append(DoubleConv(feature * 2, feature))
|
||||
|
||||
# Final output layer
|
||||
self.final_conv = nn.Conv2d(features[0], out_channels, kernel_size=1)
|
||||
|
||||
def forward(self, x):
|
||||
skip_connections = []
|
||||
|
||||
# Encoder
|
||||
for encode in self.encoder:
|
||||
x = encode(x)
|
||||
skip_connections.append(x)
|
||||
x = self.pool(x)
|
||||
|
||||
# Bottleneck
|
||||
x = self.bottleneck(x)
|
||||
|
||||
# Decoder
|
||||
skip_connections = skip_connections[::-1]
|
||||
|
||||
for idx in range(0, len(self.decoder), 2):
|
||||
x = self.decoder[idx](x) # Upsample
|
||||
skip_connection = skip_connections[idx // 2]
|
||||
|
||||
# Handle size mismatch
|
||||
if x.shape != skip_connection.shape:
|
||||
x = nn.functional.interpolate(x, size=skip_connection.shape[2:])
|
||||
|
||||
concat_skip = torch.cat((skip_connection, x), dim=1)
|
||||
x = self.decoder[idx + 1](concat_skip) # Double conv
|
||||
|
||||
return self.final_conv(x)
|
||||
|
||||
|
||||
# ============ TRAINING FUNCTIONS ============
|
||||
|
||||
def train_epoch(model, dataloader, criterion, optimizer, device):
|
||||
"""Train for one epoch"""
|
||||
model.train()
|
||||
total_loss = 0
|
||||
|
||||
pbar = tqdm(dataloader, desc="Training")
|
||||
for batch_idx, (inputs, targets) in enumerate(pbar):
|
||||
inputs = inputs.to(device)
|
||||
targets = targets.to(device)
|
||||
|
||||
# Forward pass
|
||||
optimizer.zero_grad()
|
||||
outputs = model(inputs)
|
||||
loss = criterion(outputs, targets)
|
||||
|
||||
# Backward pass
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
total_loss += loss.item()
|
||||
pbar.set_postfix({'loss': loss.item()})
|
||||
|
||||
return total_loss / len(dataloader)
|
||||
|
||||
|
||||
def validate(model, dataloader, criterion, device):
|
||||
"""Validate model"""
|
||||
model.eval()
|
||||
total_loss = 0
|
||||
|
||||
with torch.no_grad():
|
||||
for inputs, targets in tqdm(dataloader, desc="Validation"):
|
||||
inputs = inputs.to(device)
|
||||
targets = targets.to(device)
|
||||
|
||||
outputs = model(inputs)
|
||||
loss = criterion(outputs, targets)
|
||||
total_loss += loss.item()
|
||||
|
||||
return total_loss / len(dataloader)
|
||||
|
||||
|
||||
def visualize_results(model, dataset, device, num_samples=3):
|
||||
"""Visualize cloud removal results"""
|
||||
model.eval()
|
||||
|
||||
fig, axes = plt.subplots(num_samples, 3, figsize=(15, 5 * num_samples))
|
||||
|
||||
with torch.no_grad():
|
||||
for i in range(num_samples):
|
||||
idx = np.random.randint(0, len(dataset))
|
||||
input_data, target = dataset[idx]
|
||||
|
||||
input_data = input_data.unsqueeze(0).to(device)
|
||||
output = model(input_data)
|
||||
|
||||
# Convert to numpy
|
||||
input_rgb = input_data[0, :3, :, :].cpu().numpy().transpose(1, 2, 0)
|
||||
target_rgb = target[:3, :, :].cpu().numpy().transpose(1, 2, 0)
|
||||
output_rgb = output[0, :3, :, :].cpu().numpy().transpose(1, 2, 0)
|
||||
|
||||
# Clip to [0, 1]
|
||||
input_rgb = np.clip(input_rgb * 3, 0, 1) # Enhance for visualization
|
||||
target_rgb = np.clip(target_rgb * 3, 0, 1)
|
||||
output_rgb = np.clip(output_rgb * 3, 0, 1)
|
||||
|
||||
if num_samples == 1:
|
||||
axes[0].imshow(input_rgb)
|
||||
axes[0].set_title("Input (Cloudy)")
|
||||
axes[0].axis('off')
|
||||
|
||||
axes[1].imshow(output_rgb)
|
||||
axes[1].set_title("Output (Predicted)")
|
||||
axes[1].axis('off')
|
||||
|
||||
axes[2].imshow(target_rgb)
|
||||
axes[2].set_title("Target (Clean)")
|
||||
axes[2].axis('off')
|
||||
else:
|
||||
axes[i, 0].imshow(input_rgb)
|
||||
axes[i, 0].set_title(f"Sample {i+1}: Input (Cloudy)")
|
||||
axes[i, 0].axis('off')
|
||||
|
||||
axes[i, 1].imshow(output_rgb)
|
||||
axes[i, 1].set_title(f"Sample {i+1}: Output (Predicted)")
|
||||
axes[i, 1].axis('off')
|
||||
|
||||
axes[i, 2].imshow(target_rgb)
|
||||
axes[i, 2].set_title(f"Sample {i+1}: Target (Clean)")
|
||||
axes[i, 2].axis('off')
|
||||
|
||||
plt.tight_layout()
|
||||
return fig
|
||||
|
||||
|
||||
# ============ MAIN TRAINING SCRIPT ============
|
||||
|
||||
def train_cloud_removal_model(
|
||||
data_dir="winter_dataset",
|
||||
use_s1=True,
|
||||
batch_size=8,
|
||||
num_epochs=50,
|
||||
learning_rate=1e-4,
|
||||
device="cuda" if torch.cuda.is_available() else "cpu",
|
||||
save_dir="model_train"
|
||||
):
|
||||
"""
|
||||
Train cloud removal model
|
||||
|
||||
Args:
|
||||
data_dir: Thư mục chứa dữ liệu SEN12MS-CR
|
||||
use_s1: Có sử dụng S1 radar data không
|
||||
batch_size: Batch size
|
||||
num_epochs: Số epochs
|
||||
learning_rate: Learning rate
|
||||
device: 'cuda' hoặc 'cpu'
|
||||
save_dir: Thư mục lưu model
|
||||
"""
|
||||
|
||||
print("=" * 70)
|
||||
print("🌥️ CLOUD REMOVAL MODEL TRAINING")
|
||||
print("=" * 70)
|
||||
print(f"Data directory: {data_dir}")
|
||||
print(f"Use S1 (SAR): {use_s1}")
|
||||
print(f"Device: {device}")
|
||||
print(f"Batch size: {batch_size}")
|
||||
print(f"Epochs: {num_epochs}")
|
||||
print(f"Learning rate: {learning_rate}")
|
||||
print("=" * 70)
|
||||
|
||||
# Create dataset
|
||||
print("\n📂 Loading dataset...")
|
||||
|
||||
# Use RGB + NIR bands for training (B02, B03, B04, B08)
|
||||
s2_bands = [S2Bands.B02, S2Bands.B03, S2Bands.B04, S2Bands.B08]
|
||||
|
||||
dataset = CloudRemovalDataset(
|
||||
base_dir=data_dir,
|
||||
season=Seasons.WINTER,
|
||||
use_s1=use_s1,
|
||||
s2_bands=s2_bands,
|
||||
normalize=True
|
||||
)
|
||||
|
||||
# Split train/val
|
||||
train_size = int(0.8 * len(dataset))
|
||||
val_size = len(dataset) - train_size
|
||||
train_dataset, val_dataset = torch.utils.data.random_split(
|
||||
dataset, [train_size, val_size]
|
||||
)
|
||||
|
||||
print(f"Train samples: {len(train_dataset)}")
|
||||
print(f"Val samples: {len(val_dataset)}")
|
||||
|
||||
# Create dataloaders
|
||||
train_loader = DataLoader(
|
||||
train_dataset,
|
||||
batch_size=batch_size,
|
||||
shuffle=True,
|
||||
num_workers=4,
|
||||
pin_memory=True if device == "cuda" else False
|
||||
)
|
||||
|
||||
val_loader = DataLoader(
|
||||
val_dataset,
|
||||
batch_size=batch_size,
|
||||
shuffle=False,
|
||||
num_workers=4,
|
||||
pin_memory=True if device == "cuda" else False
|
||||
)
|
||||
|
||||
# Create model
|
||||
print("\n🏗️ Creating U-Net model...")
|
||||
in_channels = len(s2_bands) + (2 if use_s1 else 0) # S2 + S1 (VV, VH)
|
||||
out_channels = len(s2_bands)
|
||||
|
||||
model = UNet(in_channels=in_channels, out_channels=out_channels)
|
||||
model = model.to(device)
|
||||
|
||||
print(f"Input channels: {in_channels}")
|
||||
print(f"Output channels: {out_channels}")
|
||||
print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}")
|
||||
|
||||
# Loss and optimizer
|
||||
criterion = nn.L1Loss() # MAE loss
|
||||
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
|
||||
scheduler = optim.lr_scheduler.ReduceLROnPlateau(
|
||||
optimizer, mode='min', factor=0.5, patience=5
|
||||
)
|
||||
|
||||
# Training loop
|
||||
print("\n🚀 Starting training...")
|
||||
best_val_loss = float('inf')
|
||||
train_losses = []
|
||||
val_losses = []
|
||||
|
||||
for epoch in range(num_epochs):
|
||||
print(f"\n{'='*70}")
|
||||
print(f"Epoch {epoch + 1}/{num_epochs}")
|
||||
print(f"{'='*70}")
|
||||
|
||||
# Train
|
||||
train_loss = train_epoch(model, train_loader, criterion, optimizer, device)
|
||||
train_losses.append(train_loss)
|
||||
|
||||
# Validate
|
||||
val_loss = validate(model, val_loader, criterion, device)
|
||||
val_losses.append(val_loss)
|
||||
|
||||
# Update learning rate
|
||||
scheduler.step(val_loss)
|
||||
|
||||
print(f"\nEpoch {epoch + 1} Summary:")
|
||||
print(f" Train Loss: {train_loss:.6f}")
|
||||
print(f" Val Loss: {val_loss:.6f}")
|
||||
|
||||
# Save best model
|
||||
if val_loss < best_val_loss:
|
||||
best_val_loss = val_loss
|
||||
save_path = Path(save_dir) / "cloud_removal_unet_best.pth"
|
||||
save_path.parent.mkdir(exist_ok=True)
|
||||
|
||||
torch.save({
|
||||
'epoch': epoch,
|
||||
'model_state_dict': model.state_dict(),
|
||||
'optimizer_state_dict': optimizer.state_dict(),
|
||||
'train_loss': train_loss,
|
||||
'val_loss': val_loss,
|
||||
'use_s1': use_s1,
|
||||
'in_channels': in_channels,
|
||||
'out_channels': out_channels
|
||||
}, save_path)
|
||||
|
||||
print(f" 💾 Saved best model: {save_path}")
|
||||
|
||||
# Visualize every 10 epochs
|
||||
if (epoch + 1) % 10 == 0:
|
||||
print("\n📊 Generating visualizations...")
|
||||
fig = visualize_results(model, val_dataset, device, num_samples=3)
|
||||
|
||||
viz_path = Path(save_dir) / f"cloud_removal_epoch_{epoch+1}.png"
|
||||
fig.savefig(viz_path, dpi=150, bbox_inches='tight')
|
||||
plt.close(fig)
|
||||
|
||||
print(f" 💾 Saved visualization: {viz_path}")
|
||||
|
||||
# Plot training curves
|
||||
print("\n📈 Plotting training curves...")
|
||||
fig, ax = plt.subplots(figsize=(10, 6))
|
||||
ax.plot(train_losses, label='Train Loss')
|
||||
ax.plot(val_losses, label='Val Loss')
|
||||
ax.set_xlabel('Epoch')
|
||||
ax.set_ylabel('Loss (MAE)')
|
||||
ax.set_title('Cloud Removal Training Progress')
|
||||
ax.legend()
|
||||
ax.grid(True)
|
||||
|
||||
curve_path = Path(save_dir) / "training_curves.png"
|
||||
fig.savefig(curve_path, dpi=150, bbox_inches='tight')
|
||||
plt.close(fig)
|
||||
|
||||
print(f" 💾 Saved training curves: {curve_path}")
|
||||
|
||||
# Final summary
|
||||
print("\n" + "=" * 70)
|
||||
print("✅ TRAINING COMPLETED!")
|
||||
print("=" * 70)
|
||||
print(f"Best validation loss: {best_val_loss:.6f}")
|
||||
print(f"Model saved to: {Path(save_dir) / 'cloud_removal_unet_best.pth'}")
|
||||
print("=" * 70)
|
||||
|
||||
return model, train_losses, val_losses
|
||||
|
||||
|
||||
# ============ INFERENCE FUNCTION ============
|
||||
|
||||
def apply_cloud_removal(model_path, cloudy_image, s1_data=None, device="cuda"):
|
||||
"""
|
||||
Áp dụng model để khử mây cho một ảnh
|
||||
|
||||
Args:
|
||||
model_path: Đường dẫn đến model đã train
|
||||
cloudy_image: Ảnh S2 bị mây [C, H, W]
|
||||
s1_data: Dữ liệu S1 (optional) [2, H, W]
|
||||
device: 'cuda' hoặc 'cpu'
|
||||
|
||||
Returns:
|
||||
cleaned_image: Ảnh đã khử mây [C, H, W]
|
||||
"""
|
||||
# Load model
|
||||
checkpoint = torch.load(model_path, map_location=device)
|
||||
|
||||
model = UNet(
|
||||
in_channels=checkpoint['in_channels'],
|
||||
out_channels=checkpoint['out_channels']
|
||||
)
|
||||
model.load_state_dict(checkpoint['model_state_dict'])
|
||||
model = model.to(device)
|
||||
model.eval()
|
||||
|
||||
# Prepare input
|
||||
input_tensor = torch.from_numpy(cloudy_image).float().unsqueeze(0).to(device)
|
||||
|
||||
if checkpoint['use_s1'] and s1_data is not None:
|
||||
s1_tensor = torch.from_numpy(s1_data).float().unsqueeze(0).to(device)
|
||||
input_tensor = torch.cat([input_tensor, s1_tensor], dim=1)
|
||||
|
||||
# Inference
|
||||
with torch.no_grad():
|
||||
output = model(input_tensor)
|
||||
|
||||
cleaned_image = output[0].cpu().numpy()
|
||||
|
||||
return cleaned_image
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Train model
|
||||
model, train_losses, val_losses = train_cloud_removal_model(
|
||||
data_dir="winter_dataset",
|
||||
use_s1=True,
|
||||
batch_size=8,
|
||||
num_epochs=50,
|
||||
learning_rate=1e-4
|
||||
)
|
||||
Reference in New Issue
Block a user