import os
import time
import shutil
import numpy as np
import rasterio
from rasterio.windows import from_bounds
import cv2
import ee
import geemap
from google.colab import drive
from sklearn.model_selection import train_test_split

os.system('pip install -q geedim geemap')

drive.mount('/content/drive', force_remount=True)

# CONFIGURATION

MY_PROJECT_ID = '[REDACTED_FOR_SECURITY]' #[REDACTED_FOR_SECURITY]
ASSET_ID = 'projects/[REDACTED_FOR_SECURITY]/assets/WheatMasks/binary_Punjab_Wheat_2024_Sieved'
SAVE_DIR = '/content/drive/MyDrive/SatMAE_PartialFT_Results_1/'

if not os.path.exists(SAVE_DIR):
    os.makedirs(SAVE_DIR)

# EARTH ENGINE INIT
try:
    ee.Initialize(project=MY_PROJECT_ID)
    print("Earth Engine Initialized.")
except Exception as e:
    ee.Authenticate()
    ee.Initialize(project=MY_PROJECT_ID)
    print("Earth Engine Authenticated & Initialized.")

# CONSTANTS
PATCH_SIZE = 224
S2_SCALE = 10000.0
TIME_WINDOWS = [
    ('2024-10-15', '2024-11-15'),
    ('2025-01-01', '2025-01-31'),
    ('2025-02-15', '2025-03-15')
]

def generate_satmae_partial_data():
    print("Starting SatMAE Partial FT Data Generation...")

    # 1. Mask
    mask_img = ee.Image(ASSET_ID)
    roi_geom = mask_img.geometry()
    mask_file = 'local_mask_satmae.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.04
        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.BBox(cx-offset, cy-offset, cx+offset, cy+offset)

    # 2. Imagery
    stack = []
    for i, (start, end) in enumerate(TIME_WINDOWS):
        fname = f'satmae_time_{i}.tif'
        attempts = 0
        while not os.path.exists(fname) and attempts < 3:
            try:
                print(f"Downloading T{i+1}: {start} to {end}...")
                img = ee.ImageCollection('COPERNICUS/S2_SR_HARMONIZED').filterBounds(small_roi).filterDate(start, end).median().select(['B2','B3','B4','B8','B11','B12'])
                geemap.download_ee_image(img, fname, region=small_roi, scale=10, crs='EPSG:4326', overwrite=True)
            except:
                attempts += 1
                time.sleep(2)

        if not os.path.exists(fname):
            with rasterio.open(mask_file) as src:
                 profile = src.profile
                 profile.update(count=6, dtype=rasterio.float32)
                 with rasterio.open(fname, 'w', **profile) as dst:
                     dst.write(np.zeros((6, target_h, target_w), dtype=np.float32))

        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)
            arr = np.clip(arr / S2_SCALE, 0, 1).astype(np.float32)
            stack.append(arr)

    # 3. Tiling
    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.mean(mask_p) < 0.01 or np.isnan(img_p).any(): 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"Dataset Generated. Shape: {X.shape}")
    X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=42)

    np.save(os.path.join(SAVE_DIR, 'train_x.npy'), X_train)
    np.save(os.path.join(SAVE_DIR, 'train_y.npy'), y_train)
    np.save(os.path.join(SAVE_DIR, 'val_x.npy'), X_val)
    np.save(os.path.join(SAVE_DIR, 'val_y.npy'), y_val)
    print("Done.")

generate_satmae_partial_data()

Mounted at /content/drive
Earth Engine Authenticated & Initialized.
Starting SatMAE Partial FT Data Generation...

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

...Masks/binary_Punjab_Wheat_2024_Sieved:   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
/usr/local/lib/python3.12/dist-packages/geedim/image.py:254: RuntimeWarning: Couldn't find STAC entry for: 'projects/phd0-473604/assets/WheatMasks/binary_Punjab_Wheat_2024_Sieved'.
  return STACClient().get(self.id)

Downloading T1: 2024-10-15 to 2024-11-15...

  0%|          |0/12 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
/usr/local/lib/python3.12/dist-packages/geedim/image.py:254: RuntimeWarning: Couldn't find STAC entry for: 'None'.
  return STACClient().get(self.id)

Downloading T2: 2025-01-01 to 2025-01-31...

  0%|          |0/12 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

import torch
import torch.nn as nn
from huggingface_hub import hf_hub_download

