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.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

# 1. Setup
drive.mount('/content/drive', force_remount=True)
os.system('pip install -q rasterio geopandas timm segmentation-models-pytorch huggingface_hub geedim')

try:
    ee.Initialize(project='[REDACTED_FOR_SECURITY]')
except:
    ee.Authenticate()
    ee.Initialize(project='[REDACTED_FOR_SECURITY]')

# 2. Config - FORCED CPU
device = torch.device('cpu')
print(f" GPU Limit Reached. Training on: {device}")

SAVE_DIR = '/content/drive/MyDrive/SatMAE_CPU_Results/'
if not os.path.exists(SAVE_DIR): os.makedirs(SAVE_DIR)

# Reduced Batch Size for CPU RAM safety
BATCH_SIZE = 4
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')
]

# 3. Data Loader
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)

    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 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)
    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)

    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()

Mounted at /content/drive
 GPU Limit Reached. Training on: cpu
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:googleapiclient.http:Sleeping 0.05 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 1.90 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.48 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.12 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 1.01 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
/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)

Stacking 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
WARNING:googleapiclient.http:Sleeping 1.20 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.08 seconds before retry 1 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/satmae-2026/thumbnails?fields=name&alt=json, after 429
/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

  0%|          |0/48 tiles [00:00<?]

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).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):
        # Already on CPU,
        with torch.no_grad():
            gt_np = gt.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, 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)

class SatMAEPatchEmbed(nn.Module):
    def __init__(self, in_chans=8, embed_dim=768, patch_size=16):
        super().__init__()
        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
    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.proj(x).flatten(2).transpose(1, 2)
        x = x.reshape(B, T, -1, x.shape[-1])
        return x

class SatMAEBackbone(nn.Module):
    def __init__(self, num_frames=3, in_chans=8, embed_dim=768, depth=12, num_heads=12):
        super().__init__()
        self.patch_embed = SatMAEPatchEmbed(in_chans=in_chans, embed_dim=embed_dim)
        num_patches = (224 // 16) ** 2
        self.pos_embed = nn.Parameter(torch.zeros(1, 1, num_patches + 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))

        encoder_layer = nn.TransformerEncoderLayer(d_model=embed_dim, nhead=num_heads, dim_feedforward=embed_dim*4, activation="gelu", batch_first=True, norm_first=True)
        self.blocks = nn.TransformerEncoder(encoder_layer, num_layers=depth)
        self.norm = nn.LayerNorm(embed_dim)

    def forward(self, x):
        x = self.patch_embed(x)
        B, T, N, D = x.shape
        x = x + self.time_embed
        x = x.reshape(B, T*N, D)
        spatial_pos = self.pos_embed[:, :, 1:, :].expand(B, T, -1, -1).reshape(B, T*N, D)
        x = x + spatial_pos
        cls_token = self.cls_token.expand(B, -1, -1, -1).reshape(B, 1, D) + self.pos_embed[:, :, 0, :].expand(B, 1, D)
        x = torch.cat((cls_token, x), dim=1)
        x = self.blocks(x)
        x = self.norm(x)
        return x

class SatMAE_Efficient(nn.Module):
    def __init__(self, num_frames=3, embed_dim=768):
        super().__init__()
        print(" Initializing SatMAE (Standard)...")
        self.backbone = SatMAEBackbone(num_frames=num_frames, embed_dim=embed_dim)

        try:
            print(" Loading 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.backbone.patch_embed.proj.weight.data = new_w
            self.backbone.patch_embed.proj.bias.data = sd['vit.embeddings.patch_embeddings.projection.bias']
            print("Weights Loaded.")
        except:
            print(" Weights missing, using random init.")

        # Partial Freeze
        for p in self.backbone.blocks.parameters(): p.requires_grad = False
        self.backbone.patch_embed.proj.weight.requires_grad = True
        self.backbone.time_embed.requires_grad = True

        self.temp_agg = nn.Conv2d(embed_dim * num_frames, embed_dim, kernel_size=1)
        self.decoder = nn.Sequential(
            nn.Upsample(scale_factor=2), nn.Conv2d(768, 256, 3, 1, 1), nn.BatchNorm2d(256), nn.GELU(),
            nn.Upsample(scale_factor=2), nn.Conv2d(256, 128, 3, 1, 1), nn.BatchNorm2d(128), nn.GELU(),
            nn.Upsample(scale_factor=2), nn.Conv2d(128, 64, 3, 1, 1), nn.BatchNorm2d(64), nn.GELU(),
            nn.Upsample(scale_factor=2), nn.Conv2d(64, 32, 3, 1, 1), nn.BatchNorm2d(32), nn.GELU(),
            nn.Conv2d(32, 1, 1)
        )

    def forward(self, x):
        features = self.backbone(x)[:, 1:, :]
        B, L, D = features.shape
        features = features.view(B, 3, 14, 14, D).permute(0, 4, 1, 2, 3).flatten(1, 2)
        features = self.temp_agg(features)
        return self.decoder(features)

model = SatMAE_Efficient().to(device) # device is cpu
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))
t_ds, v_ds = random_split(ds, [tr_sz, len(ds)-tr_sz])
train_loader = DataLoader(t_ds, BATCH_SIZE, shuffle=True)
val_loader = DataLoader(v_ds, BATCH_SIZE, shuffle=False)

