# CELL 1: SETUP
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 scipy.spatial.distance import directed_hausdorff
from huggingface_hub import hf_hub_download
import ee
import geemap
from google.colab import drive

# 1. Environment
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 (Optimized for Fine-Tuning)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f" Training on: {device}")

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

BATCH_SIZE = 8
EPOCHS = 50
LR_MAX = 1e-4   # Peak LR
LR_MIN = 1e-6   # Min LR
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 Loading
def get_satmae_data():
    print("1. 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)
        target_h, target_w = mask.shape
        small_roi = ee.Geometry.Rectangle([cx-offset, cy-offset, cx+offset, cy+offset], proj=str(src.crs), geodesic=False)

    stack = []
    print("2. 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

    print("3. Creating Patches...")
    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
 Training on: cuda
1. 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: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.59 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)

2. 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
/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<?]

# CELL 2:  COMPOUND LOSS


class DiceLoss(nn.Module):
    def __init__(self, smooth=1e-6):
        super().__init__()
        self.smooth = smooth
    def forward(self, inputs, targets):
        inputs = torch.sigmoid(inputs).view(-1)
        targets = targets.view(-1)
        inter = (inputs * targets).sum()
        dice = (2. * inter + self.smooth) / (inputs.sum() + targets.sum() + self.smooth)
        return 1 - dice

class FocalTverskyLoss(nn.Module):
    def __init__(self, alpha=0.7, beta=0.3, gamma=2.0, smooth=1e-6):
        super().__init__()
        self.alpha = alpha
        self.beta = beta
        self.gamma = gamma
        self.smooth = smooth

    def forward(self, inputs, targets):
        inputs = torch.sigmoid(inputs).view(-1)
        targets = targets.view(-1)

        # Tversky Index
        TP = (inputs * targets).sum()
        FP = ((1-targets) * inputs).sum()
        FN = (targets * (1-inputs)).sum()

        Tversky = (TP + self.smooth) / (TP + self.alpha*FP + self.beta*FN + self.smooth)
        return (1 - Tversky)**self.gamma

class HausdorffDTLoss(nn.Module):
    def __init__(self, alpha=2.0):
        super().__init__()
        self.alpha = alpha

    def forward(self, pred, gt):
        # 1. Move to CPU for Scipy
        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] > 0.5).astype(np.uint8)

                if mask.sum() == 0: continue

                # d_in: dist to nearest zero (background)
                # d_out: dist to nearest one (foreground)
                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: Weighted MSE based on distance map
        loss = torch.mean((probs - gt) ** 2 * (1 + self.alpha * torch.abs(dist_map)))
        return loss

class CompoundLoss(nn.Module):
    def __init__(self):
        super().__init__()
        self.dice = DiceLoss()
        self.boundary = HausdorffDTLoss(alpha=2.0)
        self.focal = FocalTverskyLoss()

    def forward(self, preds, targets):
        # Weights: 50% Dice, 30% Boundary, 20% Focal
        return 0.5*self.dice(preds, targets) + \
               0.3*self.boundary(preds, targets) + \
               0.2*self.focal(preds, targets)

# CELL 3: SATMAE MODEL ( PARTIAL FINE-TUNE)

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_FineTune(nn.Module):
    def __init__(self, num_frames=3, embed_dim=768):
        super().__init__()
        print(" Initializing SatMAE (Partial Fine-Tune Mode)...")
        self.backbone = SatMAEBackbone(num_frames=num_frames, embed_dim=embed_dim)

        # 1. Load Pre-trained Weights
        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.")

        # 2. FREEZING STRATEGY

        # A. Freeze EVERYTHING first
        for p in self.backbone.blocks.parameters(): p.requires_grad = False

        # B. Unfreeze Inputs (Adapt to Sentinel-2/1)
        self.backbone.patch_embed.proj.weight.requires_grad = True
        self.backbone.time_embed.requires_grad = True

        # C. Unfreeze Last 2 Layers (Fixed Syntax)
        # We access .layers (ModuleList) -> slice it -> loop through layers -> loop through params
        for layer in self.backbone.blocks.layers[-2:]:
            for p in layer.parameters():
                p.requires_grad = True

        # 3. Standard Decoder
        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)