class SatMAEPatchEmbed(nn.Module):
    def __init__(self, in_chans=6, 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=6, 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 SatMAESegmentation(nn.Module):
    def __init__(self, num_frames=3, in_chans=6, embed_dim=768):
        super().__init__()
        print("Initializing SatMAE Partial FT Model...")
        self.num_frames = num_frames
        self.backbone = SatMAEBackbone(num_frames=num_frames, in_chans=in_chans, embed_dim=embed_dim)

        # --- AUTOMATIC WEIGHT LOADING ---
        print("Downloading Pretrained MAE Weights (facebook/vit-mae-base)...")
        try:
            p = hf_hub_download("facebook/vit-mae-base", "pytorch_model.bin")
            sd = torch.load(p, map_location='cpu')

            new_sd = {}
            for k, v in sd.items():

                # Adapt Patch Embeddings (3 Ch -> 6 Ch)
                if 'patch_embed.proj.weight' in k:
                    print(f"   Adapting Weights: {v.shape} -> 6 Channels")
                    new_w = torch.zeros(768, in_chans, 16, 16)
                    new_w[:, :3, :, :] = v  # Copy RGB
                    # Initialize Bands 4-6 with average of RGB
                    new_w[:, 3:, :, :] = v.mean(dim=1, keepdim=True).repeat(1, in_chans-3, 1, 1)
                    new_sd['backbone.patch_embed.proj.weight'] = new_w
                    continue

                if 'patch_embed.proj.bias' in k:
                    new_sd['backbone.patch_embed.proj.bias'] = v
                    continue

                # Backbone Block Mapping
                if 'blocks.' in k:
                    new_k = f"backbone.{k.replace('blocks.', 'blocks.layers.')}"
                    new_k = new_k.replace('mlp.fc1', 'linear1')
                    new_k = new_k.replace('mlp.fc2', 'linear2')
                    new_sd[new_k] = v
                elif 'norm.' in k:
                    new_sd[f"backbone.{k}"] = v
                elif 'pos_embed' in k:
                    if v.shape == self.backbone.pos_embed.shape:
                        new_sd['backbone.pos_embed'] = v
                elif 'cls_token' in k:
                    new_sd['backbone.cls_token'] = v

            missing, unexpected = self.backbone.load_state_dict(new_sd, strict=False)
            print(f"Weights Loaded! (Matched {len(new_sd)} keys)")

        except Exception as e:
            print(f"Weight Download Failed: {e}. Using Random Initialization.")

        # --- PARTIAL FREEZE STRATEGY (10 Frozen / 2 Unfrozen) ---
        print("Applying Partial Freeze Strategy...")

        # 1. Freeze ALL Blocks initially
        for param in self.backbone.blocks.parameters():
            param.requires_grad = False

        # 2. Unfreeze Last 2 Blocks (Layers 10 and 11)
        print("Unfreezing Last 2 Encoder Blocks...")
        for layer in self.backbone.blocks.layers[-2:]:
            for param in layer.parameters():
                param.requires_grad = True

        # 3. Unfreeze Embeddings & Norms
        self.backbone.patch_embed.proj.weight.requires_grad = True
        self.backbone.pos_embed.requires_grad = True
        self.backbone.time_embed.requires_grad = True
        self.backbone.cls_token.requires_grad = True
        self.backbone.norm.weight.requires_grad = True
        self.backbone.norm.bias.requires_grad = True

        # 4. Decoder
        self.temporal_agg = nn.Conv2d(embed_dim * num_frames, embed_dim, kernel_size=1)
        self.decoder = nn.Sequential(
            nn.Upsample(scale_factor=2), nn.Conv2d(embed_dim, 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
        H_p = 14
        features = features.view(B, self.num_frames, H_p, H_p, D)
        features = features.permute(0, 1, 4, 2, 3).reshape(B, self.num_frames * D, H_p, H_p)
        features = self.temporal_agg(features)
        return self.decoder(features)

import torch.optim as optim
import torch.optim.swa_utils as swa_utils
from torch.utils.data import Dataset, DataLoader
from scipy.ndimage import distance_transform_edt as distance
import numpy as np
import glob

def apply_augmentation(x, y):
    if np.random.rand() > 0.5: x = torch.flip(x, [4]); y = torch.flip(y, [3])
    if np.random.rand() > 0.5: x = torch.flip(x, [3]); y = torch.flip(y, [2])
    return x, y

def manage_rolling_checkpoints(save_dir, keep_k=5):
    files = sorted(glob.glob(os.path.join(save_dir, "epoch_*.pth")), key=os.path.getmtime)
    if len(files) > keep_k:
        for f in files[:-keep_k]: os.remove(f)

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

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] > 0.5).astype(np.uint8)
                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)
        return torch.mean((torch.sigmoid(pred) - gt) ** 2 * (1 + self.alpha * torch.abs(dist_map)))

class CompoundLoss(nn.Module):
    def __init__(self): super().__init__(); self.dice = DiceLoss(); self.boundary = HausdorffDTLoss()
    def forward(self, p, t): return 0.7*self.dice(p, t) + 0.3*self.boundary(p, t)

import torch
import torch.nn as nn
import torch.optim as optim
import torch.optim.swa_utils as swa_utils
from torch.utils.data import Dataset, DataLoader
import numpy as np
import os
import glob
import re
from google.colab import auth
from googleapiclient.discovery import build
from googleapiclient.errors import HttpError
from huggingface_hub import hf_hub_download

# ==========================================
# 1. TRASH CLEANER SETUP
# ==========================================
try:
    print("Authenticating for Trash Cleaner...")
    auth.authenticate_user()
except:
    print(" Authentication skipped/failed. Trash cleaning might not work.")

def empty_trash_specific(keyword="epoch"):
    """Permanently deletes files from Drive Trash matching the keyword."""
    try:
        service = build('drive', 'v3')
        query = f"trashed = true and name contains '{keyword}'"
        results = service.files().list(q=query, fields="files(id, name)").execute()
        items = results.get('files', [])

        if not items:
            print(f"    Trash is clean.")
            return

        print(f"    Deleting {len(items)} '{keyword}' files from Trash...")
        for item in items:
            try:
                service.files().delete(fileId=item['id']).execute()
            except:
                pass
        print("    Trash Purged.")
    except Exception as e:
        print(f"    Trash Error: {e}")

# ==========================================
# 2. MODEL DEFINITION (PARTIAL FINE-TUNE)
# ==========================================