print(f" Starting CPU Training ({EPOCHS} Epochs). This will be slow...")
history = []

for ep in range(EPOCHS):
    start = time.time()
    model.train()
    train_loss = 0

    for x, y in train_loader:
        # No .to(device) needed if everything is already CPU tensors
        optimizer.zero_grad()
        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:
            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))

    duration = time.time() - start
    print(f"Ep {ep+1} | Train: {avg_t:.4f} | Val: {avg_v:.4f} | Time: {duration:.1f}s")

torch.save(model.state_dict(), SAVE_DIR + "SatMAE_CPU.pth")
print(" Training Complete.")

🏗️ Initializing SatMAE (Standard)...

/usr/local/lib/python3.12/dist-packages/torch/nn/modules/transformer.py:392: UserWarning: enable_nested_tensor is True, but self.use_nested_tensor is False because encoder_layer.norm_first was True
  warnings.warn(

 Loading 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 Loaded.
 Starting CPU Training (50 Epochs). This will be slow...
Ep 1 | Train: 0.8907 | Val: 1.0165 | Time: 211.4s
Ep 2 | Train: 0.6234 | Val: 0.9342 | Time: 193.5s
Ep 3 | Train: 0.5258 | Val: 0.8174 | Time: 192.9s
Ep 4 | Train: 0.4793 | Val: 0.7793 | Time: 194.3s
Ep 5 | Train: 0.4557 | Val: 0.5353 | Time: 193.6s
Ep 6 | Train: 0.4288 | Val: 0.4738 | Time: 188.2s
Ep 7 | Train: 0.4290 | Val: 0.4689 | Time: 190.1s
Ep 8 | Train: 0.4132 | Val: 0.4364 | Time: 188.7s
Ep 9 | Train: 0.4098 | Val: 0.3938 | Time: 190.0s

import matplotlib.pyplot as plt
from sklearn.metrics import accuracy_score, f1_score, jaccard_score, precision_score, recall_score, confusion_matrix

def visualize_and_evaluate(model, loader):
    print(f"Evaluating...")
    model.eval()

    # 1. VISUALIZATION
    x_batch, y_batch = next(iter(loader))
    with torch.no_grad():
        logits = model(x_batch)
        preds = (torch.sigmoid(logits) > 0.5).float()

    fig, axes = plt.subplots(3, 3, figsize=(12, 12))
    cols = ["Input (Peak)", "Ground Truth", "Prediction"]
    for ax, col in zip(axes[0], cols): ax.set_title(col, fontweight='bold')

    for i in range(3):
        rgb = x_batch[i, [2, 1, 0], 1, :, :].permute(1, 2, 0).numpy()
        rgb = np.clip(rgb * 3.5, 0, 1)
        axes[i, 0].imshow(rgb); axes[i, 0].axis('off')
        axes[i, 1].imshow(y_batch[i, 0], cmap='gray'); axes[i, 1].axis('off')
        axes[i, 2].imshow(preds[i, 0], cmap='gray'); axes[i, 2].axis('off')
    plt.tight_layout()
    plt.show()

    # 2. METRICS
    all_preds, all_targets = [], []
    with torch.no_grad():
        for x, y in loader:
            logits = model(x)
            preds = (torch.sigmoid(logits) > 0.5).float().numpy().flatten()
            targets = y.numpy().flatten()
            all_preds.extend(preds)
            all_targets.extend(targets)

    y_p = np.array(all_preds).astype(int)
    y_t = np.array(all_targets).astype(int)

    metrics = {
        "IoU": round(jaccard_score(y_t, y_p, average='binary'), 4),
        "F1": round(f1_score(y_t, y_p, average='binary'), 4),
        "Precision": round(precision_score(y_t, y_p, average='binary'), 4),
        "Recall": round(recall_score(y_t, y_p, average='binary'), 4)
    }

    print("\n--- FINAL METRICS ---")
    print(metrics)

    with open(f"{SAVE_DIR}SatMAE_CPU_Metrics.json", 'w') as f:
        json.dump(metrics, f, indent=4)

visualize_and_evaluate(model, val_loader)