# CELL 4: TRAINING & BOUNDARY EVALUATION


def calculate_boundary_iou(gt_mask, pred_mask, dilation=5):
    """ Calculates IoU only in the 5px border region """
    gt_mask = gt_mask.astype(np.uint8)
    pred_mask = pred_mask.astype(np.uint8)
    kernel = np.ones((dilation, dilation), np.uint8)
    gt_b = cv2.dilate(gt_mask, kernel) - cv2.erode(gt_mask, kernel)
    pred_b = cv2.dilate(pred_mask, kernel) - cv2.erode(pred_mask, kernel)

    inter = np.logical_and(gt_b, pred_b).sum()
    union = np.logical_or(gt_b, pred_b).sum()
    if union == 0: return 1.0
    return inter / union

# Setup
model = SatMAE_FineTune().to(device)
criterion = CompoundLoss()
optimizer = optim.AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=LR_MAX)
scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=10, T_mult=2, eta_min=LR_MIN)

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 Advanced Fine-Tuning ({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()
        preds = model(x)
        loss = criterion(preds, y)
        loss.backward()
        optimizer.step()
        train_loss += loss.item()

    # Update Scheduler
    scheduler.step()
    current_lr = scheduler.get_last_lr()[0]

    # Validation & Metrics
    model.eval()
    val_loss = 0
    b_ious = []

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

            # Quick Boundary Check (Sample 1 per batch to save time)
            p_bin = (torch.sigmoid(preds) > 0.5).cpu().numpy().astype(np.uint8)
            y_bin = y.cpu().numpy().astype(np.uint8)
            b_ious.append(calculate_boundary_iou(y_bin[0,0], p_bin[0,0]))

    avg_t = train_loss / len(train_loader)
    avg_v = val_loss / len(val_loader)
    avg_b_iou = np.mean(b_ious)
    history.append(avg_v)

    if (ep+1) % 5 == 0:
        print(f"Ep {ep+1} | LR: {current_lr:.2e} | Train: {avg_t:.4f} | Val: {avg_v:.4f} | Boundary IoU: {avg_b_iou:.4f}")

torch.save(model.state_dict(), SAVE_DIR + "SatMAE_Advanced_FT_Final.pth")

🏗️ Initializing SatMAE (Partial Fine-Tune Mode)...
📥 Loading Weights...
✅ Weights Loaded.
 Starting Advanced Fine-Tuning (50 Epochs)...
Ep 5 | LR: 5.05e-05 | Train: 0.4842 | Val: 0.9113 | Boundary IoU: 0.1137
Ep 10 | LR: 1.00e-04 | Train: 0.4332 | Val: 0.5578 | Boundary IoU: 0.2069
Ep 15 | LR: 8.55e-05 | Train: 0.3724 | Val: 0.3377 | Boundary IoU: 0.2240
Ep 20 | LR: 5.05e-05 | Train: 0.3509 | Val: 0.3124 | Boundary IoU: 0.2332
Ep 25 | LR: 1.55e-05 | Train: 0.3424 | Val: 0.3216 | Boundary IoU: 0.2242
Ep 30 | LR: 1.00e-04 | Train: 0.3392 | Val: 0.3190 | Boundary IoU: 0.2375
Ep 35 | LR: 9.62e-05 | Train: 0.3278 | Val: 0.3009 | Boundary IoU: 0.2366
Ep 40 | LR: 8.55e-05 | Train: 0.3131 | Val: 0.2898 | Boundary IoU: 0.2462
Ep 45 | LR: 6.94e-05 | Train: 0.2992 | Val: 0.2984 | Boundary IoU: 0.2772
Ep 50 | LR: 5.05e-05 | Train: 0.2910 | Val: 0.3563 | Boundary IoU: 0.2887

# ==========================================
# CELL 5: VISUALIZATION & METRICS REPORTING
# ==========================================
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
import numpy as np
from datetime import datetime
from scipy.spatial.distance import directed_hausdorff
from sklearn.metrics import accuracy_score, f1_score, jaccard_score, precision_score, recall_score

# --- HELPER FUNCTIONS ---
def calculate_boundary_iou(gt_mask, pred_mask, dilation=5):
    """Calculates IoU only in the 5px border region."""
    import cv2
    gt_mask = gt_mask.astype(np.uint8)
    pred_mask = pred_mask.astype(np.uint8)
    kernel = np.ones((dilation, dilation), np.uint8)

    gt_b = cv2.dilate(gt_mask, kernel) - cv2.erode(gt_mask, kernel)
    pred_b = cv2.dilate(pred_mask, kernel) - cv2.erode(pred_mask, kernel)

    inter = np.logical_and(gt_b, pred_b).sum()
    union = np.logical_or(gt_b, pred_b).sum()
    if union == 0: return 1.0
    return inter / union

def symmetric_hausdorff(mask_pred, mask_gt):
    """
    Computes Symmetric Hausdorff Distance.
    FIX: directed_hausdorff returns (distance, idx1, idx2), so we take [0].
    """
    # Get coordinates of True pixels
    coords_pred = np.argwhere(mask_pred)
    coords_gt = np.argwhere(mask_gt)

    # If either is empty, return a penalty (diagonal of 224x224 image)
    if len(coords_pred) == 0 or len(coords_gt) == 0:
        return 316.0 # approx sqrt(224^2 + 224^2)

    # Calculate directed distances
    d_pg = directed_hausdorff(coords_pred, coords_gt)[0] # Fix: Access first element
    d_gp = directed_hausdorff(coords_gt, coords_pred)[0] # Fix: Access first element

    return max(d_pg, d_gp)

def generate_final_report(model, loader, history, save_dir, model_name="SatMAE_Full_FT"):
    print(f" Generating Final Report for {model_name}...")
    model.eval()

    # 1. VISUALIZATION (Save to Drive)
    # --------------------------------
    try:
        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()

        # Plot 4 Samples
        fig, axes = plt.subplots(4, 3, figsize=(12, 16))
        cols = ["Input (RGB - Peak Season)", "Ground Truth", "Prediction"]
        for ax, col in zip(axes[0], cols): ax.set_title(col, fontsize=14, fontweight='bold')

        for i in range(4):
            if i >= len(x_batch): break

            # Extract RGB (B4, B3, B2) -> Indices [2, 1, 0]
            rgb = x_batch[i, [2, 1, 0], 1, :, :].permute(1, 2, 0).cpu().numpy()
            rgb = np.clip(rgb * 3.5, 0, 1) # Brighten

            gt_img = y_batch[i, 0].cpu().numpy()
            pred_img = preds[i, 0].cpu().numpy()

            # Sample IoU
            inter = np.logical_and(gt_img, pred_img).sum()
            union = np.logical_or(gt_img, pred_img).sum()
            iou = inter / (union + 1e-6)

            axes[i, 0].imshow(rgb)
            axes[i, 1].imshow(gt_img, cmap='gray')
            axes[i, 2].imshow(pred_img, cmap='gray')
            axes[i, 2].text(5, 20, f"IoU: {iou:.2f}", color='lime', fontweight='bold', bbox=dict(facecolor='black', alpha=0.7))

            for ax in axes[i]: ax.axis('off')

        plt.tight_layout()
        viz_path = f"{save_dir}{model_name}_Visuals.png"
        plt.savefig(viz_path)
        plt.show()
        print(f"Visualizations saved to: {viz_path}")
    except Exception as e:
        print(f" Visualization skipped due to error: {e}")

    # 2. FULL DATASET METRICS
    # -----------------------
    print(" Calculating Metrics on Full Validation Set...")
    all_preds, all_targets = [], []
    boundary_ious = []
    hausdorff_dists = []

    with torch.no_grad():
        for x, y in loader:
            x = x.to(device)
            logits = model(x)
            p_batch = (torch.sigmoid(logits) > 0.5).float().cpu().numpy()
            y_batch_np = y.numpy()

            # Per-Image Metrics
            for j in range(len(y_batch_np)):
                p = p_batch[j, 0].astype(np.uint8)
                t = y_batch_np[j, 0].astype(np.uint8)

                boundary_ious.append(calculate_boundary_iou(t, p))

                if np.sum(t) > 0 and np.sum(p) > 0:
                    hausdorff_dists.append(symmetric_hausdorff(p, t))

            all_preds.extend(p_batch.flatten())
            all_targets.extend(y_batch_np.flatten())

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

    metrics = {
        "Model": model_name,
        "Date": datetime.now().strftime("%Y-%m-%d %H:%M"),
        "Standard_IoU": round(jaccard_score(y_t, y_p, average='binary'), 4),
        "Boundary_IoU": round(np.mean(boundary_ious), 4),
        "Hausdorff_Dist_px": round(np.mean(hausdorff_dists), 2),
        "F1_Score": 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),
        "Accuracy": round(accuracy_score(y_t, y_p), 4)
    }

    # 3. PLOT LOSS CURVE
    # ------------------
    if len(history) > 0:
        plt.figure(figsize=(10, 5))
        plt.plot(history, label='Validation Loss', color='orange', linewidth=2)
        plt.title(f'{model_name} Training Loss Curve')
        plt.xlabel('Epochs')
        plt.ylabel('Loss')
        plt.legend()
        plt.grid(True, alpha=0.3)
        loss_path = f"{save_dir}{model_name}_LossCurve.png"
        plt.savefig(loss_path)
        plt.show()

    # 4. SAVE TO DRIVE
    # ----------------
    json_path = f"{save_dir}{model_name}_Metrics.json"
    with open(json_path, 'w') as f:
        json.dump(metrics, f, indent=4)

    print("\nFINAL REPORT CARD:")
    print(json.dumps(metrics, indent=4))
    print(f" All files saved to {save_dir}")

# EXECUTE
generate_final_report(model, val_loader, history, SAVE_DIR, "SatMAE_Full_FineTune")

 Generating Final Report for SatMAE_Full_FineTune...

Visualizations saved to: /content/drive/MyDrive/SatMAE_Advanced_FT/SatMAE_Full_FineTune_Visuals.png
 Calculating Metrics on Full Validation Set...

FINAL REPORT CARD:
{
    "Model": "SatMAE_Full_FineTune",
    "Date": "2026-01-13 04:27",
    "Standard_IoU": 0.8213,
    "Boundary_IoU": 0.3582,
    "Hausdorff_Dist_px": 21.24,
    "F1_Score": 0.9019,
    "Precision": 0.9425,
    "Recall": 0.8646,
    "Accuracy": 0.8679
}
 All files saved to /content/drive/MyDrive/SatMAE_Advanced_FT/

# ==========================================
# CELL 5: VISUALIZATION & METRICS REPORTING
# ==========================================
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
import numpy as np
import json
import time
from datetime import datetime
from scipy.spatial.distance import directed_hausdorff
from sklearn.metrics import accuracy_score, f1_score, jaccard_score, precision_score, recall_score

# --- HELPER FUNCTIONS ---
def calculate_boundary_iou(gt_mask, pred_mask, dilation=5):
    """Calculates IoU only in the 5px border region."""
    import cv2
    gt_mask = gt_mask.astype(np.uint8)
    pred_mask = pred_mask.astype(np.uint8)
    kernel = np.ones((dilation, dilation), np.uint8)

    gt_b = cv2.dilate(gt_mask, kernel) - cv2.erode(gt_mask, kernel)
    pred_b = cv2.dilate(pred_mask, kernel) - cv2.erode(pred_mask, kernel)

    inter = np.logical_and(gt_b, pred_b).sum()
    union = np.logical_or(gt_b, pred_b).sum()
    if union == 0: return 1.0
    return inter / union

def symmetric_hausdorff(mask_pred, mask_gt):
    """
    Computes Symmetric Hausdorff Distance.
    Returns max(d(Pred, GT), d(GT, Pred)).
    """
    coords_pred = np.argwhere(mask_pred)
    coords_gt = np.argwhere(mask_gt)

    # Penalty if empty
    if len(coords_pred) == 0 or len(coords_gt) == 0:
        return 316.0 # approx diagonal of 224x224

    d_pg = directed_hausdorff(coords_pred, coords_gt)[0]
    d_gp = directed_hausdorff(coords_gt, coords_pred)[0]

    return max(d_pg, d_gp)

def generate_final_report(model, loader, history, save_dir, model_name="SatMAE_Full_FT"):
    print(f" Generating Final Report for {model_name}...")
    model.eval()

    # -------------------------------------------------------
    # 1. VISUALIZATION (Sample Predictions)
    # -------------------------------------------------------
    try:
        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()

        # Plot 4 Samples
        fig, axes = plt.subplots(4, 3, figsize=(12, 16))
        cols = ["Input (RGB - Peak Season)", "Ground Truth", "Prediction"]
        for ax, col in zip(axes[0], cols): ax.set_title(col, fontsize=14, fontweight='bold')

        for i in range(4):
            if i >= len(x_batch): break

            # Extract RGB (B4, B3, B2) -> Indices [2, 1, 0] from Time Step 1
            rgb = x_batch[i, [2, 1, 0], 1, :, :].permute(1, 2, 0).cpu().numpy()
            rgb = np.clip(rgb * 3.5, 0, 1) # Brighten

            gt_img = y_batch[i, 0].cpu().numpy()
            pred_img = preds[i, 0].cpu().numpy()

            # Sample IoU for Display
            inter = np.logical_and(gt_img, pred_img).sum()
            union = np.logical_or(gt_img, pred_img).sum()
            iou = inter / (union + 1e-6)

            axes[i, 0].imshow(rgb)
            axes[i, 1].imshow(gt_img, cmap='gray')
            axes[i, 2].imshow(pred_img, cmap='gray')
            axes[i, 2].text(5, 20, f"IoU: {iou:.2f}", color='lime', fontweight='bold', bbox=dict(facecolor='black', alpha=0.7))

            for ax in axes[i]: ax.axis('off')

        plt.tight_layout()
        viz_path = f"{save_dir}{model_name}_Visuals.png"
        plt.savefig(viz_path)
        plt.show()
        print(f" Visualizations saved to: {viz_path}")
    except Exception as e:
        print(f" Visualization skipped error: {e}")

    # -------------------------------------------------------
    # 2. FULL DATASET METRICS (With FPS)
    # -------------------------------------------------------
    print(" Calculating Metrics & FPS on Validation Set...")
    all_preds, all_targets = [], []
    boundary_ious = []
    hausdorff_dists = []

    # Start Timer
    start_time = time.time()

    with torch.no_grad():
        for x, y in loader:
            x = x.to(device)
            logits = model(x)
            p_batch = (torch.sigmoid(logits) > 0.5).float().cpu().numpy()
            y_batch_np = y.numpy()

            # Per-Image Metrics
            for j in range(len(y_batch_np)):
                p = p_batch[j, 0].astype(np.uint8)
                t = y_batch_np[j, 0].astype(np.uint8)

                boundary_ious.append(calculate_boundary_iou(t, p))

                if np.sum(t) > 0 and np.sum(p) > 0:
                    hausdorff_dists.append(symmetric_hausdorff(p, t))

            all_preds.extend(p_batch.flatten())
            all_targets.extend(y_batch_np.flatten())

    # Stop Timer
    end_time = time.time()
    total_time = end_time - start_time
    total_images = len(loader.dataset)
    fps = total_images / (total_time + 1e-6)

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

    metrics = {
        "Model": model_name,
        "Date": datetime.now().strftime("%Y-%m-%d %H:%M"),
        "Standard_IoU": round(jaccard_score(y_t, y_p, average='binary'), 4),
        "Boundary_IoU": round(np.mean(boundary_ious), 4),
        "Hausdorff_Dist_px": round(np.mean(hausdorff_dists), 2),
        "F1_Score": 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),
        "Accuracy": round(accuracy_score(y_t, y_p), 4),
        "Inference_Time_Sec": round(total_time, 2),
        "FPS": round(fps, 2)
    }

    # -------------------------------------------------------
    # 3. DUAL LOSS GRAPH (Train vs Validation)
    # -------------------------------------------------------
    if isinstance(history, dict) and 'train_loss' in history:
        train_loss = history['train_loss']
        val_loss = history['val_loss']
        epochs = range(1, len(train_loss) + 1)

        plt.figure(figsize=(10, 6))
        plt.plot(epochs, train_loss, 'b-', label='Training Loss', linewidth=2)
        plt.plot(epochs, val_loss, 'r--', label='Validation Loss', linewidth=2)

        plt.title(f'{model_name} Learning Curve', fontsize=16)
        plt.xlabel('Epochs', fontsize=12)
        plt.ylabel('Compound Loss', fontsize=12)
        plt.legend(fontsize=12)
        plt.grid(True, alpha=0.3)

        loss_path = f"{save_dir}{model_name}_DualLossCurve.png"
        plt.savefig(loss_path)
        plt.show()
        print(f" Loss Graph saved to: {loss_path}")
    else:
        print(" History format not recognized for dual plotting.")

    # -------------------------------------------------------
    # 4. SAVE RESULTS
    # -------------------------------------------------------
    json_path = f"{save_dir}{model_name}_Metrics.json"
    with open(json_path, 'w') as f:
        json.dump(metrics, f, indent=4)

    print("\n FINAL REPORT CARD:")
    print(json.dumps(metrics, indent=4))
    print(f"All files saved to {save_dir}")