class SatMAEPatchEmbed(nn.Module):
    def __init__(self, in_chans=6, 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=6, 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 SatMAESegmentation(nn.Module):
    def __init__(self, num_frames=3, in_chans=6, embed_dim=768):
        super().__init__()
        print(" Initializing SatMAE Partial FT Model...")
        self.num_frames = num_frames
        self.backbone = SatMAEBackbone(num_frames=num_frames, in_chans=in_chans, embed_dim=embed_dim)

        # --- AUTOMATIC WEIGHT LOADING ---
        print(" Downloading Pretrained MAE Weights (facebook/vit-mae-base)...")
        try:
            p = hf_hub_download("facebook/vit-mae-base", "pytorch_model.bin")
            sd = torch.load(p, map_location='cpu')

            new_sd = {}
            for k, v in sd.items():
                if 'patch_embed.proj.weight' in k:
                    new_w = torch.zeros(768, in_chans, 16, 16)
                    new_w[:, :3, :, :] = v  # Copy RGB
                    new_w[:, 3:, :, :] = v.mean(dim=1, keepdim=True).repeat(1, in_chans-3, 1, 1)
                    new_sd['backbone.patch_embed.proj.weight'] = new_w
                elif 'patch_embed.proj.bias' in k:
                    new_sd['backbone.patch_embed.proj.bias'] = v
                elif 'blocks.' in k:
                    new_k = f"backbone.{k.replace('blocks.', 'blocks.layers.')}"
                    new_k = new_k.replace('mlp.fc1', 'linear1').replace('mlp.fc2', 'linear2')
                    new_sd[new_k] = v
                elif 'norm.' in k: new_sd[f"backbone.{k}"] = v
                elif 'pos_embed' in k:
                    if v.shape == self.backbone.pos_embed.shape: new_sd['backbone.pos_embed'] = v
                elif 'cls_token' in k: new_sd['backbone.cls_token'] = v

            self.backbone.load_state_dict(new_sd, strict=False)
            print(f" Weights Loaded!")
        except Exception as e:
            print(f" Weight Download Failed: {e}. Random Init.")

        # --- PARTIAL FREEZE STRATEGY (10 Frozen / 2 Unfrozen) ---
        print("Applying Partial Freeze Strategy...")

        # 1. Freeze ALL Blocks initially
        for param in self.backbone.blocks.parameters():
            param.requires_grad = False

        # 2. Unfreeze Last 2 Blocks (Layers 10 and 11)
        print(" Unfreezing Last 2 Encoder Blocks...")
        for layer in self.backbone.blocks.layers[-2:]:
            for param in layer.parameters():
                param.requires_grad = True

        # 3. Unfreeze Embeddings & Norms
        self.backbone.patch_embed.proj.weight.requires_grad = True
        self.backbone.pos_embed.requires_grad = True
        self.backbone.time_embed.requires_grad = True
        self.backbone.cls_token.requires_grad = True
        self.backbone.norm.weight.requires_grad = True
        self.backbone.norm.bias.requires_grad = True

        # 4. Decoder
        self.temporal_agg = nn.Conv2d(embed_dim * num_frames, embed_dim, kernel_size=1)
        self.decoder = nn.Sequential(
            nn.Upsample(scale_factor=2), nn.Conv2d(embed_dim, 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, self.num_frames, 14, 14, D)
        features = features.permute(0, 1, 4, 2, 3).reshape(B, self.num_frames * D, 14, 14)
        features = self.temporal_agg(features)
        return self.decoder(features)

# ==========================================
# 3. CONFIGURATION & SETUP
# ==========================================
SAVE_DIR = '/content/drive/MyDrive/SatMAE_PartialFT_Results_1/'
CHECKPOINT_PATH = os.path.join(SAVE_DIR, "checkpoint_satmae_partial.pth")
BATCH_SIZE = 8
EPOCHS = 5000
PATIENCE_TRIGGER = 150
SWA_DURATION = 50

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Device: {device}")

model = SatMAESegmentation(num_frames=3, in_chans=6).to(device)
criterion = CompoundLoss()
optimizer = optim.AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=1e-4)
scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=50, T_mult=2)

swa_model = swa_utils.AveragedModel(model)
swa_scheduler = swa_utils.SWALR(optimizer, swa_lr=5e-5)

class MmapDataset(Dataset):
    def __init__(self, x_path, y_path):
        self.data = np.load(x_path, mmap_mode='r')
        self.target = np.load(y_path, mmap_mode='r')
    def __len__(self): return len(self.data)
    def __getitem__(self, index):
        return torch.from_numpy(self.data[index].copy()), torch.from_numpy(self.target[index].copy())

if not os.path.exists(SAVE_DIR): os.makedirs(SAVE_DIR)
train_ds = MmapDataset(os.path.join(SAVE_DIR, 'train_x.npy'), os.path.join(SAVE_DIR, 'train_y.npy'))
val_ds = MmapDataset(os.path.join(SAVE_DIR, 'val_x.npy'), os.path.join(SAVE_DIR, 'val_y.npy'))
train_loader = DataLoader(train_ds, BATCH_SIZE, shuffle=True)
val_loader = DataLoader(val_ds, BATCH_SIZE, shuffle=False)

# ==========================================
# 4. ROBUST RESUME LOGIC
# ==========================================
start_epoch = 0; best_loss = float('inf'); patience_counter = 0; swa_active = False; swa_epoch_counter = 0

if os.path.exists(CHECKPOINT_PATH):
    try:
        print(f"Attempting to load main checkpoint: {CHECKPOINT_PATH}")
        ckpt = torch.load(CHECKPOINT_PATH, map_location=device)
        model.load_state_dict(ckpt['model_state_dict'])
        optimizer.load_state_dict(ckpt['optimizer_state_dict'])
        start_epoch = ckpt['epoch'] + 1
        best_loss = ckpt['best_loss']
        patience_counter = ckpt.get('patience_counter', 0)
        swa_active = ckpt.get('swa_active', False)
        swa_epoch_counter = ckpt.get('swa_epoch_counter', 0)
        print(" Main checkpoint loaded successfully.")
    except Exception as e:
        print(f" Main checkpoint corrupted. Searching for backups...")
        epoch_files = glob.glob(os.path.join(SAVE_DIR, "epoch_*.pth"))
        if epoch_files:
            latest_file = max(epoch_files, key=os.path.getmtime)
            print(f" Rescuing from backup: {latest_file}")
            try:
                ckpt = torch.load(latest_file, map_location=device)
                model.load_state_dict(ckpt)
                match = re.search(r'epoch_(\d+).pth', latest_file)
                start_epoch = int(match.group(1)) if match else 0
                print(f"Rescue successful. Resuming from Epoch {start_epoch}.")
            except:
                print("Backup failed. Starting from Scratch.")
        else:
            print(" No backups found. Starting from Scratch.")

# ==========================================
# 5. TRAINING LOOP
# ==========================================
print(f" Starting Training. Max Epochs: {EPOCHS}")

for ep in range(start_epoch, EPOCHS):
    model.train(); train_loss = 0
    for x, y in train_loader:
        x, y = x.to(device), y.to(device)
        x, y = apply_augmentation(x, y)
        optimizer.zero_grad()
        loss = criterion(model(x), 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)
            val_loss += criterion(model(x), y).item()

    avg_t = train_loss / len(train_loader); avg_v = val_loss / len(val_loader)
    status_msg = ""

    if swa_active:
        swa_model.update_parameters(model); swa_scheduler.step()
        swa_epoch_counter += 1; status_msg = f"SWA Phase ({swa_epoch_counter}/{SWA_DURATION})"
        if swa_epoch_counter >= SWA_DURATION:
            print(" SWA Complete. Saving & Stopping.")
            swa_utils.update_bn(train_loader, swa_model, device=device)
            torch.save(swa_model.state_dict(), os.path.join(SAVE_DIR, "satmae_partial_swa_final.pth"))
            break
    else:
        scheduler.step()
        if avg_v < best_loss:
            best_loss = avg_v; patience_counter = 0; status_msg = "New Best Model!"
            torch.save(model.state_dict(), os.path.join(SAVE_DIR, "best_model.pth"))
        else:
            patience_counter += 1; status_msg = f"No Improvement ({patience_counter}/{PATIENCE_TRIGGER})"

        if patience_counter >= PATIENCE_TRIGGER:
            print(f" Patience limit reached. Triggering SWA.")
            swa_active = True; swa_epoch_counter = 0

    print(f"Epoch {ep+1} | Train: {avg_t:.4f} | Val: {avg_v:.4f} | {status_msg}")

    # Save Epoch Checkpoint
    epoch_ckpt = os.path.join(SAVE_DIR, f"epoch_{ep+1}.pth")
    torch.save(model.state_dict(), epoch_ckpt)

    # Maintenance
    manage_rolling_checkpoints(SAVE_DIR, keep_k=5)
    if (ep + 1) % 10 == 0:
        print(" Maintenance: Emptying 'epoch' files from Drive Trash...")
        empty_trash_specific(keyword="epoch")

    # Save Resume State
    torch.save({
        'epoch': ep,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'best_loss': best_loss,
        'patience_counter': patience_counter,
        'swa_active': swa_active,
        'swa_epoch_counter': swa_epoch_counter
    }, CHECKPOINT_PATH)

