In [1]:
import os
import time
import json
import numpy as np
import cv2
import rasterio
from rasterio.windows import from_bounds
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, random_split
from scipy.ndimage import distance_transform_edt as distance
from huggingface_hub import hf_hub_download
import ee
import geemap
from google.colab import drive
# Mount Drive
drive.mount('/content/drive', force_remount=True)
os.system('pip install -q rasterio geopandas timm segmentation-models-pytorch huggingface_hub geedim')
# Authenticate GEE
try:
ee.Initialize(project='[REDACTED_FOR_SECURITY]')
except:
ee.Authenticate()
ee.Initialize(project='[REDACTED_FOR_SECURITY]')
# Configuration
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Device: {device}")
if str(device) == 'cuda':
print(f"GPU: {torch.cuda.get_device_name(0)}")
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
# Paths and Hyperparameters
SAVE_DIR = '/content/drive/MyDrive/SatMAE_Advanced_Results/'
if not os.path.exists(SAVE_DIR): os.makedirs(SAVE_DIR)
BATCH_SIZE = 8
EPOCHS = 50
LR = 1e-4
PATCH_SIZE = 224
ASSET_ID = 'projects/[REDACTED_FOR_SECURITY]/assets/Punjab_Mask_2024_NEW'
TIME_WINDOWS = [
('2024-11-01', '2024-11-30'),
('2025-02-15', '2025-03-15'),
('2025-04-01', '2025-04-15')
]
Mounted at /content/drive Device: cuda GPU: Tesla T4
/usr/local/lib/python3.12/dist-packages/torch/backends/__init__.py:46: UserWarning: Please use the new API settings to control TF32 behavior, such as torch.backends.cudnn.conv.fp32_precision = 'tf32' or torch.backends.cuda.matmul.fp32_precision = 'ieee'. Old settings, e.g, torch.backends.cuda.matmul.allow_tf32 = True, torch.backends.cudnn.allow_tf32 = True, allowTF32CuDNN() and allowTF32CuBLAS() will be deprecated after Pytorch 2.9. Please see https://pytorch.org/docs/main/notes/cuda.html#tensorfloat-32-tf32-on-ampere-and-later-devices (Triggered internally at /pytorch/aten/src/ATen/Context.cpp:80.) self.setter(val)
In [2]:
def get_satmae_data():
print("Ingesting Asset...")
mask_img = ee.Image(ASSET_ID)
roi_geom = mask_img.geometry()
mask_file = 'local_mask.tif'
if not os.path.exists(mask_file):
geemap.download_ee_image(mask_img, mask_file, region=roi_geom, scale=10, crs='EPSG:4326', overwrite=True)
print("Processing Subset...")
with rasterio.open(mask_file) as src:
b = src.bounds
cx, cy = (b.left + b.right)/2, (b.bottom + b.top)/2
offset = 0.06
window = from_bounds(cx-offset, cy-offset, cx+offset, cy+offset, src.transform)
mask = src.read(1, window=window)
mask = np.where(mask > 0, 1.0, 0.0).astype(np.float32)
small_roi = ee.Geometry.Rectangle([cx-offset, cy-offset, cx+offset, cy+offset], proj=str(src.crs), geodesic=False)
target_h, target_w = mask.shape
stack = []
print("Stacking 3 Time Steps...")
for i, (start, end) in enumerate(TIME_WINDOWS):
fname = f'time_{i}.tif'
if not os.path.exists(fname):
s2 = ee.ImageCollection('COPERNICUS/S2_SR_HARMONIZED').filterBounds(small_roi).filterDate(start, end).median().select(['B2','B3','B4','B8','B11','B12'])
s1 = ee.ImageCollection('COPERNICUS/S1_GRD').filterBounds(small_roi).filterDate(start, end).mean().select(['VV','VH'])
fused = ee.Image.cat([s2, s1]).clip(small_roi)
geemap.download_ee_image(fused, fname, region=small_roi, scale=10, crs='EPSG:4326', overwrite=True)
with rasterio.open(fname) as src:
arr = src.read()
arr = np.transpose(arr, (1, 2, 0))
if arr.shape[:2] != (target_h, target_w):
arr = cv2.resize(arr, (target_w, target_h), interpolation=cv2.INTER_LINEAR)
s2_n = np.clip(arr[:,:,:6] / 5000.0, 0, 1)
s1_n = np.clip((arr[:,:,6:] - (-25.0)) / (0.0 - (-25.0)), 0, 1)
stack.append(np.concatenate([s2_n, s1_n], axis=2))
full_cube = np.stack(stack, axis=2)
print("Tiling...")
x_out, y_out = [], []
stride = PATCH_SIZE
for y in range(0, target_h, stride):
for x in range(0, target_w, stride):
img_p = full_cube[y:y+stride, x:x+stride]
mask_p = mask[y:y+stride, x:x+stride]
if img_p.shape[0] != PATCH_SIZE or img_p.shape[1] != PATCH_SIZE: continue
if np.min(img_p) < 0: continue
x_out.append(img_p)
y_out.append(mask_p)
if len(x_out) == 0: raise ValueError("No valid patches found.")
X = np.array(x_out, dtype=np.float32).transpose(0, 4, 3, 1, 2)
y = np.array(y_out, dtype=np.float32)[:, None, :, :]
print(f"Data Ready. Shape: {X.shape}")
return torch.tensor(X), torch.tensor(y)
X_data, y_data = get_satmae_data()
Ingesting Asset...
/usr/local/lib/python3.12/dist-packages/geemap/common.py:12471: FutureWarning: 'BaseImage' is deprecated and will be removed in a future release. Please use the 'ee.Image.gd' accessor instead. img = gd.download.BaseImage(image)
...tmae-2026/assets/Punjab_Mask_2024_NEW: 0%| |0/585 tiles [00:00<?]
WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:googleapiclient.http:Sleeping 0.49 seconds before retry 1 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/satmae-2026/thumbnails?fields=name&alt=json, after 429 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:googleapiclient.http:Sleeping 1.30 seconds before retry 1 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/satmae-2026/thumbnails?fields=name&alt=json, after 429 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 /usr/local/lib/python3.12/dist-packages/geedim/image.py:254: RuntimeWarning: Couldn't find STAC entry for: 'projects/satmae-2026/assets/Punjab_Mask_2024_NEW'. return STACClient().get(self.id)
Processing Subset... Stacking 3 Time Steps...
0%| |0/48 tiles [00:00<?]
WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 /usr/local/lib/python3.12/dist-packages/geedim/image.py:254: RuntimeWarning: Couldn't find STAC entry for: 'None'. return STACClient().get(self.id)
0%| |0/48 tiles [00:00<?]
WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:googleapiclient.http:Sleeping 1.86 seconds before retry 1 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/satmae-2026/thumbnails?fields=name&alt=json, after 429 WARNING:googleapiclient.http:Sleeping 0.67 seconds before retry 1 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/satmae-2026/thumbnails?fields=name&alt=json, after 429 WARNING:googleapiclient.http:Sleeping 3.84 seconds before retry 2 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/satmae-2026/thumbnails?fields=name&alt=json, after 429
0%| |0/48 tiles [00:00<?]
WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10
Tiling... Data Ready. Shape: (25, 8, 3, 224, 224)
In [3]:
class DiceLoss(nn.Module):
def __init__(self, smooth=1e-6):
super(DiceLoss, self).__init__()
self.smooth = smooth
def forward(self, inputs, targets):
inputs = torch.sigmoid(inputs)
inputs = inputs.view(-1)
targets = targets.view(-1)
intersection = (inputs * targets).sum()
dice = (2. * intersection + self.smooth) / (inputs.sum() + targets.sum() + self.smooth)
return 1 - dice
class HausdorffDTLoss(nn.Module):
def __init__(self, alpha=2.0):
super().__init__()
self.alpha = alpha
def forward(self, pred, gt):
with torch.no_grad():
gt_np = gt.cpu().numpy()
dist_map = np.zeros_like(gt_np)
for i in range(len(gt_np)):
mask = gt_np[i, 0]
if mask.sum() == 0: continue
d_in = distance(mask)
d_out = distance(1 - mask)
dist_map[i, 0] = (d_out - d_in)
dist_map = torch.tensor(dist_map, device=pred.device, dtype=torch.float32)
probs = torch.sigmoid(pred)
loss = torch.mean((probs - gt) ** 2 * (1 + self.alpha * torch.abs(dist_map)))
return loss
class CompositeLoss(nn.Module):
def __init__(self):
super().__init__()
self.dice = DiceLoss()
self.hd = HausdorffDTLoss(alpha=2.0)
self.bce = nn.BCEWithLogitsLoss()
def forward(self, preds, targets):
return 0.4*self.dice(preds, targets) + 0.4*self.bce(preds, targets) + 0.2*self.hd(preds, targets)
In [4]:
class FlashAttentionBlock(nn.Module):
def __init__(self, dim, num_heads):
super().__init__()
self.num_heads = num_heads
self.head_dim = dim // num_heads
self.scale = self.head_dim ** -0.5
self.qkv = nn.Linear(dim, dim * 3, bias=True)
self.proj = nn.Linear(dim, dim)
self.norm1 = nn.LayerNorm(dim)
self.norm2 = nn.LayerNorm(dim)
self.mlp = nn.Sequential(
nn.Linear(dim, dim * 4),
nn.GELU(),
nn.Linear(dim * 4, dim)
)
def forward(self, x):
B, N, C = x.shape
qkv = self.qkv(self.norm1(x)).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
q, k, v = qkv[0], qkv[1], qkv[2]
with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=True):
attn = F.scaled_dot_product_attention(q, k, v)
attn = attn.transpose(1, 2).reshape(B, N, C)
x = x + self.proj(attn)
x = x + self.mlp(self.norm2(x))
return x
class SatMAEPlusPlus_Encoder(nn.Module):
def __init__(self, num_frames=3, in_chans=8, embed_dim=768, depth=12, num_heads=12):
super().__init__()
self.patch_embed = nn.Conv2d(in_chans, embed_dim, kernel_size=16, stride=16)
self.pos_embed = nn.Parameter(torch.zeros(1, 1, 196+1, embed_dim))
self.time_embed = nn.Parameter(torch.zeros(1, num_frames, 1, embed_dim))
self.cls_token = nn.Parameter(torch.zeros(1, 1, 1, embed_dim))
self.blocks = nn.ModuleList([
FlashAttentionBlock(embed_dim, num_heads) for _ in range(depth)
])
self.norm = nn.LayerNorm(embed_dim)
def forward(self, x):
B, C, T, H, W = x.shape
x = x.permute(0, 2, 1, 3, 4).reshape(B * T, C, H, W)
x = self.patch_embed(x).flatten(2).transpose(1, 2)
x = x.reshape(B, T, -1, 768)
x = x + self.time_embed
x = x.reshape(B, T*196, 768)
pos = self.pos_embed[:, :, 1:, :].expand(B, T, -1, -1).reshape(B, T*196, 768)
x = x + pos
cls = self.cls_token.expand(B, -1, -1, -1).reshape(B, 1, 768) + self.pos_embed[:, :, 0, :].expand(B, 1, 768)
x = torch.cat((cls, x), dim=1)
for blk in self.blocks:
x = blk(x)
x = self.norm(x)
return x
class RefinementBlock(nn.Module):
def __init__(self, in_c, out_c):
super().__init__()
self.block = nn.Sequential(
nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
nn.Conv2d(in_c, out_c, kernel_size=3, padding=1),
nn.BatchNorm2d(out_c),
nn.GELU(),
nn.Conv2d(out_c, out_c, kernel_size=3, padding=1),
nn.BatchNorm2d(out_c),
nn.GELU()
)
def forward(self, x):
return self.block(x)
class FeaturePyramidDecoder(nn.Module):
def __init__(self, embed_dim, num_frames):
super().__init__()
self.temporal_fuse = nn.Conv2d(embed_dim * num_frames, embed_dim, kernel_size=1)
self.stage1 = RefinementBlock(embed_dim, 256)
self.stage2 = RefinementBlock(256, 128)
self.stage3 = RefinementBlock(128, 64)
self.stage4 = RefinementBlock(64, 32)
self.final_head = nn.Conv2d(32, 1, kernel_size=1)
def forward(self, x, B, T):
x = x[:, 1:, :]
x = x.reshape(B, T, 14, 14, 768).permute(0, 1, 4, 2, 3).reshape(B, T*768, 14, 14)
x = self.temporal_fuse(x)
x = self.stage1(x)
x = self.stage2(x)
x = self.stage3(x)
x = self.stage4(x)
return self.final_head(x)
class SatMAE_PlusPlus(nn.Module):
def __init__(self):
super().__init__()
print("Assembling SatMAE++")
self.encoder = SatMAEPlusPlus_Encoder()
self.decoder = FeaturePyramidDecoder(embed_dim=768, num_frames=3)
try:
print("Loading Google ViT Weights...")
p = hf_hub_download("google/vit-base-patch16-224", "pytorch_model.bin")
sd = torch.load(p, map_location='cpu')
w = sd['vit.embeddings.patch_embeddings.projection.weight']
new_w = torch.zeros(768, 8, 16, 16)
new_w[:, :3] = w
new_w[:, 3:] = w.mean(1, keepdim=True).repeat(1, 5, 1, 1)
self.encoder.patch_embed.weight.data = new_w
self.encoder.patch_embed.bias.data = sd['vit.embeddings.patch_embeddings.projection.bias']
for i in range(12):
prefix = f'vit.encoder.layer.{i}.'
blk = self.encoder.blocks[i]
q = sd[prefix + 'attention.attention.query.weight']
k = sd[prefix + 'attention.attention.key.weight']
v = sd[prefix + 'attention.attention.value.weight']
qb = sd[prefix + 'attention.attention.query.bias']
kb = sd[prefix + 'attention.attention.key.bias']
vb = sd[prefix + 'attention.attention.value.bias']
blk.qkv.weight.data = torch.cat([q, k, v], dim=0)
blk.qkv.bias.data = torch.cat([qb, kb, vb], dim=0)
blk.proj.weight.data = sd[prefix + 'attention.output.dense.weight']
blk.proj.bias.data = sd[prefix + 'attention.output.dense.bias']
blk.norm1.weight.data = sd[prefix + 'layernorm_before.weight']
blk.norm1.bias.data = sd[prefix + 'layernorm_before.bias']
blk.norm2.weight.data = sd[prefix + 'layernorm_after.weight']
blk.norm2.bias.data = sd[prefix + 'layernorm_after.bias']
blk.mlp[0].weight.data = sd[prefix + 'intermediate.dense.weight']
blk.mlp[0].bias.data = sd[prefix + 'intermediate.dense.bias']
blk.mlp[2].weight.data = sd[prefix + 'output.dense.weight']
blk.mlp[2].bias.data = sd[prefix + 'output.dense.bias']
print("Weights Adapted and Loaded Successfully.")
except Exception as e:
print(f"Weight loading failed: {e}. Using Random Init.")
for p in self.encoder.parameters(): p.requires_grad = False
self.encoder.patch_embed.weight.requires_grad = True
self.encoder.time_embed.requires_grad = True
for p in self.encoder.blocks[-2:].parameters(): p.requires_grad = True
def forward(self, x):
B, C, T, H, W = x.shape
features = self.encoder(x)
return self.decoder(features, B, T)
In [5]:
model = SatMAE_PlusPlus().to(device)
optimizer = optim.AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=LR)
criterion = CompositeLoss()
ds = TensorDataset(X_data, y_data)
tr_sz = int(0.85 * len(ds))
train_ds, val_ds = random_split(ds, [tr_sz, len(ds)-tr_sz])
train_loader = DataLoader(train_ds, BATCH_SIZE, shuffle=True)
val_loader = DataLoader(val_ds, BATCH_SIZE, shuffle=False)
print(f"Starting Training ({EPOCHS} Epochs)...")
history = []
for ep in range(EPOCHS):
model.train()
train_loss = 0
for x, y in train_loader:
x, y = x.to(device), y.to(device)
optimizer.zero_grad()
with torch.cuda.amp.autocast():
preds = model(x)
loss = criterion(preds, y)
loss.backward()
optimizer.step()
train_loss += loss.item()
model.eval()
val_loss = 0
with torch.no_grad():
for x, y in val_loader:
x, y = x.to(device), y.to(device)
preds = model(x)
val_loss += criterion(preds, y).item()
avg_t = train_loss / len(train_loader)
avg_v = val_loss / len(val_loader)
history.append((avg_t, avg_v))
if (ep+1) % 5 == 0:
print(f"Ep {ep+1} | Train Loss: {avg_t:.4f} | Val Loss: {avg_v:.4f}")
torch.save(model.state_dict(), SAVE_DIR + "SatMAE_PlusPlus_Final.pth")
print("Training Complete and Model Saved.")
Assembling SatMAE++ Loading Google ViT Weights...
/usr/local/lib/python3.12/dist-packages/huggingface_hub/utils/_auth.py:94: UserWarning: The secret `HF_TOKEN` does not exist in your Colab secrets. To authenticate with the Hugging Face Hub, create a token in your settings tab (https://huggingface.co/settings/tokens), set it as secret in your Google Colab and restart your session. You will be able to reuse this secret in all of your notebooks. Please note that authentication is recommended but still optional to access public models or datasets. warnings.warn(
pytorch_model.bin: 0%| | 0.00/346M [00:00<?, ?B/s]
Weights Adapted and Loaded Successfully. Starting Training (50 Epochs)...
/tmp/ipython-input-3025469134.py:22: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.
with torch.cuda.amp.autocast():
/usr/lib/python3.12/contextlib.py:105: FutureWarning: `torch.backends.cuda.sdp_kernel()` is deprecated. In the future, this context manager will be removed. Please see `torch.nn.attention.sdpa_kernel()` for the new context manager, with updated signature.
self.gen = func(*args, **kwds)
Ep 5 | Train Loss: 0.8669 | Val Loss: 1.0325 Ep 10 | Train Loss: 0.6625 | Val Loss: 0.7839 Ep 15 | Train Loss: 0.5679 | Val Loss: 0.6473 Ep 20 | Train Loss: 0.5182 | Val Loss: 0.5999 Ep 25 | Train Loss: 0.4840 | Val Loss: 0.5945 Ep 30 | Train Loss: 0.4643 | Val Loss: 0.5735 Ep 35 | Train Loss: 0.4482 | Val Loss: 0.5551 Ep 40 | Train Loss: 0.4300 | Val Loss: 0.5676 Ep 45 | Train Loss: 0.4242 | Val Loss: 0.5430 Ep 50 | Train Loss: 0.3946 | Val Loss: 0.5571 Training Complete and Model Saved.
In [6]:
import matplotlib.pyplot as plt
def plot_training_history(history):
train_loss = [h[0] for h in history]
val_loss = [h[1] for h in history]
epochs = range(1, len(history) + 1)
plt.figure(figsize=(10, 5))
plt.plot(epochs, train_loss, label='Training Loss', color='blue')
plt.plot(epochs, val_loss, label='Validation Loss', color='orange', linestyle='--')
plt.title('SatMAE++ Training Progress')
plt.xlabel('Epochs')
plt.ylabel('Composite Loss')
plt.legend()
plt.grid(True, alpha=0.3)
plt.show()
def visualize_predictions(model, loader, device, num_samples=3):
model.eval()
x_batch, y_batch = next(iter(loader))
x_batch, y_batch = x_batch.to(device), y_batch.to(device)
with torch.no_grad():
logits = model(x_batch)
preds = (torch.sigmoid(logits) > 0.5).float().cpu()
# Move inputs/targets to CPU for plotting
x_batch = x_batch.cpu()
y_batch = y_batch.cpu()
fig, axes = plt.subplots(num_samples, 3, figsize=(12, 4 * num_samples))
cols = ["Input (RGB - Peak Season)", "Ground Truth", "Prediction"]
# Set column titles
for ax, col in zip(axes[0], cols):
ax.set_title(col, fontweight='bold')
for i in range(num_samples):
# Extract RGB from Time Step 1 (Peak Growth)
# Data is (C, T, H, W) -> (8, 3, 224, 224)
# Indices: 0=B2(Blue), 1=B3(Green), 2=B4(Red)
# We need [2, 1, 0] for RGB
rgb = x_batch[i, [2, 1, 0], 1, :, :].permute(1, 2, 0).numpy()
# Brighten image for display (Inputs were normalized by 5000)
rgb = np.clip(rgb * 3.5, 0, 1)
# Plot Input
axes[i, 0].imshow(rgb)
axes[i, 0].axis('off')
# Plot Ground Truth
axes[i, 1].imshow(y_batch[i, 0], cmap='gray')
axes[i, 1].axis('off')
# Plot Prediction
axes[i, 2].imshow(preds[i, 0], cmap='gray')
axes[i, 2].axis('off')
# Calculate Sample IoU
intersection = (preds[i] * y_batch[i]).sum()
union = preds[i].sum() + y_batch[i].sum() - intersection
iou = intersection / (union + 1e-6)
axes[i, 2].text(5, 20, f"IoU: {iou:.2f}", color='white', fontweight='bold', bbox=dict(facecolor='black', alpha=0.5))
plt.tight_layout()
plt.show()
# Execute Visualization
if len(history) > 0:
plot_training_history(history)
visualize_predictions(model, val_loader, device)
else:
print("No training history found. Run training cell first.")
In [7]:
from sklearn.metrics import accuracy_score, f1_score, jaccard_score, precision_score, recall_score, confusion_matrix
import pandas as pd
def evaluate_and_save_metrics(model, loader, device, save_dir, model_name="SatMAE_PlusPlus"):
print(f"Calculating Metrics for {model_name}...")
model.eval()
all_preds = []
all_targets = []
# Inference Loop
start_time = time.time()
with torch.no_grad():
for x, y in loader:
x = x.to(device)
logits = model(x)
probs = torch.sigmoid(logits)
# Threshold at 0.5
preds = (probs > 0.5).float().cpu().numpy().flatten()
targets = y.cpu().numpy().flatten()
all_preds.extend(preds)
all_targets.extend(targets)
end_time = time.time()
total_time = end_time - start_time
fps = len(loader.dataset) / (total_time + 1e-6)
# Convert to integer arrays for metric calculation
y_pred = np.array(all_preds).astype(int)
y_true = np.array(all_targets).astype(int)
# Calculate Metrics
metrics = {
"IoU": round(jaccard_score(y_true, y_pred, average='binary'), 4),
"F1_Score": round(f1_score(y_true, y_pred, average='binary'), 4),
"Pixel_Accuracy": round(accuracy_score(y_true, y_pred), 4),
"Precision": round(precision_score(y_true, y_pred, average='binary'), 4),
"Recall": round(recall_score(y_true, y_pred, average='binary'), 4),
"FPS": round(fps, 2),
"Inference_Time_Sec": round(total_time, 2)
}
# Confusion Matrix
tn, fp, fn, tp = confusion_matrix(y_true, y_pred).ravel()
metrics["Confusion_Matrix"] = {
"TP": int(tp), "FP": int(fp),
"FN": int(fn), "TN": int(tn)
}
# Print Report
print("\n--- FINAL PERFORMANCE REPORT ---")
print(f"IoU Score: {metrics['IoU']}")
print(f"F1 Score: {metrics['F1_Score']}")
print(f"Precision: {metrics['Precision']}")
print(f"Recall: {metrics['Recall']}")
print(f"Inference Speed: {metrics['FPS']} FPS")
print("-" * 30)
# Save Metrics to JSON
json_path = f"{save_dir}{model_name}_Metrics.json"
with open(json_path, 'w') as f:
json.dump(metrics, f, indent=4)
print(f"Metrics saved to: {json_path}")
# Save History to CSV
if 'history' in globals() and len(history) > 0:
csv_path = f"{save_dir}{model_name}_History.csv"
df = pd.DataFrame(history, columns=['Train_Loss', 'Val_Loss'])
df['Epoch'] = range(1, len(history) + 1)
df.to_csv(csv_path, index=False)
print(f"History saved to: {csv_path}")
# Run Evaluation
evaluate_and_save_metrics(model, val_loader, device, SAVE_DIR)
Calculating Metrics for SatMAE_PlusPlus...
/usr/lib/python3.12/contextlib.py:105: FutureWarning: `torch.backends.cuda.sdp_kernel()` is deprecated. In the future, this context manager will be removed. Please see `torch.nn.attention.sdpa_kernel()` for the new context manager, with updated signature. self.gen = func(*args, **kwds)
--- FINAL PERFORMANCE REPORT --- IoU Score: 0.7165 F1 Score: 0.8348 Precision: 0.8437 Recall: 0.8262 Inference Speed: 15.84 FPS ------------------------------ Metrics saved to: /content/drive/MyDrive/SatMAE_Advanced_Results/SatMAE_PlusPlus_Metrics.json History saved to: /content/drive/MyDrive/SatMAE_Advanced_Results/SatMAE_PlusPlus_History.csv