In [ ]:
#CELL_1
import ee
import geemap
import datetime
import os
import numpy as np
import rasterio
import time
from sklearn.model_selection import train_test_split
# --- 1. INITIALIZATION ---
try:
ee.Initialize(project='[REDACTED_FOR_SECURITY]')
except:
ee.Authenticate()
ee.Initialize(project='[REDACTED_FOR_SECURITY]')
# Configuration
year = 2021
START_DATE = f'{year}-10-16'
END_DATE = f'{year + 1}-04-16'
ASSET_ID = '[REDACTED_FOR_SECURITY]'
# NOTE: Ensure this folder exists in your actual Google Drive root
DRIVE_FOLDER = 'SatMAE_Scratch_Results_12Frames'
SAVE_DIR = f'/content/drive/MyDrive/{DRIVE_FOLDER}/'
if not os.path.exists(SAVE_DIR):
os.makedirs(SAVE_DIR)
# Assets
wheat_mask = ee.Image(ASSET_ID)
roi = wheat_mask.geometry()
s2Bands = ['B2', 'B3', 'B4', 'B5', 'B6', 'B7', 'B8', 'B8A', 'B11', 'B12']
MAX_CLOUD_PROB = 70
# Sentinel-2 Processing
s2Raw = (ee.ImageCollection('COPERNICUS/S2_SR_HARMONIZED')
.filterDate(START_DATE, END_DATE).filterBounds(roi)
.filter(ee.Filter.lt('CLOUDY_PIXEL_PERCENTAGE', 100))
.map(lambda img: img.updateMask(
img.select('MSK_SNWPRB').lt(1)
.add(img.select('SCL').gte(4).And(img.select('SCL').lte(6)))
)))
s2Clouds = ee.ImageCollection('COPERNICUS/S2_CLOUD_PROBABILITY').filterDate(START_DATE, END_DATE).filterBounds(roi)
s2Joined = ee.ImageCollection(ee.Join.saveFirst('cloud_mask_img').apply(
s2Raw, s2Clouds, ee.Filter.equals(leftField='system:index', rightField='system:index')))
def processS2(img):
isCloud = ee.Image(img.get('cloud_mask_img')).select('probability').gt(MAX_CLOUD_PROB)
optical = img.select(s2Bands).multiply(0.025)
ndvi = optical.normalizedDifference(['B8', 'B4']).rename('NDVI')
combined_mask = isCloud.Not()
return (optical.addBands(ndvi).updateMask(combined_mask).copyProperties(img, ['system:time_start']))
s2Clean = s2Joined.map(processS2)
# Fortnightly Composites
intervals = []
cur = datetime.datetime.fromisoformat(START_DATE)
end = datetime.datetime.fromisoformat(END_DATE)
while cur < end:
if cur.day == 1:
nxt = cur + datetime.timedelta(days=14)
str_lbl = f"{cur.year}_{cur.month:02d}_1"
else:
next_month = cur.replace(day=1) + datetime.timedelta(days=32)
nxt = next_month.replace(day=1)
str_lbl = f"{cur.year}_{cur.month:02d}_2"
if nxt > end: nxt = end
intervals.append([cur.strftime('%Y-%m-%d'), (nxt).strftime('%Y-%m-%d'), str_lbl])
cur = nxt
intervals = intervals[:12]
print(f"Generating {len(intervals)} Time Frames...")
def makeComposite(item):
start, end = ee.Date(ee.List(item).get(0)), ee.Date(ee.List(item).get(1))
return s2Clean.filterDate(start, end).qualityMosaic('NDVI').select(s2Bands)\
.clamp(0, 250).toByte().unmask(255)\
.set('fortnight_label', ee.List(item).get(2))
fortnightlyCol = ee.ImageCollection.fromImages(ee.List(intervals).map(makeComposite))
# --- ROBUST EXPORT LOGIC ---
def monitor_task(task):
print(f"Submitting Task: {task.config['description']}...")
task.start()
while True:
status = task.status()
state = status['state']
if state in ['COMPLETED', 'FAILED', 'CANCELLED']:
print(f"\nTask Finished with state: {state}")
if state == 'FAILED':
print(f"Error Message: {status.get('error_message', 'Unknown Error')}")
raise RuntimeError("GEE Export Failed.")
break
print(f"Status: {state}...", end='\r')
time.sleep(10)
def wait_for_file(filepath, timeout=120):
"""Waits for a file to appear in Colab Drive mount."""
print(f"Waiting for Drive sync: {os.path.basename(filepath)}...")
start = time.time()
while not os.path.exists(filepath):
if time.time() - start > timeout:
raise FileNotFoundError(f"Timeout: {filepath} did not appear after {timeout}s.")
time.sleep(5)
# Force a directory refresh
try: os.listdir(os.path.dirname(filepath))
except: pass
print(f"Found: {filepath}")
return True
def generate_local_dataset():
print("Starting Robust Data Generation (Batch Export)...")
centroid = roi.centroid()
buffer_poly = centroid.buffer(2500).bounds()
# 1. Export Stack
stack_name = 'local_stack'
stack_file = os.path.join(SAVE_DIR, f'{stack_name}.tif')
if not os.path.exists(stack_file):
task_stack = ee.batch.Export.image.toDrive(
image=fortnightlyCol.toBands(),
description='Export_Stack_12Frames',
folder=DRIVE_FOLDER,
fileNamePrefix=stack_name,
region=buffer_poly,
scale=10,
crs='EPSG:4326',
fileFormat='GeoTIFF',
maxPixels=1e9
)
monitor_task(task_stack)
else:
print(f"Found existing stack export: {stack_file}")
# 2. Export Mask
mask_name = 'local_mask'
mask_file = os.path.join(SAVE_DIR, f'{mask_name}.tif')
if not os.path.exists(mask_file):
task_mask = ee.batch.Export.image.toDrive(
image=wheat_mask,
description='Export_Mask',
folder=DRIVE_FOLDER,
fileNamePrefix=mask_name,
region=buffer_poly,
scale=10,
crs='EPSG:4326',
fileFormat='GeoTIFF',
maxPixels=1e9
)
monitor_task(task_mask)
else:
print(f"Found existing mask export: {mask_file}")
# 3. Sync Wait
wait_for_file(stack_file)
wait_for_file(mask_file)
# 4. Process
print("Processing GeoTIFFs to Numpy...")
with rasterio.open(stack_file) as src_stack, rasterio.open(mask_file) as src_mask:
img_data = src_stack.read()
mask_data = src_mask.read(1)
H, W = mask_data.shape
if img_data.shape[0] < 120:
print(f"Warning: Got {img_data.shape[0]} bands. Padding with zeros.")
padding = np.zeros((120 - img_data.shape[0], H, W), dtype=img_data.dtype)
img_data = np.concatenate([img_data, padding], axis=0)
img_reshaped = img_data[:120].reshape(12, 10, H, W)
X_list, y_list = [], []
if H < 224 or W < 224:
print("ROI too small.")
return
else:
for r in range(0, H-224, 112):
for c in range(0, W-224, 112):
chip_x = img_reshaped[:, :, r:r+224, c:c+224]
chip_y = mask_data[r:r+224, c:c+224]
if np.mean(chip_x == 255) > 0.30: continue
if np.mean(chip_y > 0) > 0.01:
X_list.append(chip_x)
y_list.append(chip_y)
if len(X_list) == 0: raise RuntimeError("No valid tiles found.")
X = np.array(X_list).astype(np.uint8)
y = np.array(y_list).astype(np.uint8)[:, None, :, :]
split_idx = int(0.8 * len(X))
X_train, X_val = X[:split_idx], X[split_idx:]
y_train, y_val = y[:split_idx], y[split_idx:]
print(f"Spatial Split: Train={len(X_train)}, Val={len(X_val)}")
np.save(os.path.join(SAVE_DIR, 'train_x.npy'), X_train)
np.save(os.path.join(SAVE_DIR, 'train_y.npy'), y_train)
np.save(os.path.join(SAVE_DIR, 'val_x.npy'), X_val)
np.save(os.path.join(SAVE_DIR, 'val_y.npy'), y_val)
print("Done.")
generate_local_dataset()
In [ ]:
# this code is used to process the already downloaded drive files so we can run this cell withouth running first cell if we delete the runtime (dummy cell)
''' The Fix (Gradient Accumulation): The code I provided uses ACCUM_STEPS = 16.
It calculates the error for Image 1, but does not update the weights.
It adds Image 2, Image 3... up to Image 16.
Only then does it update the weights.
Result: Mathematically, the model "thinks" it used Batch Size 16. You get the stability of a large batch with the memory usage of a small batch.'''
import ee
import geemap
import os
import numpy as np
import rasterio
from google.colab import drive
# --- 1. MOUNT DRIVE ---
# This allows us to access the files you already saved
drive.mount('/content/drive')
# --- 2. INITIALIZE EARTH ENGINE ---
try:
ee.Initialize(project='[REDACTED_FOR_SECURITY]')
print(" Earth Engine Initialized.")
except:
ee.Authenticate()
ee.Initialize(project='[REDACTED_FOR_SECURITY]')
print(" Earth Engine Authenticated & Initialized.")
# --- 3. CONFIGURATION ---
DRIVE_FOLDER = 'SatMAE_Scratch_Results_12Frames'
SAVE_DIR = f'/content/drive/MyDrive/{DRIVE_FOLDER}/'
stack_file = os.path.join(SAVE_DIR, 'local_stack.tif')
mask_file = os.path.join(SAVE_DIR, 'local_mask.tif')
# --- 4. PROCESSING LOGIC (TIF -> NPY) ---
def process_existing_files():
if not os.path.exists(stack_file) or not os.path.exists(mask_file):
print(f" Error: Files not found in {SAVE_DIR}")
print("Please check your Google Drive folder name.")
return
print(f" Found Stack: {os.path.basename(stack_file)}")
print(f" Found Mask: {os.path.basename(mask_file)}")
print(" Processing into Training Data (Spatial Split + NoData Filter)...")
with rasterio.open(stack_file) as src_stack, rasterio.open(mask_file) as src_mask:
img_data = src_stack.read()
mask_data = src_mask.read(1)
# Padding Check (Ensure 120 channels: 12 frames * 10 bands)
H, W = mask_data.shape
if img_data.shape[0] < 120:
print(f" Warning: Got {img_data.shape[0]} bands. Padding with zeros.")
padding = np.zeros((120 - img_data.shape[0], H, W), dtype=img_data.dtype)
img_data = np.concatenate([img_data, padding], axis=0)
img_reshaped = img_data[:120].reshape(12, 10, H, W)
# Chip Generation
X_list, y_list = [], []
if H < 224 or W < 224:
print(" Error: ROI is smaller than 224x224.")
return
else:
# Stride 112 = 50% Overlap
for r in range(0, H-224, 112):
for c in range(0, W-224, 112):
chip_x = img_reshaped[:, :, r:r+224, c:c+224]
chip_y = mask_data[r:r+224, c:c+224]
# Filter Garbage (Discard if >30% is NoData/255)
if np.mean(chip_x == 255) > 0.30: continue
# Filter Empty Labels (Keep if >1% wheat)
if np.mean(chip_y > 0) > 0.01:
X_list.append(chip_x)
y_list.append(chip_y)
if len(X_list) == 0:
print(" No valid tiles found! (Mask might be empty or threshold too strict).")
return
# Spatial Split (Top 80% Train, Bottom 20% Val)
X = np.array(X_list).astype(np.uint8)
y = np.array(y_list).astype(np.uint8)[:, None, :, :]
split_idx = int(0.8 * len(X))
X_train, X_val = X[:split_idx], X[split_idx:]
y_train, y_val = y[:split_idx], y[split_idx:]
print("-" * 30)
print(f" Data Ready!")
print(f" Train Tiles: {len(X_train)}")
print(f" Val Tiles: {len(X_val)}")
# Save to Drive
np.save(os.path.join(SAVE_DIR, 'train_x.npy'), X_train)
np.save(os.path.join(SAVE_DIR, 'train_y.npy'), y_train)
np.save(os.path.join(SAVE_DIR, 'val_x.npy'), X_val)
np.save(os.path.join(SAVE_DIR, 'val_y.npy'), y_val)
print(" .npy files saved successfully.")
# Run the processing
process_existing_files()
#------------------------------------------------------------------------------------------------------
import os
import numpy as np
import rasterio
from google.colab import drive
# --- 1. MOUNT DRIVE ---
# We need to see your Google Drive files
drive.mount('/content/drive', force_remount=True)
# --- 2. CONFIGURATION ---
# This must match exactly where your files are
DRIVE_FOLDER = 'SatMAE_Scratch_Results_12Frames'
SAVE_DIR = f'/content/drive/MyDrive/{DRIVE_FOLDER}/'
stack_file = os.path.join(SAVE_DIR, 'local_stack.tif')
mask_file = os.path.join(SAVE_DIR, 'local_mask.tif')
# --- 3. PROCESSING FUNCTION (No GEE required) ---
def generate_dataset_from_drive():
print(f" Checking folder: {SAVE_DIR}")
if not os.path.exists(stack_file) or not os.path.exists(mask_file):
print(f" Error: Files not found!")
print(f" Looking for: {stack_file}")
print(" Please check if the folder name is correct.")
return
print(f"Found Stack: {os.path.basename(stack_file)}")
print(f" Found Mask: {os.path.basename(mask_file)}")
print(" Processing GeoTIFFs into NPY (Training Data)...")
# Open the files directly from Drive
with rasterio.open(stack_file) as src_stack, rasterio.open(mask_file) as src_mask:
img_data = src_stack.read()
mask_data = src_mask.read(1)
# 1. Padding Check (Ensure 120 channels)
H, W = mask_data.shape
if img_data.shape[0] < 120:
print(f" Warning: Got {img_data.shape[0]} bands. Padding with zeros.")
padding = np.zeros((120 - img_data.shape[0], H, W), dtype=img_data.dtype)
img_data = np.concatenate([img_data, padding], axis=0)
img_reshaped = img_data[:120].reshape(12, 10, H, W)
# 2. Chip Generation
X_list, y_list = [], []
if H < 224 or W < 224:
print(" Error: Image is smaller than 224x224. Cannot create tiles.")
return
else:
# Stride 112 = 50% Overlap
for r in range(0, H-224, 112):
for c in range(0, W-224, 112):
chip_x = img_reshaped[:, :, r:r+224, c:c+224]
chip_y = mask_data[r:r+224, c:c+224]
# Filter Garbage (255)
if np.mean(chip_x == 255) > 0.30: continue
# Filter Empty Labels
if np.mean(chip_y > 0) > 0.01:
X_list.append(chip_x)
y_list.append(chip_y)
if len(X_list) == 0:
print(" No valid tiles found! (Check if mask has wheat pixels).")
return
# 3. Spatial Split (Top 80% Train, Bottom 20% Val)
X = np.array(X_list).astype(np.uint8)
y = np.array(y_list).astype(np.uint8)[:, None, :, :]
split_idx = int(0.8 * len(X))
X_train, X_val = X[:split_idx], X[split_idx:]
y_train, y_val = y[:split_idx], y[split_idx:]
print("-" * 30)
print(f" Data Processed Successfully!")
print(f" Train Chips: {len(X_train)}")
print(f" Val Chips: {len(X_val)}")
# 4. Save NPY files back to Drive
np.save(os.path.join(SAVE_DIR, 'train_x.npy'), X_train)
np.save(os.path.join(SAVE_DIR, 'train_y.npy'), y_train)
np.save(os.path.join(SAVE_DIR, 'val_x.npy'), X_val)
np.save(os.path.join(SAVE_DIR, 'val_y.npy'), y_val)
print(" .npy files saved.")
# --- RUN IT ---
generate_dataset_from_drive()
In [ ]:
#CELL_2
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torch.optim.swa_utils import AveragedModel, SWALR, update_bn
from tqdm.auto import tqdm
# --- 1. METRICS & LOSS (GPU Optimized) ---
def compute_metrics(pred_probs, targets, threshold=0.5):
pred_mask = (pred_probs > threshold).float()
intersection = (pred_mask * targets).sum()
union = pred_mask.sum() + targets.sum() - intersection
iou = (intersection + 1e-6) / (union + 1e-6)
dice = (2 * intersection + 1e-6) / (pred_mask.sum() + targets.sum() + 1e-6)
return iou.item(), dice.item()
class FastHausdorffLoss(nn.Module):
"""
GPU-Efficient Approximation of Hausdorff Loss.
Uses MaxPool to approximate morphological erosion/dilation for boundary detection.
"""
def __init__(self):
super().__init__()
def forward(self, pred, gt):
# 1. Get Probabilities
probs = torch.sigmoid(pred)
# 2. Extract Edges (approximate gradient) using MaxPool
# Edge = Image - ERODE(Image)
p_bin = probs
t_bin = gt
# Invert -> MaxPool -> Invert = Erosion
p_eroded = -F.max_pool2d(-p_bin, kernel_size=3, stride=1, padding=1)
t_eroded = -F.max_pool2d(-t_bin, kernel_size=3, stride=1, padding=1)
p_edge = p_bin - p_eroded
t_edge = t_bin - t_eroded
# 3. Mean Absolute Error between edges
# This penalizes mismatch in boundaries
return (p_edge - t_edge).abs().mean()
class CompoundLoss(nn.Module):
def forward(self, inputs, targets):
# Dice
probs = torch.sigmoid(inputs)
inter = (probs * targets).sum()
dice_loss = 1 - (2. * inter / (probs.sum() + targets.sum() + 1e-6))
# Fast Hausdorff
fast_hd = FastHausdorffLoss()(inputs, targets)
return 0.7 * dice_loss + 0.3 * fast_hd
class TestTimeAugmentation:
def __init__(self, model): self.model = model
def apply(self, x):
preds = []
with torch.no_grad():
preds.append(torch.sigmoid(self.model(x)))
preds.append(torch.flip(torch.sigmoid(self.model(torch.flip(x, [-1]))), [-1]))
preds.append(torch.flip(torch.sigmoid(self.model(torch.flip(x, [-2]))), [-2]))
preds.append(torch.flip(torch.sigmoid(self.model(torch.flip(x, [-2,-1]))), [-2,-1]))
return torch.stack(preds).mean(dim=0)
# --- 2. MODEL ARCHITECTURE ---
class SatMAEPatchEmbed(nn.Module):
def __init__(self, img_size=224, patch_size=16, in_chans=10, embed_dim=768):
super().__init__()
self.num_patches = (img_size // patch_size) ** 2
self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
def forward(self, x):
B, T, C, H, W = x.shape
x = x.reshape(B * T, C, H, W)
x = self.proj(x).flatten(2).transpose(1, 2)
return x.reshape(B, T, -1, 768)
class SatMAEEncoder(nn.Module):
def __init__(self, img_size=224, patch_size=16, in_chans=10, num_frames=12, embed_dim=768, depth=12, num_heads=12):
super().__init__()
self.patch_embed = SatMAEPatchEmbed(img_size, patch_size, in_chans, embed_dim)
self.pos_embed_spatial = nn.Parameter(torch.zeros(1, 1, self.patch_embed.num_patches, embed_dim))
self.pos_embed_temporal = nn.Parameter(torch.zeros(1, num_frames, 1, embed_dim))
# Scratch Init
nn.init.trunc_normal_(self.pos_embed_spatial, std=0.02)
nn.init.trunc_normal_(self.pos_embed_temporal, std=0.02)
self.blocks = nn.TransformerEncoder(nn.TransformerEncoderLayer(embed_dim, num_heads, int(embed_dim*4), 0.1, 'gelu', batch_first=True), depth)
self.norm = nn.LayerNorm(embed_dim)
self.apply(self._init_weights)
def _init_weights(self, m):
if isinstance(m, nn.Linear):
nn.init.trunc_normal_(m.weight, std=0.02)
if m.bias is not None: nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.LayerNorm):
nn.init.constant_(m.bias, 0)
nn.init.constant_(m.weight, 1.0)
elif isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
def forward(self, x):
B, T, C, H, W = x.shape
x = self.patch_embed(x)
x = x + self.pos_embed_spatial + self.pos_embed_temporal
x = x.reshape(B, T * x.shape[2], -1)
x = self.norm(self.blocks(x))
return x.reshape(B, T, 14, 14, -1)
class SatMAESegmentation(nn.Module):
def __init__(self, img_size=224, patch_size=16, in_chans=10, num_frames=12, embed_dim=768):
super().__init__()
self.num_frames = num_frames; self.chans = in_chans
self.encoder = SatMAEEncoder(img_size, patch_size, in_chans, num_frames, embed_dim)
self.up1 = nn.ConvTranspose2d(embed_dim, 256, 2, 2)
self.conv1 = nn.Sequential(nn.Conv2d(256, 256, 3, 1, 1), nn.BatchNorm2d(256), nn.ReLU())
self.up2 = nn.ConvTranspose2d(256, 128, 2, 2)
self.conv2 = nn.Sequential(nn.Conv2d(128, 128, 3, 1, 1), nn.BatchNorm2d(128), nn.ReLU())
self.up3 = nn.ConvTranspose2d(128, 64, 2, 2)
self.conv3 = nn.Sequential(nn.Conv2d(64, 64, 3, 1, 1), nn.BatchNorm2d(64), nn.ReLU())
self.up4 = nn.ConvTranspose2d(64, 32, 2, 2)
self.conv4 = nn.Sequential(nn.Conv2d(32, 32, 3, 1, 1), nn.BatchNorm2d(32), nn.ReLU())
self.final = nn.Conv2d(32, 1, 1)
def forward(self, x):
if x.ndim == 4:
B, _, H, W = x.shape
x = x.reshape(B, self.num_frames, self.chans, H, W)
features = self.encoder(x)
x = features.mean(dim=1).permute(0, 3, 1, 2)
x = self.conv1(self.up1(x))
x = self.conv2(self.up2(x))
x = self.conv3(self.up3(x))
x = self.conv4(self.up4(x))
return self.final(x)
# --- 3. TRAINER ---
class SatMAETrainer:
def __init__(self, model, loaders, device, lr=1e-4):
self.model = model; self.loaders = loaders; self.device = device
self.optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)
self.scheduler = optim.lr_scheduler.CosineAnnealingLR(self.optimizer, T_max=150)
self.swa_model = AveragedModel(model)
self.swa_scheduler = SWALR(self.optimizer, swa_lr=lr)
self.criterion = CompoundLoss()
self.tta = TestTimeAugmentation(model)
self.history = {'train_loss': [], 'val_iou': []}
def fit(self, epochs=200):
print(f"Starting Scratch Training ({epochs} epochs)...")
use_swa = False
swa_start_epoch = 151
patience = 30
no_improv = 0
best_iou = 0
for ep in range(epochs):
self.model.train()
t_loss = 0
pbar = tqdm(self.loaders['train'], desc=f"Epoch {ep+1}/{epochs}", leave=False)
for batch in pbar:
if isinstance(batch, (list, tuple)): x, y = batch[0], batch[1]
else: x, y = batch
x, y = x.to(self.device), y.to(self.device)
self.optimizer.zero_grad()
logits = self.model(x)
loss = self.criterion(logits, y)
loss.backward()
self.optimizer.step()
t_loss += loss.item()
pbar.set_postfix({'loss': f"{loss.item():.4f}"})
# Validation
eval_model = self.swa_model if use_swa else self.model
eval_model.eval()
v_iou_sum = 0
with torch.no_grad():
for batch in self.loaders['val']:
if isinstance(batch, (list, tuple)): x, y = batch[0], batch[1]
else: x, y = batch
x, y = x.to(self.device), y.to(self.device)
if use_swa: probs = torch.sigmoid(eval_model(x))
else: probs = self.tta.apply(x)
iou, _ = compute_metrics(probs, y)
v_iou_sum += iou
avg_loss = t_loss / len(self.loaders['train'])
avg_iou = v_iou_sum / len(self.loaders['val']) if len(self.loaders['val']) > 0 else 0
self.history['train_loss'].append(avg_loss)
self.history['val_iou'].append(avg_iou)
print(f"Ep {ep+1} | Loss: {avg_loss:.4f} | Val IoU: {avg_iou:.4f} | SWA: {use_swa}")
if (ep + 1) % 5 == 0:
torch.save(self.model.state_dict(), f"checkpoint_ep{ep+1}.pth")
if not use_swa:
self.scheduler.step()
if avg_iou > best_iou:
best_iou = avg_iou
no_improv = 0
torch.save(self.model.state_dict(), 'best_model.pth')
print(f" >>> New Best IoU: {best_iou:.4f}")
else:
no_improv += 1
if no_improv >= patience:
print(f" !! No Improvement. Triggering SWA.")
use_swa = True
swa_start_epoch = ep + 1
if ep + 1 >= swa_start_epoch: use_swa = True
else:
self.swa_model.update_parameters(self.model)
self.swa_scheduler.step()
if use_swa:
update_bn(self.loaders['train'], self.swa_model, device=self.device)
torch.save(self.swa_model.module.state_dict(), 'final_swa_model.pth')
print("Training Complete. SWA Model Saved.")
In [ ]:
#CELL_3
import torch
import numpy as np
import os
from torch.utils.data import Dataset, DataLoader
# --- 1. MEMORY EFFICIENT DATASET (Optimized) ---
class SatMAEDataset(Dataset):
def __init__(self, x_path, y_path):
# Load in mmap mode (Zero RAM usage initially)
self.x_data = np.load(x_path, mmap_mode='r')
self.y_data = np.load(y_path, mmap_mode='r')
def __len__(self):
return len(self.x_data)
def __getitem__(self, idx):
# 1. Read ONLY this specific index into RAM
x_np = self.x_data[idx] # Shape: (12, 10, 224, 224)
y_np = self.y_data[idx] # Shape: (1, 224, 224)
# 2. NUMPY OPTIMIZATION (Faster than torch.where)
# Create a float copy for scaling
x_float = x_np.astype(np.float32)
y_float = y_np.astype(np.float32)
# Zero out '255' (NoData) directly in numpy
# This prevents "Bright White" artifacts in the model
x_float[x_np == 255] = 0.0
y_float[y_np == 255] = 0.0
# Scale 0-250 -> 0-1
x_float /= 250.0
# 3. Convert to Tensor
return torch.from_numpy(x_float), torch.from_numpy(y_float)
# --- 2. EXECUTION ---
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using Device: {device}")
# Initialize Model (12 Frames, 10 Channels)
model = SatMAESegmentation(img_size=224, patch_size=16, in_chans=10, num_frames=12).to(device)
# Load Data
SAVE_DIR = '/content/drive/MyDrive/SatMAE_Scratch_Results_12Frames/'
try:
print("Initializing Data Loaders...")
train_ds = SatMAEDataset(os.path.join(SAVE_DIR, 'train_x.npy'), os.path.join(SAVE_DIR, 'train_y.npy'))
val_ds = SatMAEDataset(os.path.join(SAVE_DIR, 'val_x.npy'), os.path.join(SAVE_DIR, 'val_y.npy'))
# Num_workers=2 allows the CPU to fetch the next batch while GPU trains
loaders = {
'train': DataLoader(train_ds, batch_size=16, shuffle=True, num_workers=2),
'val': DataLoader(val_ds, batch_size=16, shuffle=False, num_workers=2)
}
# Start Training
print("Starting Training Loop...")
trainer = SatMAETrainer(model, loaders, device)
trainer.fit(epochs=200)
except FileNotFoundError:
print(f"ERROR: Data files not found in {SAVE_DIR}. Run Cell 1 (GEE Generation) first.")
except Exception as e:
import traceback
traceback.print_exc()
print(f"An error occurred: {e}")
In [ ]:
#CELL_4
import torch
import torch.optim as optim
from torch.optim.swa_utils import AveragedModel, SWALR, update_bn
import os
import csv
import pandas as pd
import matplotlib.pyplot as plt
# --- CONFIGURATION ---
SAVE_DIR = '/content/drive/MyDrive/SatMAE_Scratch_Results_12Frames/'
CHECKPOINT_PATH = os.path.join(SAVE_DIR, "checkpoint_latest.pth")
BEST_MODEL_PATH = os.path.join(SAVE_DIR, "best_model.pth")
LOG_PATH = os.path.join(SAVE_DIR, "training_log.csv")
BATCH_SIZE = 16
MAX_EPOCHS = 5000
PATIENCE_TRIGGER = 100 # If no improvement for 100 epochs -> Trigger SWA
SWA_DURATION = 50 # Run SWA for 50 epochs then STOP
# Ensure Directory
if not os.path.exists(SAVE_DIR): os.makedirs(SAVE_DIR)
# --- SETUP ---
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Device: {device}")
# Model (12 Frames, 12 Channels)
model = SatMAESegmentation(img_size=224, patch_size=16, in_chans=12, num_frames=12).to(device)
criterion = CompoundLoss()
optimizer = optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-4)
scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=50, T_mult=2)
# SWA Setup
swa_model = AveragedModel(model)
swa_scheduler = SWALR(optimizer, swa_lr=5e-5)
# Loaders (Re-using from Cell 3)
# train_ds, val_ds must be defined from previous cell
loaders = {
'train': DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=2),
'val': DataLoader(val_ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=2)
}
# --- RESUME LOGIC ---
start_epoch = 0
best_iou = 0.0
patience_counter = 0
swa_active = False
swa_epoch_counter = 0
history = {'epoch': [], 'train_loss': [], 'val_loss': [], 'val_iou': []}
if os.path.exists(CHECKPOINT_PATH):
print("Found checkpoint. Resuming...")
ckpt = torch.load(CHECKPOINT_PATH, map_location=device)
model.load_state_dict(ckpt['model_state_dict'])
optimizer.load_state_dict(ckpt['optimizer_state_dict'])
scheduler.load_state_dict(ckpt['scheduler_state_dict'])
start_epoch = ckpt['epoch'] + 1
best_iou = ckpt['best_iou']
patience_counter = ckpt['patience_counter']
swa_active = ckpt['swa_active']
swa_epoch_counter = ckpt['swa_epoch_counter']
if swa_active:
swa_model.load_state_dict(ckpt['swa_state_dict'])
swa_scheduler.load_state_dict(ckpt['swa_scheduler_state_dict'])
# Load history from CSV
if os.path.exists(LOG_PATH):
df = pd.read_csv(LOG_PATH)
history['epoch'] = df['epoch'].tolist()
history['train_loss'] = df['train_loss'].tolist()
history['val_loss'] = df['val_loss'].tolist()
history['val_iou'] = df['val_iou'].tolist()
print(f"Resumed at Epoch {start_epoch}. Best IoU: {best_iou:.4f}")
# --- TRAINING LOOP ---
print(f" Starting Training. Max: {MAX_EPOCHS} Eps. Patience: {PATIENCE_TRIGGER}")
try:
for ep in range(start_epoch, MAX_EPOCHS):
model.train()
train_loss = 0
# Training Step
for x, y in tqdm(loaders['train'], desc=f"Ep {ep+1}", leave=False):
x, y = x.to(device), y.to(device)
optimizer.zero_grad()
preds = model(x)
loss = criterion(preds, y)
loss.backward()
optimizer.step()
train_loss += loss.item()
# Validation Step
eval_model = swa_model if swa_active else model
eval_model.eval()
val_loss = 0
val_iou_sum = 0
with torch.no_grad():
for x, y in loaders['val']:
x, y = x.to(device), y.to(device)
# Forward
if swa_active: preds = eval_model(x)
else: preds = model(x) # No TTA for speed during training loop
val_loss += criterion(preds, y).item()
# Metrics
probs = torch.sigmoid(preds)
iou, _ = compute_metrics(probs, y)
val_iou_sum += iou
# Stats
avg_t = train_loss / len(loaders['train'])
avg_v = val_loss / len(loaders['val'])
avg_iou = val_iou_sum / len(loaders['val'])
# Log to History
history['epoch'].append(ep+1)
history['train_loss'].append(avg_t)
history['val_loss'].append(avg_v)
history['val_iou'].append(avg_iou)
# Append to CSV immediately (Crash Safety)
with open(LOG_PATH, 'a', newline='') as f:
writer = csv.writer(f)
if ep == 0 and not os.path.exists(LOG_PATH):
writer.writerow(['epoch', 'train_loss', 'val_loss', 'val_iou'])
writer.writerow([ep+1, avg_t, avg_v, avg_iou])
# Status Message
status = ""
# --- LOGIC BRANCHING ---
if swa_active:
# SWA PHASE
swa_model.update_parameters(model)
swa_scheduler.step()
swa_epoch_counter += 1
status = f"SWA Mode ({swa_epoch_counter}/{SWA_DURATION})"
# Stop Condition
if swa_epoch_counter >= SWA_DURATION:
print(f"Epoch {ep+1} | T: {avg_t:.4f} | V: {avg_v:.4f} | IoU: {avg_iou:.4f} | {status}")
print(" SWA Complete. Saving Final Model.")
update_bn(loaders['train'], swa_model, device=device)
torch.save(swa_model.module.state_dict(), os.path.join(SAVE_DIR, "final_swa_model.pth"))
break # EXIT LOOP
else:
# NORMAL PHASE
scheduler.step()
if avg_iou > best_iou:
best_iou = avg_iou
patience_counter = 0
torch.save(model.state_dict(), BEST_MODEL_PATH)
status = f" Best IoU!"
else:
patience_counter += 1
status = f"No Improv ({patience_counter}/{PATIENCE_TRIGGER})"
# Trigger SWA?
if patience_counter >= PATIENCE_TRIGGER:
print(f" Patience Limit Reached. Triggering SWA for {SWA_DURATION} epochs...")
swa_active = True
swa_epoch_counter = 0
# Load best weights before starting SWA to ensure stability
model.load_state_dict(torch.load(BEST_MODEL_PATH))
print(f"Ep {ep+1} | T: {avg_t:.4f} | V: {avg_v:.4f} | IoU: {avg_iou:.4f} | {status}")
# --- SAVE ROLLING CHECKPOINT ---
# Overwrites the same file every epoch. Saves storage.
torch.save({
'epoch': ep,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'scheduler_state_dict': scheduler.state_dict(),
'swa_state_dict': swa_model.state_dict() if swa_active else None,
'swa_scheduler_state_dict': swa_scheduler.state_dict() if swa_active else None,
'best_iou': best_iou,
'patience_counter': patience_counter,
'swa_active': swa_active,
'swa_epoch_counter': swa_epoch_counter
}, CHECKPOINT_PATH)
except KeyboardInterrupt:
print("Training Interrupted. Checkpoint Saved.")
# Plot Results
plt.figure(figsize=(10, 5))
plt.plot(history['train_loss'], label='Train Loss')
plt.plot(history['val_loss'], label='Val Loss')
plt.legend()
plt.title("Training Curves")
plt.savefig(os.path.join(SAVE_DIR, "loss_curve.png"))
plt.show()