Authenticating for Trash Cleaner...
Device: cuda
 Initializing SatMAE Partial FT Model...
 Downloading Pretrained MAE Weights (facebook/vit-mae-base)...
 Weights Loaded!
Applying Partial Freeze Strategy...
 Unfreezing Last 2 Encoder Blocks...
Attempting to load main checkpoint: /content/drive/MyDrive/SatMAE_PartialFT_Results_1/checkpoint_satmae_partial.pth
 Main checkpoint corrupted. Searching for backups...
 Rescuing from backup: /content/drive/MyDrive/SatMAE_PartialFT_Results_1/epoch_121.pth
Rescue successful. Resuming from Epoch 121.
 Starting Training. Max Epochs: 5000
Epoch 122 | Train: 0.3627 | Val: 0.6601 | New Best Model!
Epoch 123 | Train: 0.4789 | Val: 4.2347 | No Improvement (1/150)
Epoch 124 | Train: 0.3872 | Val: 4.2443 | No Improvement (2/150)
Epoch 125 | Train: 0.3794 | Val: 3.4719 | No Improvement (3/150)
Epoch 126 | Train: 0.3666 | Val: 0.4990 | New Best Model!
Epoch 127 | Train: 0.3661 | Val: 0.2869 | New Best Model!
Epoch 128 | Train: 0.3667 | Val: 0.3365 | No Improvement (1/150)
Epoch 129 | Train: 0.3614 | Val: 0.3680 | No Improvement (2/150)
Epoch 130 | Train: 0.3578 | Val: 0.3907 | No Improvement (3/150)

WARNING:google_auth_httplib2:httplib2 transport does not support per-request timeout. Set the timeout when constructing the httplib2.Http instance.
WARNING:google_auth_httplib2:httplib2 transport does not support per-request timeout. Set the timeout when constructing the httplib2.Http instance.

 Maintenance: Emptying 'epoch' files from Drive Trash...
    Deleting 8 'epoch' files from Trash...
    Trash Purged.
Epoch 131 | Train: 0.3573 | Val: 0.4085 | No Improvement (4/150)
Epoch 132 | Train: 0.3565 | Val: 0.4167 | No Improvement (5/150)
Epoch 133 | Train: 0.3562 | Val: 0.4043 | No Improvement (6/150)
Epoch 134 | Train: 0.3529 | Val: 0.3796 | No Improvement (7/150)
Epoch 135 | Train: 0.3549 | Val: 0.3592 | No Improvement (8/150)
Epoch 136 | Train: 0.3533 | Val: 0.3494 | No Improvement (9/150)
Epoch 137 | Train: 0.3483 | Val: 0.3384 | No Improvement (10/150)
Epoch 138 | Train: 0.3512 | Val: 0.3212 | No Improvement (11/150)
Epoch 139 | Train: 0.3499 | Val: 0.3071 | No Improvement (12/150)
Epoch 140 | Train: 0.3487 | Val: 0.2961 | No Improvement (13/150)

WARNING:google_auth_httplib2:httplib2 transport does not support per-request timeout. Set the timeout when constructing the httplib2.Http instance.
WARNING:google_auth_httplib2:httplib2 transport does not support per-request timeout. Set the timeout when constructing the httplib2.Http instance.

 Maintenance: Emptying 'epoch' files from Drive Trash...
    Deleting 11 'epoch' files from Trash...
    Trash Purged.
Epoch 141 | Train: 0.3441 | Val: 0.2913 | No Improvement (14/150)
Epoch 142 | Train: 0.3477 | Val: 0.2945 | No Improvement (15/150)
Epoch 143 | Train: 0.3445 | Val: 0.3013 | No Improvement (16/150)
Epoch 144 | Train: 0.3405 | Val: 0.3099 | No Improvement (17/150)
Epoch 145 | Train: 0.3427 | Val: 0.3141 | No Improvement (18/150)
Epoch 146 | Train: 0.3402 | Val: 0.3201 | No Improvement (19/150)
Epoch 147 | Train: 0.3437 | Val: 0.3234 | No Improvement (20/150)
Epoch 148 | Train: 0.3419 | Val: 0.3338 | No Improvement (21/150)
Epoch 149 | Train: 0.3421 | Val: 0.3346 | No Improvement (22/150)
Epoch 150 | Train: 0.3364 | Val: 0.3322 | No Improvement (23/150)

WARNING:google_auth_httplib2:httplib2 transport does not support per-request timeout. Set the timeout when constructing the httplib2.Http instance.
WARNING:google_auth_httplib2:httplib2 transport does not support per-request timeout. Set the timeout when constructing the httplib2.Http instance.

 Maintenance: Emptying 'epoch' files from Drive Trash...
    Deleting 10 'epoch' files from Trash...
    Trash Purged.
Epoch 151 | Train: 0.3386 | Val: 0.3287 | No Improvement (24/150)
Epoch 152 | Train: 0.3377 | Val: 0.3190 | No Improvement (25/150)
Epoch 153 | Train: 0.3393 | Val: 0.3070 | No Improvement (26/150)
Epoch 154 | Train: 0.3393 | Val: 0.2969 | No Improvement (27/150)
Epoch 155 | Train: 0.3355 | Val: 0.2927 | No Improvement (28/150)
Epoch 156 | Train: 0.3360 | Val: 0.2920 | No Improvement (29/150)
Epoch 157 | Train: 0.3367 | Val: 0.2944 | No Improvement (30/150)
Epoch 158 | Train: 0.3377 | Val: 0.3001 | No Improvement (31/150)
Epoch 159 | Train: 0.3385 | Val: 0.3078 | No Improvement (32/150)
Epoch 160 | Train: 0.3348 | Val: 0.3150 | No Improvement (33/150)