# EXECUTE
generate_final_report(model, val_loader, history, SAVE_DIR, "SatMAE_Full_FineTune")

 Generating Final Report for SatMAE_Full_FineTune...

 Visualizations saved to: /content/drive/MyDrive/SatMAE_Advanced_FT/SatMAE_Full_FineTune_Visuals.png
 Calculating Metrics & FPS on Validation Set...
 History format not recognized for dual plotting.

 FINAL REPORT CARD:
{
    "Model": "SatMAE_Full_FineTune",
    "Date": "2026-01-13 04:34",
    "Standard_IoU": 0.8213,
    "Boundary_IoU": 0.3582,
    "Hausdorff_Dist_px": 21.24,
    "F1_Score": 0.9019,
    "Precision": 0.9425,
    "Recall": 0.8646,
    "Accuracy": 0.8679,
    "Inference_Time_Sec": 0.49,
    "FPS": 8.16
}
All files saved to /content/drive/MyDrive/SatMAE_Advanced_FT/

#Summary of the Work
#here i froze the first 10 block but unfroze the last 2 transformer blocks , it is just for experimenting but the result i got was poor than our satmae-finetuned model     FINAL STANDARD METRICS:
{"results of our Satmae finetunned with Compound loss "
    "Model": "SatMAE_Code1_Fair",
    "Pixel_Accuracy": 0.8757,-
    "IoU_Score": 0.8464,
    "F1_Score": 0.9168,
    "Precision": 0.8918,
    "Recall": 0.9433

    "Standard_IoU": 0.8464,
    "Boundary_IoU": 0.3481,
    "Hausdorff_Dist": 18.72,
    "FPS": 6.46
}
{

  # --- FINAL PERFORMANCE REPORT of our Satmae finetunned withouth Compounnd loss (hausdroff loss)
 Pixel Accuracy:   0.8732  (Overall correct pixels)
 IoU Score:        0.8358  (Intersection over Union - Target > 0.70)
 F1 / Dice Score:  0.9106   (Harmonic Mean of Precision & Recall)
 Precision:        0.9456  (Low False Positives)
 Recall:           0.8781  (Low False Negatives)

 Confusion Matrix Counts:
   True Wheat (TP):  129,585 pixels
   False Wheat (FP): 7,461 pixels (Over-segmentation)
   Missed Wheat (FN):17,991 pixels (Under-segmentation)
   Background (TN):  45,667 pixels

}
{"Results of the experimented 10 layer frozen and last 2 layer unfrozen"
    "Model": "SatMAE_partial_FineTune",
    "Date": "2026-01-13 04:34",
    "Standard_IoU": 0.8213,
    "Boundary_IoU": 0.3582,
    "Hausdorff_Dist_px": 21.24,
    "F1_Score": 0.9019,
    "Precision": 0.9425,
    "Recall": 0.8646,
    "Accuracy": 0.8679,
    "Inference_Time_Sec": 0.49,
    "FPS": 8.16
}

#After that i start working on Satmae finetunned with compound loss with 500 epochs but the results were poor than the 50 epochs model thenn i change the time 3 to 6 months and increase channel to 8 to 10 and implementing new methods and techniques to overcome the overfitting and over training issues
# Right now that colab notebook is running , when it will be completed i will share to you