WARNING:google_auth_httplib2:httplib2 transport does not support per-request timeout. Set the timeout when constructing the httplib2.Http instance.
WARNING:google_auth_httplib2:httplib2 transport does not support per-request timeout. Set the timeout when constructing the httplib2.Http instance.

 Maintenance: Emptying 'epoch' files from Drive Trash...
    Deleting 10 'epoch' files from Trash...
    Trash Purged.
Epoch 161 | Train: 0.3368 | Val: 0.3210 | No Improvement (34/150)
Epoch 162 | Train: 0.3373 | Val: 0.3269 | No Improvement (35/150)
Epoch 163 | Train: 0.3345 | Val: 0.3302 | No Improvement (36/150)
Epoch 164 | Train: 0.3349 | Val: 0.3301 | No Improvement (37/150)
Epoch 165 | Train: 0.3353 | Val: 0.3296 | No Improvement (38/150)
Epoch 166 | Train: 0.3326 | Val: 0.3280 | No Improvement (39/150)
Epoch 167 | Train: 0.3339 | Val: 0.3261 | No Improvement (40/150)
Epoch 168 | Train: 0.3354 | Val: 0.3242 | No Improvement (41/150)
Epoch 169 | Train: 0.3357 | Val: 0.3226 | No Improvement (42/150)
Epoch 170 | Train: 0.3321 | Val: 0.3210 | No Improvement (43/150)

WARNING:google_auth_httplib2:httplib2 transport does not support per-request timeout. Set the timeout when constructing the httplib2.Http instance.
WARNING:google_auth_httplib2:httplib2 transport does not support per-request timeout. Set the timeout when constructing the httplib2.Http instance.

import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
import os
import time
import cv2
from scipy.spatial.distance import directed_hausdorff
from sklearn.metrics import precision_score, recall_score, f1_score, accuracy_score, jaccard_score
from torch.utils.data import DataLoader, Dataset
from google.colab import drive

# 1. PATHS AND DEVICE SETUP
if not os.path.exists('/content/drive'):
    drive.mount('/content/drive', force_remount=True)

# TARGET DIRECTORY FOR PARTIAL FINETUNING
SAVE_DIR = '/content/drive/MyDrive/SatMAE_PartialFT_Results_1/'
CHECKPOINT_PATH = os.path.join(SAVE_DIR, "checkpoint_satmae_partial.pth")
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 2. DATA LOADER (MMAP for memory efficiency)
class MmapDataset(Dataset):
    def __init__(self, x_path, y_path):
        self.data = np.load(x_path, mmap_mode='r')
        self.target = np.load(y_path, mmap_mode='r')
    def __len__(self): return len(self.data)
    def __getitem__(self, index):
        return torch.from_numpy(self.data[index].copy()), torch.from_numpy(self.target[index].copy())

# 3. ADVANCED METRIC HELPERS
def calculate_boundary_iou(gt_mask, pred_mask):
    # Detect edges to calculate IoU specifically on the boundaries
    gt_edges = cv2.Canny(gt_mask.astype(np.uint8)*255, 100, 200) > 0
    pred_edges = cv2.Canny(pred_mask.astype(np.uint8)*255, 100, 200) > 0
    intersection = np.logical_and(gt_edges, pred_edges).sum()
    union = np.logical_or(gt_edges, pred_edges).sum()
    return intersection / union if union > 0 else 0.0

def calculate_hausdorff(gt_mask, pred_mask):
    # Measure the maximum distance between the predicted boundary and the ground truth boundary
    gt_coords = np.argwhere(gt_mask)
    pred_coords = np.argwhere(pred_mask)
    if len(gt_coords) == 0 or len(pred_coords) == 0: return 0.0
    return max(directed_hausdorff(gt_coords, pred_coords)[0], directed_hausdorff(pred_coords, gt_coords)[0])

# 4. EXECUTION FUNCTION
def evaluate_partial_ft(model_dir, checkpoint_path, num_samples=3):
    # Load Validation Data
    val_x = os.path.join(model_dir, 'val_x.npy')
    val_y = os.path.join(model_dir, 'val_y.npy')
    val_ds = MmapDataset(val_x, val_y)
    val_loader = DataLoader(val_ds, batch_size=8, shuffle=False)

    # Initialize Model
    model = SatMAESegmentation(num_frames=3, in_chans=6).to(DEVICE)

    # Load Weights
    if os.path.exists(checkpoint_path):
        print(f" Loading Checkpoint: {checkpoint_path}")
        ckpt = torch.load(checkpoint_path, map_location=DEVICE)
        model.load_state_dict(ckpt['model_state_dict'])
    else:
        print(" Error: Checkpoint file not found.")
        return

    model.eval()
    all_preds, all_targets = [], []
    boundary_ious, hausdorff_dists = [], []
    total_frames = 0
    start_time = time.time()

    print(" Calculating metrics...")
    with torch.no_grad():
        for x, y in val_loader:
            x = x.to(DEVICE)
            probs = torch.sigmoid(model(x))
            preds_bin = (probs > 0.5).cpu().numpy().astype(np.uint8)
            targets_bin = y.numpy().astype(np.uint8)
            total_frames += x.size(0)

            for i in range(preds_bin.shape[0]):
                p, t = preds_bin[i, 0], targets_bin[i, 0]
                boundary_ious.append(calculate_boundary_iou(t, p))
                hausdorff_dists.append(calculate_hausdorff(t, p))
                all_preds.append(p.flatten())
                all_targets.append(t.flatten())

    # Final Report Calculation
    fps = total_frames / (time.time() - start_time)
    y_p, y_t = np.concatenate(all_preds), np.concatenate(all_targets)

    print("\n" + "="*50)
    print("       SATMAE PARTIAL FT - PERFORMANCE       ")
    print("="*50)
    print(f"Pixel Accuracy:      {accuracy_score(y_t, y_p):.4f}")
    print(f"IoU (Jaccard):       {jaccard_score(y_t, y_p):.4f}")
    print(f"F1-Score (Dice):     {f1_score(y_t, y_p):.4f}")
    print(f"Precision:           {precision_score(y_t, y_p):.4f}")
    print(f"Recall:              {recall_score(y_t, y_p):.4f}")
    print("-" * 30)
    print(f"Boundary IoU:        {np.mean(boundary_ious):.4f}")
    print(f"Hausdorff Dist (px): {np.mean(hausdorff_dists):.2f}")
    print(f"Inference Speed:     {fps:.2f} FPS")
    print("="*50 + "\n")

    # VISUAL SAMPLES
    print(" Plotting Samples...")
    x_batch, y_batch = next(iter(val_loader))
    with torch.no_grad():
        vis_preds = (torch.sigmoid(model(x_batch.to(DEVICE))) > 0.5).cpu().numpy()

    x_np = x_batch.numpy()
    y_np = y_batch.numpy()

    fig, axs = plt.subplots(num_samples, 3, figsize=(15, 5 * num_samples))
    for i in range(num_samples):
        # Extract RGB at Time Step 0 to avoid IndexErrors
        # x_np shape: (Batch, Channels, Time, H, W)
        try:
            r, g, b = x_np[i, 2, 0], x_np[i, 1, 0], x_np[i, 0, 0]
            img = np.stack([r, g, b], axis=2)
        except:
            img = x_np[i, 0, 0] # Fallback to single band
            img = np.stack([img, img, img], axis=2)

        # Normalize 0-1
        img = (img - img.min()) / (img.max() - img.min() + 1e-6)

        axs[i, 0].imshow(img)
        axs[i, 0].set_title(f"Input Sample {i+1}")
        axs[i, 1].imshow(y_np[i, 0], cmap='gray')
        axs[i, 1].set_title("Ground Truth")
        axs[i, 2].imshow(vis_preds[i, 0], cmap='gray')

        iou = jaccard_score(y_np[i,0].flatten(), vis_preds[i,0].flatten())
        axs[i, 2].set_title(f"Prediction (IoU: {iou:.2f})")

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

    plt.tight_layout()
    plt.show()

# RUN EVALUATION
evaluate_partial_ft(SAVE_DIR, CHECKPOINT_PATH)

 Initializing SatMAE Partial FT Model...
 Downloading Pretrained MAE Weights (facebook/vit-mae-base)...
 Weights Loaded!
Applying Partial Freeze Strategy...
 Unfreezing Last 2 Encoder Blocks...
 Loading Checkpoint: /content/drive/MyDrive/SatMAE_PartialFT_Results_1/checkpoint_satmae_partial.pth

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
/tmp/ipython-input-2482868117.py in <cell line: 0>()
    143 
    144 # RUN EVALUATION
--> 145 evaluate_partial_ft(SAVE_DIR, CHECKPOINT_PATH)

/tmp/ipython-input-2482868117.py in evaluate_partial_ft(model_dir, checkpoint_path, num_samples)
     59     if os.path.exists(checkpoint_path):
     60         print(f" Loading Checkpoint: {checkpoint_path}")
---> 61         ckpt = torch.load(checkpoint_path, map_location=DEVICE)
     62         model.load_state_dict(ckpt['model_state_dict'])
     63     else:

/usr/local/lib/python3.12/dist-packages/torch/serialization.py in load(f, map_location, pickle_module, weights_only, mmap, **pickle_load_args)
   1519                 if weights_only:
   1520                     try:
-> 1521                         return _load(
   1522                             opened_zipfile,
   1523                             map_location,

/usr/local/lib/python3.12/dist-packages/torch/serialization.py in _load(zip_file, map_location, pickle_module, pickle_file, overall_storage, **pickle_load_args)
   2120     global _serialization_tls
   2121     _serialization_tls.map_location = map_location
-> 2122     result = unpickler.load()
   2123     _serialization_tls.map_location = None
   2124 

/usr/local/lib/python3.12/dist-packages/torch/_weights_only_unpickler.py in load(self)
    533                         f"Only persistent_load of storage is allowed, but got {pid[0]}"
    534                     )
--> 535                 self.append(self.persistent_load(pid))
    536             elif key[0] in [BINGET[0], LONG_BINGET[0]]:
    537                 idx = (read(1) if key[0] == BINGET[0] else unpack("<I", read(4)))[0]

/usr/local/lib/python3.12/dist-packages/torch/serialization.py in persistent_load(saved_id)
   2084         else:
   2085             nbytes = numel * torch._utils._element_size(dtype)
-> 2086             typed_storage = load_tensor(
   2087                 dtype, nbytes, key, _maybe_decode_ascii(location)
   2088             )

/usr/local/lib/python3.12/dist-packages/torch/serialization.py in load_tensor(dtype, numel, key, location)
   2037                     )
   2038             storage = (
-> 2039                 zip_file.get_storage_from_record(name, numel, torch.UntypedStorage)
   2040                 ._typed_storage()
   2041                 ._untyped_storage

RuntimeError: PytorchStreamReader failed locating file data/44: file not found

import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
import os
import time
import cv2
import glob
import re
from scipy.spatial.distance import directed_hausdorff
from sklearn.metrics import precision_score, recall_score, f1_score, accuracy_score, jaccard_score
from torch.utils.data import DataLoader, Dataset
from google.colab import drive

# 1. SETUP
if not os.path.exists('/content/drive'):
    drive.mount('/content/drive', force_remount=True)

SAVE_DIR = '/content/drive/MyDrive/SatMAE_PartialFT_Results_1/'
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 2. DATA LOADER
class MmapDataset(Dataset):
    def __init__(self, x_path, y_path):
        self.data = np.load(x_path, mmap_mode='r')
        self.target = np.load(y_path, mmap_mode='r')
    def __len__(self): return len(self.data)
    def __getitem__(self, index):
        return torch.from_numpy(self.data[index].copy()), torch.from_numpy(self.target[index].copy())

# 3. METRIC HELPERS
def calculate_boundary_iou(gt_mask, pred_mask):
    gt_edges = cv2.Canny(gt_mask.astype(np.uint8)*255, 100, 200) > 0
    pred_edges = cv2.Canny(pred_mask.astype(np.uint8)*255, 100, 200) > 0
    inter = np.logical_and(gt_edges, pred_edges).sum()
    union = np.logical_or(gt_edges, pred_edges).sum()
    return inter / union if union > 0 else 0.0

def calculate_hausdorff(gt_mask, pred_mask):
    gt_coords = np.argwhere(gt_mask)
    pred_coords = np.argwhere(pred_mask)
    if len(gt_coords) == 0 or len(pred_coords) == 0: return 0.0
    return max(directed_hausdorff(gt_coords, pred_coords)[0], directed_hausdorff(pred_coords, gt_coords)[0])

# 4. ROBUST LOAD FUNCTION
def get_best_available_weights(model_dir):
    # Order of preference: 1. Best Model, 2. Latest Epoch backup
    best_path = os.path.join(model_dir, "best_model.pth")
    if os.path.exists(best_path):
        return best_path

    epoch_files = glob.glob(os.path.join(model_dir, "epoch_*.pth"))
    if epoch_files:
        return max(epoch_files, key=os.path.getmtime)

    return None

# 5. EXECUTION
def run_evaluation(model_dir):
    weights_path = get_best_available_weights(model_dir)
    if not weights_path:
        print(" No weights found in directory.")
        return

    print(f" Loading weights from: {weights_path}")

    # Init Model
    model = SatMAESegmentation(num_frames=3, in_chans=6).to(DEVICE)
    model.load_state_dict(torch.load(weights_path, map_location=DEVICE))
    model.eval()

    # Data
    val_ds = MmapDataset(os.path.join(model_dir, 'val_x.npy'), os.path.join(model_dir, 'val_y.npy'))
    val_loader = DataLoader(val_ds, batch_size=8, shuffle=False)

    all_preds, all_targets = [], []
    boundary_ious, hausdorff_dists = [], []
    total_frames = 0
    start_time = time.time()

    print(" Evaluating...")
    with torch.no_grad():
        for x, y in val_loader:
            x = x.to(DEVICE)
            preds = torch.sigmoid(model(x))
            preds_bin = (preds > 0.5).cpu().numpy().astype(np.uint8)
            targets_bin = y.numpy().astype(np.uint8)
            total_frames += x.size(0)

            for i in range(preds_bin.shape[0]):
                p, t = preds_bin[i, 0], targets_bin[i, 0]
                boundary_ious.append(calculate_boundary_iou(t, p))
                hausdorff_dists.append(calculate_hausdorff(t, p))
                all_preds.append(p.flatten()); all_targets.append(t.flatten())

    # Report
    fps = total_frames / (time.time() - start_time)
    y_p, y_t = np.concatenate(all_preds), np.concatenate(all_targets)

    print("\n" + "="*50)
    print(f"Pixel Accuracy:      {accuracy_score(y_t, y_p):.4f}")
    print(f"IoU (Jaccard):       {jaccard_score(y_t, y_p):.4f}")
    print(f"F1-Score (Dice):     {f1_score(y_t, y_p):.4f}")
    print(f"Boundary IoU:        {np.mean(boundary_ious):.4f}")
    print(f"Hausdorff Dist:      {np.mean(hausdorff_dists):.2f}")
    print(f"Speed:               {fps:.2f} FPS")
    print("="*50)

    # Visualization
    x_batch, y_batch = next(iter(val_loader))
    with torch.no_grad():
        vis_preds = (torch.sigmoid(model(x_batch.to(DEVICE))) > 0.5).cpu().numpy()

    fig, axs = plt.subplots(3, 3, figsize=(15, 15))
    for i in range(3):
        # RGB extraction from (B, C, T, H, W) -> Using T=0
        img = x_batch[i, [2, 1, 0], 0].permute(1, 2, 0).numpy()
        img = (img - img.min()) / (img.max() - img.min() + 1e-6)

        axs[i, 0].imshow(img)
        axs[i, 0].set_title("Input RGB")
        axs[i, 1].imshow(y_batch[i, 0], cmap='gray')
        axs[i, 1].set_title("Ground Truth")
        axs[i, 2].imshow(vis_preds[i, 0], cmap='gray')
        axs[i, 2].set_title("Prediction")
        for ax in axs[i]: ax.axis('off')
    plt.show()

run_evaluation(SAVE_DIR)

🚀 Loading weights from: /content/drive/MyDrive/SatMAE_PartialFT_Results_1/best_model.pth
 Initializing SatMAE Partial FT Model...
 Downloading Pretrained MAE Weights (facebook/vit-mae-base)...
 Weights Loaded!
Applying Partial Freeze Strategy...
 Unfreezing Last 2 Encoder Blocks...
📊 Evaluating...

==================================================
Pixel Accuracy:      0.9143
IoU (Jaccard):       0.8858
F1-Score (Dice):     0.9394
Boundary IoU:        0.1070
Hausdorff Dist:      12.86
Speed:               1.49 FPS
==================================================

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
/tmp/ipython-input-458025177.py in <cell line: 0>()
    127     plt.show()
    128 
--> 129 run_evaluation(SAVE_DIR)

/tmp/ipython-input-458025177.py in run_evaluation(model_dir)
    115     for i in range(3):
    116         # RGB extraction from (B, C, T, H, W) -> Using T=0
--> 117         img = x_batch[i, [2, 1, 0], 0].permute(1, 2, 0).numpy()
    118         img = (img - img.min()) / (img.max() - img.min() + 1e-6)
    119 

IndexError: index 2 is out of bounds for dimension 0 with size 2

import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
import os
import time
import cv2
from scipy.spatial.distance import directed_hausdorff
from sklearn.metrics import (precision_score, recall_score, f1_score, accuracy_score,
                             jaccard_score, matthews_corrcoef, confusion_matrix,
                             roc_auc_score, precision_recall_curve, auc, fbeta_score)
from torch.utils.data import DataLoader, Dataset
from google.colab import drive

# 1. SETUP & PATHS
SAVE_DIR = '/content/drive/MyDrive/SatMAE_PartialFT_Results_1/'
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 2. METRIC HELPERS
def calculate_boundary_iou(gt_mask, pred_mask):
    gt_edges = cv2.Canny(gt_mask.astype(np.uint8)*255, 100, 200) > 0
    pred_edges = cv2.Canny(pred_mask.astype(np.uint8)*255, 100, 200) > 0
    inter = np.logical_and(gt_edges, pred_edges).sum()
    union = np.logical_or(gt_edges, pred_edges).sum()
    return inter / union if union > 0 else 0.0

def calculate_hausdorff(gt_mask, pred_mask):
    gt_coords = np.argwhere(gt_mask)
    pred_coords = np.argwhere(pred_mask)
    if len(gt_coords) == 0 or len(pred_coords) == 0: return 0.0
    return max(directed_hausdorff(gt_coords, pred_coords)[0], directed_hausdorff(pred_coords, gt_coords)[0])

# 3. COMPREHENSIVE EVALUATION FUNCTION
def run_final_metrics_report(model_dir):
    # Load Model Weights
    weights_path = os.path.join(model_dir, "best_model.pth")
    if not os.path.exists(weights_path):
        weights_path = os.path.join(model_dir, "satmae_finetune_swa_final.pth")

    print(f" Loading weights: {weights_path}")
    model = SatMAESegmentation(num_frames=3, in_chans=6).to(DEVICE)
    model.load_state_dict(torch.load(weights_path, map_location=DEVICE))
    model.eval()

    # Load Data
    val_ds = MmapDataset(os.path.join(model_dir, 'val_x.npy'), os.path.join(model_dir, 'val_y.npy'))
    val_loader = DataLoader(val_ds, batch_size=8, shuffle=False)

    all_preds, all_targets, all_probs = [], [], []
    boundary_ious, hausdorff_dists = [], []
    total_frames = 0
    start_time = time.time()

    print(" Evaluating Validation Set...")
    with torch.no_grad():
        for x, y in val_loader:
            x = x.to(DEVICE)
            outputs = model(x)
            probs = torch.sigmoid(outputs).cpu().numpy()
            preds_bin = (probs > 0.5).astype(np.uint8)
            targets_bin = y.numpy().astype(np.uint8)

            total_frames += x.size(0)

            for i in range(preds_bin.shape[0]):
                p, t = preds_bin[i, 0], targets_bin[i, 0]
                boundary_ious.append(calculate_boundary_iou(t, p))
                hausdorff_dists.append(calculate_hausdorff(t, p))
                all_preds.append(p.flatten())
                all_targets.append(t.flatten())
                all_probs.append(probs[i, 0].flatten())

    # Consolidate Data
    y_p = np.concatenate(all_preds)
    y_t = np.concatenate(all_targets)
    y_prob = np.concatenate(all_probs)

    # Advanced Statistical Calculations
    tn, fp, fn, tp = confusion_matrix(y_t, y_p).ravel()
    prec_pts, rec_pts, _ = precision_recall_curve(y_t, y_prob)
    pr_auc_val = auc(rec_pts, prec_pts)
    fps_val = total_frames / (time.time() - start_time)

    # PRINT DETAILED REPORT
    print("\n" + "="*60)
    print("      SATMAE FULL FINETUNE - IN-DEPTH METRICS REPORT       ")
    print("="*60)
    print(f"{'Metric':<28} | {'Value':<10}")
    print("-" * 45)
    print(f"{'Pixel Accuracy':<28} | {accuracy_score(y_t, y_p):.4f}")
    print(f"{'IoU (Jaccard Index)':<28} | {jaccard_score(y_t, y_p):.4f}")
    print(f"{'F1-Score (Dice)':<28} | {f1_score(y_t, y_p):.4f}")
    print(f"{'F2-Score (Recall Focus)':<28} | {fbeta_score(y_t, y_p, beta=2):.4f}")
    print(f"{'MCC (Matthews Corr)':<28} | {matthews_corrcoef(y_t, y_p):.4f}")
    print(f"{'ROC-AUC':<28} | {roc_auc_score(y_t, y_prob):.4f}")
    print(f"{'PR-AUC':<28} | {pr_auc_val:.4f}")
    print("-" * 45)
    print(f"{'Precision':<28} | {precision_score(y_t, y_p):.4f}")
    print(f"{'Recall (Sensitivity)':<28} | {recall_score(y_t, y_p):.4f}")
    print(f"{'Specificity':<28} | {tn / (tn + fp):.4f}")
    print(f"{'False Positive Rate':<28} | {fp / (fp + tn):.4f}")
    print("-" * 45)
    print(f"{'Boundary IoU':<28} | {np.mean(boundary_ious):.4f}")
    print(f"{'Hausdorff Distance (px)':<28} | {np.mean(hausdorff_dists):.2f}")
    print(f"{'Inference Speed':<28} | {fps_val:.2f} FPS")
    print("="*60 + "\n")

    # 4. PLOTTING BLOCK
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))

    # Confusion Matrix
    cm = confusion_matrix(y_t, y_p)
    im = ax1.imshow(cm, interpolation='nearest', cmap=plt.cm.Blues)
    ax1.set_title("Confusion Matrix")
    fig.colorbar(im, ax=ax1)
    ax1.set_xticks([0, 1])
    ax1.set_yticks([0, 1])
    ax1.set_xticklabels(['Background', 'Target'])
    ax1.set_yticklabels(['Background', 'Target'])
    ax1.set_xlabel('Predicted Label')
    ax1.set_ylabel('True Label')

    # PR Curve
    ax2.plot(rec_pts, prec_pts, color='darkgreen', lw=2, label=f'PR Curve (AUC = {pr_auc_val:.2f})')
    ax2.set_xlabel('Recall')
    ax2.set_ylabel('Precision')
    ax2.set_title('Precision-Recall Curve')
    ax2.legend(loc="lower left")
    ax2.grid(True, linestyle='--', alpha=0.5)

    plt.tight_layout()
    plt.show()

# Run the complete analysis
run_final_metrics_report(SAVE_DIR)

📥 Loading weights: /content/drive/MyDrive/SatMAE_PartialFT_Results_1/best_model.pth
 Initializing SatMAE Partial FT Model...
 Downloading Pretrained MAE Weights (facebook/vit-mae-base)...
 Weights Loaded!
Applying Partial Freeze Strategy...
 Unfreezing Last 2 Encoder Blocks...
📊 Evaluating Validation Set...

============================================================
      SATMAE FULL FINETUNE - IN-DEPTH METRICS REPORT       
============================================================
Metric                       | Value     
---------------------------------------------
Pixel Accuracy               | 0.9143
IoU (Jaccard Index)          | 0.8858
F1-Score (Dice)              | 0.9394
F2-Score (Recall Focus)      | 0.9387
MCC (Matthews Corr)          | 0.7927
ROC-AUC                      | 0.9637
PR-AUC                       | 0.9830
---------------------------------------------
Precision                    | 0.9406
Recall (Sensitivity)         | 0.9383
Specificity                  | 0.8558
False Positive Rate          | 0.1442
---------------------------------------------
Boundary IoU                 | 0.1070
Hausdorff Distance (px)      | 12.86
Inference Speed              | 2.69 FPS
============================================================