In [ ]:
import os
import time
import json
import numpy as np
import cv2
import rasterio
from rasterio.windows import from_bounds
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import TensorDataset, DataLoader, random_split
from scipy.ndimage import distance_transform_edt as distance
from scipy.spatial.distance import directed_hausdorff
from huggingface_hub import hf_hub_download
import ee
import geemap
from google.colab import drive
from sklearn.metrics import accuracy_score, f1_score, jaccard_score, precision_score, recall_score
from datetime import datetime
import matplotlib.pyplot as plt
# Mount Drive
drive.mount('/content/drive', force_remount=True)
# Initialize Earth Engine
try:
ee.Initialize(project='[REDACTED_FOR_SECURITY]')
except:
ee.Authenticate()
ee.Initialize(project='[REDACTED_FOR_SECURITY]')
# Configuration
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Running on: {device}")
SAVE_DIR = '/content/drive/MyDrive/SatMAE_Fair_Comparison/'
if not os.path.exists(SAVE_DIR): os.makedirs(SAVE_DIR)
BATCH_SIZE = 8
EPOCHS = 50
LR_MAX = 1e-4
LR_MIN = 1e-6
PATCH_SIZE = 224
ASSET_ID = 'projects/[REDACTED_FOR_SECURITY]/assets/Punjab_Mask_2024_NEW'
TIME_WINDOWS = [('2024-11-01', '2024-11-30'), ('2025-02-15', '2025-03-15'), ('2025-04-01', '2025-04-15')]
Mounted at /content/drive Running on: cuda
In [ ]:
!pip install geedim
Collecting geedim Downloading geedim-2.0.0-py3-none-any.whl.metadata (6.0 kB) Requirement already satisfied: numpy>=1.19 in /usr/local/lib/python3.12/dist-packages (from geedim) (2.0.2) Requirement already satisfied: rasterio>=1.3.8 in /usr/local/lib/python3.12/dist-packages (from geedim) (1.5.0) Requirement already satisfied: click>=8 in /usr/local/lib/python3.12/dist-packages (from geedim) (8.3.1) Requirement already satisfied: tqdm>=4.6 in /usr/local/lib/python3.12/dist-packages (from geedim) (4.67.1) Requirement already satisfied: earthengine-api>=0.1.379 in /usr/local/lib/python3.12/dist-packages (from geedim) (1.5.24) Requirement already satisfied: tabulate>=0.9 in /usr/local/lib/python3.12/dist-packages (from geedim) (0.9.0) Requirement already satisfied: fsspec>=2025.2 in /usr/local/lib/python3.12/dist-packages (from geedim) (2025.3.0) Requirement already satisfied: aiohttp>=3.11 in /usr/local/lib/python3.12/dist-packages (from geedim) (3.13.3) Requirement already satisfied: aiohappyeyeballs>=2.5.0 in /usr/local/lib/python3.12/dist-packages (from aiohttp>=3.11->geedim) (2.6.1) Requirement already satisfied: aiosignal>=1.4.0 in /usr/local/lib/python3.12/dist-packages (from aiohttp>=3.11->geedim) (1.4.0) Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.12/dist-packages (from aiohttp>=3.11->geedim) (25.4.0) Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.12/dist-packages (from aiohttp>=3.11->geedim) (1.8.0) Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.12/dist-packages (from aiohttp>=3.11->geedim) (6.7.0) Requirement already satisfied: propcache>=0.2.0 in /usr/local/lib/python3.12/dist-packages (from aiohttp>=3.11->geedim) (0.4.1) Requirement already satisfied: yarl<2.0,>=1.17.0 in /usr/local/lib/python3.12/dist-packages (from aiohttp>=3.11->geedim) (1.22.0) Requirement already satisfied: google-cloud-storage in /usr/local/lib/python3.12/dist-packages (from earthengine-api>=0.1.379->geedim) (3.7.0) Requirement already satisfied: google-api-python-client>=1.12.1 in /usr/local/lib/python3.12/dist-packages (from earthengine-api>=0.1.379->geedim) (2.187.0) Requirement already satisfied: google-auth>=1.4.1 in /usr/local/lib/python3.12/dist-packages (from earthengine-api>=0.1.379->geedim) (2.43.0) Requirement already satisfied: google-auth-httplib2>=0.0.3 in /usr/local/lib/python3.12/dist-packages (from earthengine-api>=0.1.379->geedim) (0.3.0) Requirement already satisfied: httplib2<1dev,>=0.9.2 in /usr/local/lib/python3.12/dist-packages (from earthengine-api>=0.1.379->geedim) (0.31.0) Requirement already satisfied: requests in /usr/local/lib/python3.12/dist-packages (from earthengine-api>=0.1.379->geedim) (2.32.4) Requirement already satisfied: affine in /usr/local/lib/python3.12/dist-packages (from rasterio>=1.3.8->geedim) (2.4.0) Requirement already satisfied: certifi in /usr/local/lib/python3.12/dist-packages (from rasterio>=1.3.8->geedim) (2026.1.4) Requirement already satisfied: cligj>=0.5 in /usr/local/lib/python3.12/dist-packages (from rasterio>=1.3.8->geedim) (0.7.2) Requirement already satisfied: pyparsing in /usr/local/lib/python3.12/dist-packages (from rasterio>=1.3.8->geedim) (3.3.1) Requirement already satisfied: typing-extensions>=4.2 in /usr/local/lib/python3.12/dist-packages (from aiosignal>=1.4.0->aiohttp>=3.11->geedim) (4.15.0) Requirement already satisfied: google-api-core!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.0,<3.0.0,>=1.31.5 in /usr/local/lib/python3.12/dist-packages (from google-api-python-client>=1.12.1->earthengine-api>=0.1.379->geedim) (2.29.0) Requirement already satisfied: uritemplate<5,>=3.0.1 in /usr/local/lib/python3.12/dist-packages (from google-api-python-client>=1.12.1->earthengine-api>=0.1.379->geedim) (4.2.0) Requirement already satisfied: cachetools<7.0,>=2.0.0 in /usr/local/lib/python3.12/dist-packages (from google-auth>=1.4.1->earthengine-api>=0.1.379->geedim) (6.2.4) Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.12/dist-packages (from google-auth>=1.4.1->earthengine-api>=0.1.379->geedim) (0.4.2) Requirement already satisfied: rsa<5,>=3.1.4 in /usr/local/lib/python3.12/dist-packages (from google-auth>=1.4.1->earthengine-api>=0.1.379->geedim) (4.9.1) Requirement already satisfied: idna>=2.0 in /usr/local/lib/python3.12/dist-packages (from yarl<2.0,>=1.17.0->aiohttp>=3.11->geedim) (3.11) Requirement already satisfied: google-cloud-core<3.0.0,>=2.4.2 in /usr/local/lib/python3.12/dist-packages (from google-cloud-storage->earthengine-api>=0.1.379->geedim) (2.5.0) Requirement already satisfied: google-resumable-media<3.0.0,>=2.7.2 in /usr/local/lib/python3.12/dist-packages (from google-cloud-storage->earthengine-api>=0.1.379->geedim) (2.8.0) Requirement already satisfied: google-crc32c<2.0.0,>=1.1.3 in /usr/local/lib/python3.12/dist-packages (from google-cloud-storage->earthengine-api>=0.1.379->geedim) (1.8.0) Requirement already satisfied: charset_normalizer<4,>=2 in /usr/local/lib/python3.12/dist-packages (from requests->earthengine-api>=0.1.379->geedim) (3.4.4) Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.12/dist-packages (from requests->earthengine-api>=0.1.379->geedim) (2.5.0) Requirement already satisfied: googleapis-common-protos<2.0.0,>=1.56.2 in /usr/local/lib/python3.12/dist-packages (from google-api-core!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.0,<3.0.0,>=1.31.5->google-api-python-client>=1.12.1->earthengine-api>=0.1.379->geedim) (1.72.0) Requirement already satisfied: protobuf!=3.20.0,!=3.20.1,!=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5,<7.0.0,>=3.19.5 in /usr/local/lib/python3.12/dist-packages (from google-api-core!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.0,<3.0.0,>=1.31.5->google-api-python-client>=1.12.1->earthengine-api>=0.1.379->geedim) (5.29.5) Requirement already satisfied: proto-plus<2.0.0,>=1.22.3 in /usr/local/lib/python3.12/dist-packages (from google-api-core!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.0,<3.0.0,>=1.31.5->google-api-python-client>=1.12.1->earthengine-api>=0.1.379->geedim) (1.27.0) Requirement already satisfied: pyasn1<0.7.0,>=0.6.1 in /usr/local/lib/python3.12/dist-packages (from pyasn1-modules>=0.2.1->google-auth>=1.4.1->earthengine-api>=0.1.379->geedim) (0.6.1) Downloading geedim-2.0.0-py3-none-any.whl (73 kB) ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 73.1/73.1 kB 4.0 MB/s eta 0:00:00 Installing collected packages: geedim Successfully installed geedim-2.0.0
In [ ]:
def get_satmae_data():
print("Ingesting Asset...")
mask_img = ee.Image(ASSET_ID)
roi_geom = mask_img.geometry()
mask_file = 'local_mask.tif'
if not os.path.exists(mask_file):
geemap.download_ee_image(mask_img, mask_file, region=roi_geom, scale=10, crs='EPSG:4326', overwrite=True)
with rasterio.open(mask_file) as src:
b = src.bounds
cx, cy = (b.left + b.right)/2, (b.bottom + b.top)/2
offset = 0.06
window = from_bounds(cx-offset, cy-offset, cx+offset, cy+offset, src.transform)
mask = src.read(1, window=window)
mask = np.where(mask > 0, 1.0, 0.0).astype(np.float32)
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("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("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()
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:googleapiclient.http:Sleeping 1.65 seconds before retry 1 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/satmae-2026/thumbnails?fields=name&alt=json, after 429 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:googleapiclient.http:Sleeping 0.52 seconds before retry 1 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/satmae-2026/thumbnails?fields=name&alt=json, after 429 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:googleapiclient.http:Sleeping 0.15 seconds before retry 1 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/satmae-2026/thumbnails?fields=name&alt=json, after 429 WARNING:googleapiclient.http:Sleeping 1.61 seconds before retry 1 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/satmae-2026/thumbnails?fields=name&alt=json, after 429 WARNING:googleapiclient.http:Sleeping 1.25 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 /usr/local/lib/python3.12/dist-packages/geedim/image.py:254: RuntimeWarning: Couldn't find STAC entry for: 'projects/satmae-2026/assets/Punjab_Mask_2024_NEW'. return STACClient().get(self.id)
Stacking Time Steps...
0%| |0/48 tiles [00:00<?]
WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 /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<?]
WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: 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.60 seconds before retry 1 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/satmae-2026/thumbnails?fields=name&alt=json, after 429
Creating Patches... Data Ready. Shape: (25, 8, 3, 224, 224)
In [ ]:
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)
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):
with torch.no_grad():
gt_np = gt.cpu().numpy()
dist_map = np.zeros_like(gt_np)
for i in range(len(gt_np)):
# Ensure binary mask for distance transform
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)
probs = torch.sigmoid(pred)
return torch.mean((probs - gt) ** 2 * (1 + self.alpha * torch.abs(dist_map)))
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):
# Weighted combination: 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)
In [ ]:
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 SatMAESegmentation(nn.Module):
def __init__(self, num_frames=3, embed_dim=768):
super().__init__()
print("Constructing Manual SatMAE (Code 1 Base)...")
self.backbone = SatMAEBackbone(num_frames=num_frames, embed_dim=embed_dim)
# Load Weights
try:
print("Loading Google ViT Weights...")
p = hf_hub_download("google/vit-base-patch16-224", "pytorch_model.bin")
sd = torch.load(p, map_location='cpu')
w = sd['vit.embeddings.patch_embeddings.projection.weight']
new_w = torch.zeros(768, 8, 16, 16)
new_w[:, :3] = w
new_w[:, 3:] = w.mean(1, keepdim=True).repeat(1, 5, 1, 1)
self.backbone.patch_embed.proj.weight.data = new_w
self.backbone.patch_embed.proj.bias.data = sd['vit.embeddings.patch_embeddings.projection.bias']
print("Weights Adapted.")
except:
print("Weights missing, using random init.")
# Freezing Strategy (Replicating Code 1 Logic)
print("Freezing Transformer Backbone...")
# 1. Freeze EVERYTHING initially
for param in self.backbone.parameters():
param.requires_grad = False
# 2. Unfreeze only Input and Time embeddings
self.backbone.time_embed.requires_grad = True
self.backbone.patch_embed.proj.weight.requires_grad = True
print("Unfrozen: Patch Embeddings + Time Embeddings")
# 3. 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, 3, 14, 14, D).permute(0, 4, 1, 2, 3).flatten(1, 2)
features = self.temporal_agg(features)
return self.decoder(features)
In [ ]:
model = SatMAESegmentation().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 Fair Training ({EPOCHS} Epochs)...")
history = {'train_loss': [], 'val_loss': []}
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()
scheduler.step()
model.eval()
val_loss = 0
with torch.no_grad():
for x, y in val_loader:
x, y = x.to(device), y.to(device)
preds = model(x)
val_loss += criterion(preds, y).item()
avg_t = train_loss / len(train_loader)
avg_v = val_loss / len(val_loader)
history['train_loss'].append(avg_t)
history['val_loss'].append(avg_v)
if (ep+1)%5==0:
print(f"Ep {ep+1} | Train: {avg_t:.4f} | Val: {avg_v:.4f}")
torch.save(model.state_dict(), SAVE_DIR + "SatMAE_Code1_Fair.pth")
Constructing Manual SatMAE (Code 1 Base)... Loading Google ViT Weights...
/usr/local/lib/python3.12/dist-packages/torch/nn/modules/transformer.py:392: UserWarning: enable_nested_tensor is True, but self.use_nested_tensor is False because encoder_layer.norm_first was True warnings.warn( /usr/local/lib/python3.12/dist-packages/huggingface_hub/utils/_auth.py:94: UserWarning: The secret `HF_TOKEN` does not exist in your Colab secrets. To authenticate with the Hugging Face Hub, create a token in your settings tab (https://huggingface.co/settings/tokens), set it as secret in your Google Colab and restart your session. You will be able to reuse this secret in all of your notebooks. Please note that authentication is recommended but still optional to access public models or datasets. warnings.warn(
pytorch_model.bin: 0%| | 0.00/346M [00:00<?, ?B/s]
Weights Adapted. Freezing Transformer Backbone... Unfrozen: Patch Embeddings + Time Embeddings Starting Fair Training (50 Epochs)... Ep 5 | Train: 0.4532 | Val: 0.7533 Ep 10 | Train: 0.4165 | Val: 0.4827 Ep 15 | Train: 0.3638 | Val: 0.5578 Ep 20 | Train: 0.3387 | Val: 0.5015 Ep 25 | Train: 0.3304 | Val: 0.2919 Ep 30 | Train: 0.3348 | Val: 0.3181 Ep 35 | Train: 0.3126 | Val: 0.3009 Ep 40 | Train: 0.3016 | Val: 0.2724 Ep 45 | Train: 0.2886 | Val: 0.2965 Ep 50 | Train: 0.2855 | Val: 0.2557
In [ ]:
def calculate_boundary_iou(gt_mask, pred_mask, dilation=5):
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()
return inter / (union + 1e-6)
def symmetric_hausdorff(mask_pred, mask_gt):
coords_pred = np.argwhere(mask_pred)
coords_gt = np.argwhere(mask_gt)
if len(coords_pred) == 0 or len(coords_gt) == 0: return 316.0
d_pg = directed_hausdorff(coords_pred, coords_gt)[0]
d_gp = directed_hausdorff(coords_gt, coords_pred)[0]
return max(d_pg, d_gp)
print("Running Final Fair Metrics...")
model.eval()
all_preds, all_targets = [], []
boundary_ious, hausdorff_dists = [], []
start_time = time.time()
with torch.no_grad():
for x, y in val_loader:
x = x.to(device)
logits = model(x)
p_batch = (torch.sigmoid(logits) > 0.5).float().cpu().numpy()
y_batch = y.numpy()
for j in range(len(y_batch)):
p, t = p_batch[j, 0].astype(np.uint8), y_batch[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.flatten())
total_time = time.time() - start_time
fps = len(val_loader.dataset) / (total_time + 1e-6)
y_p, y_t = np.array(all_preds).astype(int), np.array(all_targets).astype(int)
metrics = {
"Model": "SatMAE_Code1_FrozenBrain_Fair",
"Standard_IoU": round(jaccard_score(y_t, y_p, average='binary'), 4),
"Boundary_IoU": round(np.mean(boundary_ious), 4),
"Hausdorff_Dist": round(np.mean(hausdorff_dists), 2),
"FPS": round(fps, 2)
}
print("FAIR COMPARISON RESULT (Manual Code 1 Base):")
print(json.dumps(metrics, indent=4))
Running Final Fair Metrics...
FAIR COMPARISON RESULT (Manual Code 1 Base):
{
"Model": "SatMAE_Code1_FrozenBrain_Fair",
"Standard_IoU": 0.8464,
"Boundary_IoU": 0.3481,
"Hausdorff_Dist": 18.72,
"FPS": 6.46
}
In [ ]:
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import accuracy_score, f1_score, jaccard_score, precision_score, recall_score, confusion_matrix
import pandas as pd
def visualize_and_report(model, loader, history, save_dir, model_name="SatMAE_Code1_Fair"):
print(f" Generating Visual Report for {model_name}...")
model.eval()
# 1. LOSS CURVE (Train vs Validation)
if 'train_loss' in history and 'val_loss' in history:
plt.figure(figsize=(10, 6))
epochs = range(1, len(history['train_loss']) + 1)
plt.plot(epochs, history['train_loss'], 'b-', label='Training Loss', linewidth=2)
plt.plot(epochs, history['val_loss'], 'r--', label='Validation Loss', linewidth=2)
plt.title(f'{model_name} Learning Curve')
plt.xlabel('Epochs')
plt.ylabel('Compound Loss')
plt.legend()
plt.grid(True, alpha=0.3)
plt.savefig(f"{save_dir}{model_name}_LossCurve.png")
plt.show()
print(" Loss Graph Saved.")
# 2. SAMPLE VISUALIZATION (Input vs Truth vs Pred)
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 3 samples
fig, axes = plt.subplots(3, 3, figsize=(12, 12))
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(3):
if i >= len(x_batch): break
# Construct RGB from bands B4(2), B3(1), B2(0) of 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 3.5x
gt_img = y_batch[i, 0].cpu().numpy()
pred_img = preds[i, 0].cpu().numpy()
axes[i, 0].imshow(rgb)
axes[i, 1].imshow(gt_img, cmap='gray')
axes[i, 2].imshow(pred_img, cmap='gray')
for ax in axes[i]: ax.axis('off')
plt.tight_layout()
plt.savefig(f"{save_dir}{model_name}_Samples.png")
plt.show()
print(" Sample Predictions Saved.")
except Exception as e:
print(f" Visualization Error: {e}")
# -------------------------------------------------------
# 3. STANDARD METRICS (Accuracy, F1, Recall, etc.)
# -------------------------------------------------------
print(" Calculating Pixel-wise Metrics...")
all_preds, all_targets = [], []
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().flatten()
y_batch = y.numpy().flatten()
all_preds.extend(p_batch)
all_targets.extend(y_batch)
y_p = np.array(all_preds).astype(int)
y_t = np.array(all_targets).astype(int)
metrics = {
"Model": model_name,
"Pixel_Accuracy": round(accuracy_score(y_t, y_p), 4),
"IoU_Score": round(jaccard_score(y_t, y_p, average='binary'), 4),
"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)
}
# Save to JSON
with open(f"{save_dir}{model_name}_FullMetrics.json", 'w') as f:
json.dump(metrics, f, indent=4)
print("\n FINAL STANDARD METRICS:")
print(json.dumps(metrics, indent=4))
# EXECUTE
visualize_and_report(model, val_loader, history, SAVE_DIR)
Generating Visual Report for SatMAE_Code1_Fair...
Loss Graph Saved.
Sample Predictions Saved.
⏳ Calculating Pixel-wise Metrics...
FINAL STANDARD METRICS:
{
"Model": "SatMAE_Code1_Fair",
"Pixel_Accuracy": 0.8757,
"IoU_Score": 0.8464,
"F1_Score": 0.9168,
"Precision": 0.8918,
"Recall": 0.9433
}
In [ ]:
import os
import time
import json
import numpy as np
import cv2
import rasterio
from rasterio.windows import from_bounds
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import TensorDataset, DataLoader, random_split
from scipy.ndimage import distance_transform_edt as distance
from scipy.spatial.distance import directed_hausdorff
from huggingface_hub import hf_hub_download
import ee
import geemap
from google.colab import drive
from sklearn.metrics import accuracy_score, f1_score, jaccard_score, precision_score, recall_score
from datetime import datetime
import matplotlib.pyplot as plt
# 1. SETUP
drive.mount('/content/drive', force_remount=True)
try:
ee.Initialize(project='[REDACTED_FOR_SECURITY]')
except:
ee.Authenticate()
ee.Initialize(project='[REDACTED_FOR_SECURITY]')
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f" Benchmarking on: {device}")
# SAVE DIRECTORY
SAVE_DIR = '/content/drive/MyDrive/SatMAE_LongTrain_500/'
if not os.path.exists(SAVE_DIR): os.makedirs(SAVE_DIR)
# CONFIG
BATCH_SIZE = 8
EPOCHS = 500 # <--- LONG TRAINING
LR_MAX = 1e-4
LR_MIN = 1e-6
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')]
# 2. 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)
X = np.nan_to_num(X, nan=0.0)
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 🚀 Benchmarking on: cuda 1. Ingesting Asset... 2. Stacking Time Steps... 3. Creating Patches... Data Ready. Shape: (25, 8, 3, 224, 224)
In [ ]:
# 1. ADVANCED 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)
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):
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)
probs = torch.sigmoid(pred)
return torch.mean((probs - gt) ** 2 * (1 + self.alpha * torch.abs(dist_map)))
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):
return 0.5*self.dice(preds, targets) + 0.3*self.boundary(preds, targets) + 0.2*self.focal(preds, targets)
# 2. MODEL DEFINITION
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 SatMAESegmentation(nn.Module):
def __init__(self, num_frames=3, embed_dim=768):
super().__init__()
print(" Constructing Manual SatMAE (Fair Comparison)...")
self.backbone = SatMAEBackbone(num_frames=num_frames, embed_dim=embed_dim)
try:
print(" Loading Google ViT Weights...")
p = hf_hub_download("google/vit-base-patch16-224", "pytorch_model.bin")
sd = torch.load(p, map_location='cpu')
w = sd['vit.embeddings.patch_embeddings.projection.weight']
new_w = torch.zeros(768, 8, 16, 16)
new_w[:, :3] = w
new_w[:, 3:] = w.mean(1, keepdim=True).repeat(1, 5, 1, 1)
self.backbone.patch_embed.proj.weight.data = new_w
self.backbone.patch_embed.proj.bias.data = sd['vit.embeddings.patch_embeddings.projection.bias']
print(" Weights Adapted.")
except:
print(" Weights missing, using random init.")
# FREEZING (Code 1 Logic)
print(" Freezing Transformer Backbone...")
for param in self.backbone.parameters():
param.requires_grad = False
self.backbone.time_embed.requires_grad = True
self.backbone.patch_embed.proj.weight.requires_grad = True
print(" Unfrozen: Patch Embeddings + Time Embeddings")
# 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, 3, 14, 14, D).permute(0, 4, 1, 2, 3).flatten(1, 2)
features = self.temporal_agg(features)
return self.decoder(features)
In [ ]:
# TRAINING LOOP (500 EPOCHS)
model = SatMAESegmentation().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=50, T_mult=2, eta_min=LR_MIN)
# Dataset Split
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)
# CHECKPOINT MANAGER
CHECKPOINT_PATH = SAVE_DIR + "checkpoint.pth"
start_epoch = 0
history = {'train_loss': [], 'val_loss': []}
best_loss = float('inf')
# 1. Try to Resume
if os.path.exists(CHECKPOINT_PATH):
print(f" Checkpoint found! Resuming from Drive...")
checkpoint = torch.load(CHECKPOINT_PATH)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
start_epoch = checkpoint['epoch'] + 1
history = checkpoint['history']
best_loss = checkpoint['best_loss']
print(f" Resuming from Epoch {start_epoch}")
else:
print(" No checkpoint found. Starting from scratch.")
# 2. Training Loop
print(f" Starting/Resuming Training ({EPOCHS} Epochs total)...")
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)
optimizer.zero_grad()
preds = model(x)
loss = criterion(preds, y)
loss.backward()
optimizer.step()
train_loss += loss.item()
scheduler.step()
# Validation
model.eval()
val_loss = 0
with torch.no_grad():
for x, y in val_loader:
x, y = x.to(device), y.to(device)
preds = model(x)
val_loss += criterion(preds, y).item()
avg_t = train_loss / len(train_loader)
avg_v = val_loss / len(val_loader)
history['train_loss'].append(avg_t)
history['val_loss'].append(avg_v)
# Printing
if (ep+1) % 1 == 0:
print(f"Ep {ep+1}/{EPOCHS} | Train: {avg_t:.4f} | Val: {avg_v:.4f}")
# 3. SAVE BEST MODEL
if avg_v < best_loss:
best_loss = avg_v
torch.save(model.state_dict(), SAVE_DIR + "SatMAE_500_BEST.pth")
print(f" Best Model Saved (Loss: {avg_v:.4f})")
# 4. SAVE CHECKPOINT (Every 5 Epochs)
if (ep+1) % 5 == 0:
torch.save({
'epoch': ep,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'scheduler_state_dict': scheduler.state_dict(),
'history': history,
'best_loss': best_loss
}, CHECKPOINT_PATH)
print(" Checkpoint saved.")
print(" Training Complete!")
Constructing Manual SatMAE (Fair Comparison)... Loading Google ViT Weights... Weights Adapted. Freezing Transformer Backbone... Unfrozen: Patch Embeddings + Time Embeddings No checkpoint found. Starting from scratch. Starting/Resuming Training (500 Epochs total)... Ep 1/500 | Train: 1.1114 | Val: 1.1168 Best Model Saved (Loss: 1.1168) Ep 2/500 | Train: 0.8696 | Val: 1.1109 Best Model Saved (Loss: 1.1109) Ep 3/500 | Train: 0.6986 | Val: 1.0248 Best Model Saved (Loss: 1.0248) Ep 4/500 | Train: 0.6150 | Val: 0.9220 Best Model Saved (Loss: 0.9220) Ep 5/500 | Train: 0.5689 | Val: 0.8183 Best Model Saved (Loss: 0.8183) Checkpoint saved. Ep 6/500 | Train: 0.5597 | Val: 0.7745 Best Model Saved (Loss: 0.7745) Ep 7/500 | Train: 0.5206 | Val: 0.7182 Best Model Saved (Loss: 0.7182) Ep 8/500 | Train: 0.5093 | Val: 0.6697 Best Model Saved (Loss: 0.6697) Ep 9/500 | Train: 0.4956 | Val: 0.6222 Best Model Saved (Loss: 0.6222) Ep 10/500 | Train: 0.4891 | Val: 0.5733 Best Model Saved (Loss: 0.5733) Checkpoint saved. Ep 11/500 | Train: 0.4771 | Val: 0.5191 Best Model Saved (Loss: 0.5191) Ep 12/500 | Train: 0.4692 | Val: 0.5000 Best Model Saved (Loss: 0.5000) Ep 13/500 | Train: 0.4654 | Val: 0.4646 Best Model Saved (Loss: 0.4646) Ep 14/500 | Train: 0.4631 | Val: 0.4598 Best Model Saved (Loss: 0.4598) Ep 15/500 | Train: 0.4525 | Val: 0.4977 Checkpoint saved. Ep 16/500 | Train: 0.4547 | Val: 0.4641 Ep 17/500 | Train: 0.4533 | Val: 0.4419 Best Model Saved (Loss: 0.4419) Ep 18/500 | Train: 0.4444 | Val: 0.4938 Ep 19/500 | Train: 0.4441 | Val: 0.4279 Best Model Saved (Loss: 0.4279) Ep 20/500 | Train: 0.4376 | Val: 0.4903 Checkpoint saved. Ep 21/500 | Train: 0.4354 | Val: 0.4694 Ep 22/500 | Train: 0.4496 | Val: 0.4661 Ep 23/500 | Train: 0.4295 | Val: 0.5341 Ep 24/500 | Train: 0.4311 | Val: 0.4292 Ep 25/500 | Train: 0.4285 | Val: 0.4372 Checkpoint saved. Ep 26/500 | Train: 0.4285 | Val: 0.4743 Ep 27/500 | Train: 0.4293 | Val: 0.4789 Ep 28/500 | Train: 0.4258 | Val: 0.4527 Ep 29/500 | Train: 0.4195 | Val: 0.4505 Ep 30/500 | Train: 0.4312 | Val: 0.4238 Best Model Saved (Loss: 0.4238) Checkpoint saved. Ep 31/500 | Train: 0.4147 | Val: 0.4320 Ep 32/500 | Train: 0.4208 | Val: 0.4524 Ep 33/500 | Train: 0.4175 | Val: 0.4600 Ep 34/500 | Train: 0.4128 | Val: 0.4510 Ep 35/500 | Train: 0.4153 | Val: 0.4250 Checkpoint saved. Ep 36/500 | Train: 0.4117 | Val: 0.4201 Best Model Saved (Loss: 0.4201) Ep 37/500 | Train: 0.4137 | Val: 0.4342 Ep 38/500 | Train: 0.4120 | Val: 0.4207 Ep 39/500 | Train: 0.4093 | Val: 0.4157 Best Model Saved (Loss: 0.4157) Ep 40/500 | Train: 0.4136 | Val: 0.4196 Checkpoint saved. Ep 41/500 | Train: 0.4082 | Val: 0.4242 Ep 42/500 | Train: 0.4090 | Val: 0.4214 Ep 43/500 | Train: 0.4183 | Val: 0.4144 Best Model Saved (Loss: 0.4144) Ep 44/500 | Train: 0.4092 | Val: 0.4107 Best Model Saved (Loss: 0.4107) Ep 45/500 | Train: 0.4095 | Val: 0.4126 Checkpoint saved. Ep 46/500 | Train: 0.4079 | Val: 0.4152 Ep 47/500 | Train: 0.4053 | Val: 0.4157 Ep 48/500 | Train: 0.4059 | Val: 0.4175 Ep 49/500 | Train: 0.4074 | Val: 0.4183 Ep 50/500 | Train: 0.4061 | Val: 0.4177 Checkpoint saved. Ep 51/500 | Train: 0.4084 | Val: 0.4795 Ep 52/500 | Train: 0.4041 | Val: 0.3877 Best Model Saved (Loss: 0.3877) Ep 53/500 | Train: 0.4012 | Val: 0.3787 Best Model Saved (Loss: 0.3787) Ep 54/500 | Train: 0.4077 | Val: 0.4620 Ep 55/500 | Train: 0.3927 | Val: 0.3972 Checkpoint saved. Ep 56/500 | Train: 0.4039 | Val: 0.5527 Ep 57/500 | Train: 0.3897 | Val: 0.3761 Best Model Saved (Loss: 0.3761) Ep 58/500 | Train: 0.3915 | Val: 0.5613 Ep 59/500 | Train: 0.3899 | Val: 0.4326 Ep 60/500 | Train: 0.3834 | Val: 0.5095 Checkpoint saved. Ep 61/500 | Train: 0.3803 | Val: 0.3627 Best Model Saved (Loss: 0.3627) Ep 62/500 | Train: 0.3775 | Val: 0.4258 Ep 63/500 | Train: 0.3723 | Val: 0.3942 Ep 64/500 | Train: 0.3729 | Val: 0.4865 Ep 65/500 | Train: 0.3745 | Val: 0.4007 Checkpoint saved. Ep 66/500 | Train: 0.3656 | Val: 0.4485 Ep 67/500 | Train: 0.3655 | Val: 0.3595 Best Model Saved (Loss: 0.3595) Ep 68/500 | Train: 0.3636 | Val: 0.4382 Ep 69/500 | Train: 0.3637 | Val: 0.4754 Ep 70/500 | Train: 0.3640 | Val: 0.4012 Checkpoint saved. Ep 71/500 | Train: 0.3646 | Val: 0.4743 Ep 72/500 | Train: 0.3541 | Val: 0.3739 Ep 73/500 | Train: 0.3535 | Val: 0.4019 Ep 74/500 | Train: 0.3487 | Val: 0.3665 Ep 75/500 | Train: 0.3577 | Val: 0.3559 Best Model Saved (Loss: 0.3559) Checkpoint saved. Ep 76/500 | Train: 0.3438 | Val: 0.4268 Ep 77/500 | Train: 0.3441 | Val: 0.4200 Ep 78/500 | Train: 0.3387 | Val: 0.3558 Best Model Saved (Loss: 0.3558) Ep 79/500 | Train: 0.3376 | Val: 0.4020 Ep 80/500 | Train: 0.3389 | Val: 0.3576 Checkpoint saved. Ep 81/500 | Train: 0.3413 | Val: 0.3501 Best Model Saved (Loss: 0.3501) Ep 82/500 | Train: 0.3420 | Val: 0.3802 Ep 83/500 | Train: 0.3352 | Val: 0.4098 Ep 84/500 | Train: 0.3353 | Val: 0.4366 Ep 85/500 | Train: 0.3295 | Val: 0.3690 Checkpoint saved. Ep 86/500 | Train: 0.3377 | Val: 0.3950 Ep 87/500 | Train: 0.3259 | Val: 0.4308 Ep 88/500 | Train: 0.3293 | Val: 0.3573 Ep 89/500 | Train: 0.3259 | Val: 0.3768 Ep 90/500 | Train: 0.3260 | Val: 0.3433 Best Model Saved (Loss: 0.3433) Checkpoint saved. Ep 91/500 | Train: 0.3191 | Val: 0.3340 Best Model Saved (Loss: 0.3340) Ep 92/500 | Train: 0.3202 | Val: 0.3521 Ep 93/500 | Train: 0.3288 | Val: 0.4079 Ep 94/500 | Train: 0.3173 | Val: 0.3358 Ep 95/500 | Train: 0.3168 | Val: 0.3593 Checkpoint saved. Ep 96/500 | Train: 0.3191 | Val: 0.4058 Ep 97/500 | Train: 0.3220 | Val: 0.3680 Ep 98/500 | Train: 0.3217 | Val: 0.3408 Ep 99/500 | Train: 0.3131 | Val: 0.3346 Ep 100/500 | Train: 0.3103 | Val: 0.3887 Checkpoint saved. Ep 101/500 | Train: 0.3078 | Val: 0.3531 Ep 102/500 | Train: 0.3079 | Val: 0.3657 Ep 103/500 | Train: 0.3101 | Val: 0.3541 Ep 104/500 | Train: 0.3063 | Val: 0.3655 Ep 105/500 | Train: 0.3048 | Val: 0.3544 Checkpoint saved. Ep 106/500 | Train: 0.3052 | Val: 0.3589 Ep 107/500 | Train: 0.3015 | Val: 0.3447 Ep 108/500 | Train: 0.3148 | Val: 0.3484 Ep 109/500 | Train: 0.3075 | Val: 0.3404 Ep 110/500 | Train: 0.3183 | Val: 0.3641 Checkpoint saved. Ep 111/500 | Train: 0.3001 | Val: 0.3567 Ep 112/500 | Train: 0.2987 | Val: 0.3647 Ep 113/500 | Train: 0.3124 | Val: 0.3665 Ep 114/500 | Train: 0.3009 | Val: 0.3705 Ep 115/500 | Train: 0.2965 | Val: 0.3556 Checkpoint saved. Ep 116/500 | Train: 0.2965 | Val: 0.3529 Ep 117/500 | Train: 0.2961 | Val: 0.3610 Ep 118/500 | Train: 0.3023 | Val: 0.3618 Ep 119/500 | Train: 0.2954 | Val: 0.3724 Ep 120/500 | Train: 0.3081 | Val: 0.3687 Checkpoint saved. Ep 121/500 | Train: 0.3000 | Val: 0.3595 Ep 122/500 | Train: 0.2936 | Val: 0.3521 Ep 123/500 | Train: 0.2953 | Val: 0.3529 Ep 124/500 | Train: 0.2979 | Val: 0.3469 Ep 125/500 | Train: 0.2946 | Val: 0.3406 Checkpoint saved. Ep 126/500 | Train: 0.3004 | Val: 0.3479 Ep 127/500 | Train: 0.2920 | Val: 0.3631 Ep 128/500 | Train: 0.2985 | Val: 0.3530 Ep 129/500 | Train: 0.2932 | Val: 0.3477 Ep 130/500 | Train: 0.2913 | Val: 0.3561 Checkpoint saved. Ep 131/500 | Train: 0.2952 | Val: 0.3579 Ep 132/500 | Train: 0.2944 | Val: 0.3448 Ep 133/500 | Train: 0.2901 | Val: 0.3402 Ep 134/500 | Train: 0.3030 | Val: 0.3467 Ep 135/500 | Train: 0.2907 | Val: 0.3468 Checkpoint saved. Ep 136/500 | Train: 0.2896 | Val: 0.3475 Ep 137/500 | Train: 0.2921 | Val: 0.3487 Ep 138/500 | Train: 0.2892 | Val: 0.3456 Ep 139/500 | Train: 0.2903 | Val: 0.3441 Ep 140/500 | Train: 0.2936 | Val: 0.3443 Checkpoint saved. Ep 141/500 | Train: 0.2895 | Val: 0.3418 Ep 142/500 | Train: 0.2906 | Val: 0.3434 Ep 143/500 | Train: 0.2904 | Val: 0.3429 Ep 144/500 | Train: 0.2941 | Val: 0.3406 Ep 145/500 | Train: 0.2891 | Val: 0.3417 Checkpoint saved. Ep 146/500 | Train: 0.2955 | Val: 0.3395 Ep 147/500 | Train: 0.2908 | Val: 0.3390 Ep 148/500 | Train: 0.2905 | Val: 0.3394 Ep 149/500 | Train: 0.2935 | Val: 0.3417 Ep 150/500 | Train: 0.2938 | Val: 0.3431 Checkpoint saved. Ep 151/500 | Train: 0.2925 | Val: 0.4225 Ep 152/500 | Train: 0.2983 | Val: 0.5710 Ep 153/500 | Train: 0.2977 | Val: 0.4190 Ep 154/500 | Train: 0.3000 | Val: 0.5093 Ep 155/500 | Train: 0.2994 | Val: 0.4821 Checkpoint saved. Ep 156/500 | Train: 0.2939 | Val: 0.4223 Ep 157/500 | Train: 0.3051 | Val: 0.7092 Ep 158/500 | Train: 0.3012 | Val: 0.4625 Ep 159/500 | Train: 0.2998 | Val: 0.3637 Ep 160/500 | Train: 0.2903 | Val: 0.3580 Checkpoint saved. Ep 161/500 | Train: 0.2904 | Val: 0.3386 Ep 162/500 | Train: 0.2882 | Val: 0.3722 Ep 163/500 | Train: 0.2793 | Val: 0.3267 Best Model Saved (Loss: 0.3267) Ep 164/500 | Train: 0.2801 | Val: 0.3249 Best Model Saved (Loss: 0.3249) Ep 165/500 | Train: 0.2784 | Val: 0.3775 Checkpoint saved. Ep 166/500 | Train: 0.2778 | Val: 0.3718 Ep 167/500 | Train: 0.2846 | Val: 0.3043 Best Model Saved (Loss: 0.3043) Ep 168/500 | Train: 0.2694 | Val: 0.4308 Ep 169/500 | Train: 0.2696 | Val: 0.3629 Ep 170/500 | Train: 0.2669 | Val: 0.3225 Checkpoint saved. Ep 171/500 | Train: 0.2667 | Val: 0.3415 Ep 172/500 | Train: 0.2622 | Val: 0.3694 Ep 173/500 | Train: 0.2651 | Val: 0.3314 Ep 174/500 | Train: 0.2636 | Val: 0.3272 Ep 175/500 | Train: 0.2633 | Val: 0.3695 Checkpoint saved. Ep 176/500 | Train: 0.2639 | Val: 0.3303 Ep 177/500 | Train: 0.2551 | Val: 0.3609 Ep 178/500 | Train: 0.2586 | Val: 0.3584 Ep 179/500 | Train: 0.2529 | Val: 0.3316 Ep 180/500 | Train: 0.2522 | Val: 0.3269 Checkpoint saved. Ep 181/500 | Train: 0.2498 | Val: 0.3092 Ep 182/500 | Train: 0.2488 | Val: 0.3417 Ep 183/500 | Train: 0.2501 | Val: 0.2986 Best Model Saved (Loss: 0.2986) Ep 184/500 | Train: 0.2498 | Val: 0.3138 Ep 185/500 | Train: 0.2461 | Val: 0.3670 Checkpoint saved. Ep 186/500 | Train: 0.2467 | Val: 0.3077 Ep 187/500 | Train: 0.2443 | Val: 0.3068 Ep 188/500 | Train: 0.2490 | Val: 0.3402 Ep 189/500 | Train: 0.2414 | Val: 0.3612 Ep 190/500 | Train: 0.2422 | Val: 0.2898 Best Model Saved (Loss: 0.2898) Checkpoint saved. Ep 191/500 | Train: 0.2474 | Val: 0.3292 Ep 192/500 | Train: 0.2404 | Val: 0.3307 Ep 193/500 | Train: 0.2393 | Val: 0.3076 Ep 194/500 | Train: 0.2378 | Val: 0.3306 Ep 195/500 | Train: 0.2381 | Val: 0.3348 Checkpoint saved. Ep 196/500 | Train: 0.2336 | Val: 0.2969 Ep 197/500 | Train: 0.2328 | Val: 0.3303 Ep 198/500 | Train: 0.2319 | Val: 0.3225 Ep 199/500 | Train: 0.2328 | Val: 0.3285 Ep 200/500 | Train: 0.2323 | Val: 0.3153 Checkpoint saved. Ep 201/500 | Train: 0.2305 | Val: 0.3323 Ep 202/500 | Train: 0.2265 | Val: 0.3069 Ep 203/500 | Train: 0.2298 | Val: 0.3119 Ep 204/500 | Train: 0.2299 | Val: 0.3350 Ep 205/500 | Train: 0.2236 | Val: 0.2993 Checkpoint saved. Ep 206/500 | Train: 0.2225 | Val: 0.3426 Ep 207/500 | Train: 0.2210 | Val: 0.3180 Ep 208/500 | Train: 0.2266 | Val: 0.3280 Ep 209/500 | Train: 0.2328 | Val: 0.3045 Ep 210/500 | Train: 0.2230 | Val: 0.3468 Checkpoint saved. Ep 211/500 | Train: 0.2185 | Val: 0.3443 Ep 212/500 | Train: 0.2236 | Val: 0.3191 Ep 213/500 | Train: 0.2223 | Val: 0.3124 Ep 214/500 | Train: 0.2163 | Val: 0.3072 Ep 215/500 | Train: 0.2234 | Val: 0.3120 Checkpoint saved. Ep 216/500 | Train: 0.2139 | Val: 0.3480 Ep 217/500 | Train: 0.2185 | Val: 0.3261 Ep 218/500 | Train: 0.2129 | Val: 0.2990 Ep 219/500 | Train: 0.2155 | Val: 0.3008 Ep 220/500 | Train: 0.2129 | Val: 0.3193 Checkpoint saved. Ep 221/500 | Train: 0.2107 | Val: 0.3052 Ep 222/500 | Train: 0.2114 | Val: 0.3280 Ep 223/500 | Train: 0.2090 | Val: 0.3011 Ep 224/500 | Train: 0.2135 | Val: 0.3166 Ep 225/500 | Train: 0.2125 | Val: 0.2965 Checkpoint saved. Ep 226/500 | Train: 0.2086 | Val: 0.2948 Ep 227/500 | Train: 0.2178 | Val: 0.3173 Ep 228/500 | Train: 0.2103 | Val: 0.3073 Ep 229/500 | Train: 0.2037 | Val: 0.3181 Ep 230/500 | Train: 0.2101 | Val: 0.2933 Checkpoint saved. Ep 231/500 | Train: 0.2050 | Val: 0.3185 Ep 232/500 | Train: 0.2017 | Val: 0.3234 Ep 233/500 | Train: 0.2026 | Val: 0.2995 Ep 234/500 | Train: 0.2012 | Val: 0.2922 Ep 235/500 | Train: 0.2099 | Val: 0.3007 Checkpoint saved. Ep 236/500 | Train: 0.2077 | Val: 0.3018 Ep 237/500 | Train: 0.1993 | Val: 0.3036 Ep 238/500 | Train: 0.2015 | Val: 0.3091 Ep 239/500 | Train: 0.2056 | Val: 0.2945 Ep 240/500 | Train: 0.1968 | Val: 0.2991 Checkpoint saved. Ep 241/500 | Train: 0.1958 | Val: 0.2946 Ep 242/500 | Train: 0.2000 | Val: 0.2966 Ep 243/500 | Train: 0.2028 | Val: 0.3035 Ep 244/500 | Train: 0.1944 | Val: 0.2952 Ep 245/500 | Train: 0.1934 | Val: 0.2998 Checkpoint saved. Ep 246/500 | Train: 0.2005 | Val: 0.3116 Ep 247/500 | Train: 0.1943 | Val: 0.3028 Ep 248/500 | Train: 0.1917 | Val: 0.2926 Ep 249/500 | Train: 0.1919 | Val: 0.2993 Ep 250/500 | Train: 0.1964 | Val: 0.3010 Checkpoint saved. Ep 251/500 | Train: 0.1932 | Val: 0.2933 Ep 252/500 | Train: 0.1933 | Val: 0.2936 Ep 253/500 | Train: 0.1972 | Val: 0.3020 Ep 254/500 | Train: 0.1891 | Val: 0.3129 Ep 255/500 | Train: 0.1890 | Val: 0.2939 Checkpoint saved. Ep 256/500 | Train: 0.2050 | Val: 0.2915 Ep 257/500 | Train: 0.1914 | Val: 0.2981 Ep 258/500 | Train: 0.1958 | Val: 0.3151 Ep 259/500 | Train: 0.1913 | Val: 0.2889 Best Model Saved (Loss: 0.2889) Ep 260/500 | Train: 0.1901 | Val: 0.2993 Checkpoint saved. Ep 261/500 | Train: 0.1989 | Val: 0.2971 Ep 262/500 | Train: 0.1867 | Val: 0.2913 Ep 263/500 | Train: 0.1869 | Val: 0.2866 Best Model Saved (Loss: 0.2866) Ep 264/500 | Train: 0.1913 | Val: 0.3032 Ep 265/500 | Train: 0.1889 | Val: 0.2914 Checkpoint saved. Ep 266/500 | Train: 0.1849 | Val: 0.2875 Ep 267/500 | Train: 0.1898 | Val: 0.2952 Ep 268/500 | Train: 0.1840 | Val: 0.2966 Ep 269/500 | Train: 0.1905 | Val: 0.3023 Ep 270/500 | Train: 0.1845 | Val: 0.2958 Checkpoint saved. Ep 271/500 | Train: 0.1868 | Val: 0.2919 Ep 272/500 | Train: 0.1905 | Val: 0.2874 Ep 273/500 | Train: 0.1832 | Val: 0.2984 Ep 274/500 | Train: 0.1836 | Val: 0.2972 Ep 275/500 | Train: 0.1868 | Val: 0.2848 Best Model Saved (Loss: 0.2848) Checkpoint saved. Ep 276/500 | Train: 0.1865 | Val: 0.2960 Ep 277/500 | Train: 0.1818 | Val: 0.2859 Ep 278/500 | Train: 0.1877 | Val: 0.2827 Best Model Saved (Loss: 0.2827) Ep 279/500 | Train: 0.1799 | Val: 0.2836 Ep 280/500 | Train: 0.1796 | Val: 0.2894 Checkpoint saved. Ep 281/500 | Train: 0.1807 | Val: 0.2935 Ep 282/500 | Train: 0.1815 | Val: 0.2890 Ep 283/500 | Train: 0.1929 | Val: 0.2826 Best Model Saved (Loss: 0.2826) Ep 284/500 | Train: 0.1837 | Val: 0.2833 Ep 285/500 | Train: 0.1804 | Val: 0.2901 Checkpoint saved. Ep 286/500 | Train: 0.1794 | Val: 0.2932 Ep 287/500 | Train: 0.1809 | Val: 0.2931 Ep 288/500 | Train: 0.1803 | Val: 0.2926 Ep 289/500 | Train: 0.1812 | Val: 0.2841 Ep 290/500 | Train: 0.1797 | Val: 0.2809 Best Model Saved (Loss: 0.2809) Checkpoint saved. Ep 291/500 | Train: 0.1815 | Val: 0.2899 Ep 292/500 | Train: 0.1794 | Val: 0.2872 Ep 293/500 | Train: 0.1823 | Val: 0.2838 Ep 294/500 | Train: 0.1801 | Val: 0.2924 Ep 295/500 | Train: 0.1794 | Val: 0.2952 Checkpoint saved. Ep 296/500 | Train: 0.1771 | Val: 0.2820 Ep 297/500 | Train: 0.1802 | Val: 0.2826 Ep 298/500 | Train: 0.1792 | Val: 0.2877 Ep 299/500 | Train: 0.1800 | Val: 0.2875 Ep 300/500 | Train: 0.1779 | Val: 0.2866 Checkpoint saved. Ep 301/500 | Train: 0.1756 | Val: 0.2871 Ep 302/500 | Train: 0.1833 | Val: 0.2837 Ep 303/500 | Train: 0.1753 | Val: 0.2852 Ep 304/500 | Train: 0.1801 | Val: 0.2896 Ep 305/500 | Train: 0.1791 | Val: 0.2833 Checkpoint saved. Ep 306/500 | Train: 0.1752 | Val: 0.2805 Best Model Saved (Loss: 0.2805) Ep 307/500 | Train: 0.1752 | Val: 0.2866 Ep 308/500 | Train: 0.1768 | Val: 0.2899 Ep 309/500 | Train: 0.1767 | Val: 0.2872 Ep 310/500 | Train: 0.1814 | Val: 0.2877 Checkpoint saved. Ep 311/500 | Train: 0.1743 | Val: 0.2877 Ep 312/500 | Train: 0.1745 | Val: 0.2885 Ep 313/500 | Train: 0.1753 | Val: 0.2866 Ep 314/500 | Train: 0.1866 | Val: 0.2875 Ep 315/500 | Train: 0.1743 | Val: 0.2871 Checkpoint saved. Ep 316/500 | Train: 0.1811 | Val: 0.2912 Ep 317/500 | Train: 0.1751 | Val: 0.2889 Ep 318/500 | Train: 0.1763 | Val: 0.2836 Ep 319/500 | Train: 0.1821 | Val: 0.2858 Ep 320/500 | Train: 0.1780 | Val: 0.2875 Checkpoint saved. Ep 321/500 | Train: 0.1744 | Val: 0.2852 Ep 322/500 | Train: 0.1737 | Val: 0.2828 Ep 323/500 | Train: 0.1823 | Val: 0.2821 Ep 324/500 | Train: 0.1732 | Val: 0.2831 Ep 325/500 | Train: 0.1739 | Val: 0.2850 Checkpoint saved. Ep 326/500 | Train: 0.1738 | Val: 0.2851 Ep 327/500 | Train: 0.1770 | Val: 0.2850 Ep 328/500 | Train: 0.1729 | Val: 0.2839 Ep 329/500 | Train: 0.1734 | Val: 0.2840 Ep 330/500 | Train: 0.1762 | Val: 0.2842 Checkpoint saved. Ep 331/500 | Train: 0.1770 | Val: 0.2842 Ep 332/500 | Train: 0.1799 | Val: 0.2848 Ep 333/500 | Train: 0.1757 | Val: 0.2856 Ep 334/500 | Train: 0.1757 | Val: 0.2849 Ep 335/500 | Train: 0.1735 | Val: 0.2842 Checkpoint saved. Ep 336/500 | Train: 0.1764 | Val: 0.2854 Ep 337/500 | Train: 0.1746 | Val: 0.2846 Ep 338/500 | Train: 0.1752 | Val: 0.2847 Ep 339/500 | Train: 0.1783 | Val: 0.2840 Ep 340/500 | Train: 0.1760 | Val: 0.2855 Checkpoint saved. Ep 341/500 | Train: 0.1726 | Val: 0.2846 Ep 342/500 | Train: 0.1744 | Val: 0.2843 Ep 343/500 | Train: 0.1728 | Val: 0.2839 Ep 344/500 | Train: 0.1729 | Val: 0.2843 Ep 345/500 | Train: 0.1731 | Val: 0.2833 Checkpoint saved. Ep 346/500 | Train: 0.1736 | Val: 0.2826 Ep 347/500 | Train: 0.1725 | Val: 0.2819 Ep 348/500 | Train: 0.1738 | Val: 0.2817 Ep 349/500 | Train: 0.1726 | Val: 0.2815 Ep 350/500 | Train: 0.1752 | Val: 0.2802 Best Model Saved (Loss: 0.2802) Checkpoint saved. Ep 351/500 | Train: 0.1750 | Val: 0.2952 Ep 352/500 | Train: 0.1748 | Val: 0.3146 Ep 353/500 | Train: 0.1779 | Val: 0.2868 Ep 354/500 | Train: 0.1748 | Val: 0.2773 Best Model Saved (Loss: 0.2773) Ep 355/500 | Train: 0.1739 | Val: 0.3361 Checkpoint saved. Ep 356/500 | Train: 0.1811 | Val: 0.3211 Ep 357/500 | Train: 0.1744 | Val: 0.3458 Ep 358/500 | Train: 0.1732 | Val: 0.3534 Ep 359/500 | Train: 0.1737 | Val: 0.3498 Ep 360/500 | Train: 0.1759 | Val: 0.3540 Checkpoint saved. Ep 361/500 | Train: 0.1746 | Val: 0.3281 Ep 362/500 | Train: 0.1723 | Val: 0.2979 Ep 363/500 | Train: 0.1717 | Val: 0.3168 Ep 364/500 | Train: 0.1812 | Val: 0.2975 Ep 365/500 | Train: 0.1732 | Val: 0.3157 Checkpoint saved. Ep 366/500 | Train: 0.1757 | Val: 0.3168 Ep 367/500 | Train: 0.1744 | Val: 0.2959 Ep 368/500 | Train: 0.1727 | Val: 0.2839 Ep 369/500 | Train: 0.1844 | Val: 0.5944 Ep 370/500 | Train: 0.1954 | Val: 0.3088 Checkpoint saved. Ep 371/500 | Train: 0.1828 | Val: 0.2854 Ep 372/500 | Train: 0.1807 | Val: 0.3444 Ep 373/500 | Train: 0.1735 | Val: 0.3101 Ep 374/500 | Train: 0.1686 | Val: 0.3083 Ep 375/500 | Train: 0.1698 | Val: 0.2804 Checkpoint saved. Ep 376/500 | Train: 0.1676 | Val: 0.3046 Ep 377/500 | Train: 0.1633 | Val: 0.2738 Best Model Saved (Loss: 0.2738) Ep 378/500 | Train: 0.1614 | Val: 0.2691 Best Model Saved (Loss: 0.2691) Ep 379/500 | Train: 0.1636 | Val: 0.2825 Ep 380/500 | Train: 0.1589 | Val: 0.2725 Checkpoint saved. Ep 381/500 | Train: 0.1577 | Val: 0.3138 Ep 382/500 | Train: 0.1590 | Val: 0.2730 Ep 383/500 | Train: 0.1553 | Val: 0.2949 Ep 384/500 | Train: 0.1558 | Val: 0.2775 Ep 385/500 | Train: 0.1612 | Val: 0.2874 Checkpoint saved. Ep 386/500 | Train: 0.1538 | Val: 0.2974 Ep 387/500 | Train: 0.1559 | Val: 0.2758 Ep 388/500 | Train: 0.1535 | Val: 0.2827 Ep 389/500 | Train: 0.1530 | Val: 0.2875 Ep 390/500 | Train: 0.1508 | Val: 0.2762 Checkpoint saved. Ep 391/500 | Train: 0.1489 | Val: 0.2879 Ep 392/500 | Train: 0.1478 | Val: 0.2774 Ep 393/500 | Train: 0.1468 | Val: 0.2664 Best Model Saved (Loss: 0.2664) Ep 394/500 | Train: 0.1466 | Val: 0.2755 Ep 395/500 | Train: 0.1452 | Val: 0.2723 Checkpoint saved. Ep 396/500 | Train: 0.1444 | Val: 0.2684 Ep 397/500 | Train: 0.1446 | Val: 0.2791 Ep 398/500 | Train: 0.1470 | Val: 0.2735 Ep 399/500 | Train: 0.1482 | Val: 0.2794 Ep 400/500 | Train: 0.1433 | Val: 0.2800 Checkpoint saved. Ep 401/500 | Train: 0.1423 | Val: 0.2628 Best Model Saved (Loss: 0.2628) Ep 402/500 | Train: 0.1423 | Val: 0.2746 Ep 403/500 | Train: 0.1426 | Val: 0.2735 Ep 404/500 | Train: 0.1406 | Val: 0.2689 Ep 405/500 | Train: 0.1433 | Val: 0.2813 Checkpoint saved. Ep 406/500 | Train: 0.1430 | Val: 0.2679 Ep 407/500 | Train: 0.1457 | Val: 0.2743 Ep 408/500 | Train: 0.1458 | Val: 0.2878 Ep 409/500 | Train: 0.1414 | Val: 0.2888 Ep 410/500 | Train: 0.1390 | Val: 0.2718 Checkpoint saved. Ep 411/500 | Train: 0.1391 | Val: 0.2781 Ep 412/500 | Train: 0.1376 | Val: 0.2500 Best Model Saved (Loss: 0.2500) Ep 413/500 | Train: 0.1375 | Val: 0.2754 Ep 414/500 | Train: 0.1358 | Val: 0.2696 Ep 415/500 | Train: 0.1358 | Val: 0.2721 Checkpoint saved. Ep 416/500 | Train: 0.1351 | Val: 0.2665 Ep 417/500 | Train: 0.1361 | Val: 0.2781 Ep 418/500 | Train: 0.1352 | Val: 0.2601 Ep 419/500 | Train: 0.1380 | Val: 0.2809 Ep 420/500 | Train: 0.1336 | Val: 0.2686 Checkpoint saved. Ep 421/500 | Train: 0.1316 | Val: 0.2631 Ep 422/500 | Train: 0.1392 | Val: 0.2627 Ep 423/500 | Train: 0.1328 | Val: 0.2758 Ep 424/500 | Train: 0.1311 | Val: 0.2583 Ep 425/500 | Train: 0.1319 | Val: 0.2640 Checkpoint saved. Ep 426/500 | Train: 0.1340 | Val: 0.2719 Ep 427/500 | Train: 0.1294 | Val: 0.2565 Ep 428/500 | Train: 0.1329 | Val: 0.2663 Ep 429/500 | Train: 0.1276 | Val: 0.2685 Ep 430/500 | Train: 0.1292 | Val: 0.2580 Checkpoint saved. Ep 431/500 | Train: 0.1270 | Val: 0.2573 Ep 432/500 | Train: 0.1338 | Val: 0.2545 Ep 433/500 | Train: 0.1297 | Val: 0.2639 Ep 434/500 | Train: 0.1257 | Val: 0.2661 Ep 435/500 | Train: 0.1262 | Val: 0.2672 Checkpoint saved. Ep 436/500 | Train: 0.1305 | Val: 0.2637 Ep 437/500 | Train: 0.1256 | Val: 0.2670 Ep 438/500 | Train: 0.1262 | Val: 0.2592 Ep 439/500 | Train: 0.1252 | Val: 0.2607 Ep 440/500 | Train: 0.1237 | Val: 0.2540 Checkpoint saved. Ep 441/500 | Train: 0.1247 | Val: 0.2567 Ep 442/500 | Train: 0.1227 | Val: 0.2669 Ep 443/500 | Train: 0.1259 | Val: 0.2570 Ep 444/500 | Train: 0.1235 | Val: 0.2774 Ep 445/500 | Train: 0.1212 | Val: 0.2537 Checkpoint saved. Ep 446/500 | Train: 0.1220 | Val: 0.2644 Ep 447/500 | Train: 0.1204 | Val: 0.2532 Ep 448/500 | Train: 0.1235 | Val: 0.2547 Ep 449/500 | Train: 0.1244 | Val: 0.2568 Ep 450/500 | Train: 0.1205 | Val: 0.2582 Checkpoint saved. Ep 451/500 | Train: 0.1204 | Val: 0.2677 Ep 452/500 | Train: 0.1187 | Val: 0.2539 Ep 453/500 | Train: 0.1190 | Val: 0.2548 Ep 454/500 | Train: 0.1203 | Val: 0.2486 Best Model Saved (Loss: 0.2486) Ep 455/500 | Train: 0.1170 | Val: 0.2537 Checkpoint saved. Ep 456/500 | Train: 0.1172 | Val: 0.2482 Best Model Saved (Loss: 0.2482) Ep 457/500 | Train: 0.1167 | Val: 0.2515 Ep 458/500 | Train: 0.1236 | Val: 0.2559 Ep 459/500 | Train: 0.1169 | Val: 0.2487 Ep 460/500 | Train: 0.1165 | Val: 0.2533 Checkpoint saved. Ep 461/500 | Train: 0.1217 | Val: 0.2519 Ep 462/500 | Train: 0.1145 | Val: 0.2555 Ep 463/500 | Train: 0.1181 | Val: 0.2608 Ep 464/500 | Train: 0.1170 | Val: 0.2469 Best Model Saved (Loss: 0.2469) Ep 465/500 | Train: 0.1160 | Val: 0.2468 Best Model Saved (Loss: 0.2468) Checkpoint saved. Ep 466/500 | Train: 0.1177 | Val: 0.2453 Best Model Saved (Loss: 0.2453) Ep 467/500 | Train: 0.1145 | Val: 0.2569 Ep 468/500 | Train: 0.1131 | Val: 0.2624 Ep 469/500 | Train: 0.1131 | Val: 0.2487 Ep 470/500 | Train: 0.1125 | Val: 0.2559 Checkpoint saved. Ep 471/500 | Train: 0.1148 | Val: 0.2504 Ep 472/500 | Train: 0.1152 | Val: 0.2471 Ep 473/500 | Train: 0.1180 | Val: 0.2578 Ep 474/500 | Train: 0.1139 | Val: 0.2430 Best Model Saved (Loss: 0.2430) Ep 475/500 | Train: 0.1151 | Val: 0.2589 Checkpoint saved. Ep 476/500 | Train: 0.1109 | Val: 0.2537 Ep 477/500 | Train: 0.1147 | Val: 0.2612 Ep 478/500 | Train: 0.1105 | Val: 0.2589 Ep 479/500 | Train: 0.1104 | Val: 0.2488 Ep 480/500 | Train: 0.1102 | Val: 0.2574 Checkpoint saved. Ep 481/500 | Train: 0.1095 | Val: 0.2503 Ep 482/500 | Train: 0.1137 | Val: 0.2484 Ep 483/500 | Train: 0.1089 | Val: 0.2472 Ep 484/500 | Train: 0.1112 | Val: 0.2576 Ep 485/500 | Train: 0.1083 | Val: 0.2480 Checkpoint saved. Ep 486/500 | Train: 0.1076 | Val: 0.2527 Ep 487/500 | Train: 0.1083 | Val: 0.2501 Ep 488/500 | Train: 0.1080 | Val: 0.2446 Ep 489/500 | Train: 0.1085 | Val: 0.2453 Ep 490/500 | Train: 0.1062 | Val: 0.2509 Checkpoint saved. Ep 491/500 | Train: 0.1082 | Val: 0.2479 Ep 492/500 | Train: 0.1049 | Val: 0.2479 Ep 493/500 | Train: 0.1050 | Val: 0.2505 Ep 494/500 | Train: 0.1048 | Val: 0.2472 Ep 495/500 | Train: 0.1050 | Val: 0.2457 Checkpoint saved. Ep 496/500 | Train: 0.1049 | Val: 0.2452 Ep 497/500 | Train: 0.1048 | Val: 0.2456 Ep 498/500 | Train: 0.1035 | Val: 0.2451 Ep 499/500 | Train: 0.1087 | Val: 0.2448 Ep 500/500 | Train: 0.1043 | Val: 0.2460 Checkpoint saved. Training Complete!
In [ ]:
# CELL 4: FINAL EVALUATION & REPORTING
import matplotlib.pyplot as plt
import json
import seaborn as sns
from sklearn.metrics import accuracy_score, f1_score, jaccard_score, precision_score, recall_score, confusion_matrix
def generate_final_report(model, loader, history, save_dir, model_name="SatMAE_Code1_Fair_500"):
print(f"Generating Final Report for {model_name}...")
# 1. LOAD BEST MODEL
# ---------------------------------
best_path = save_dir + "SatMAE_500_BEST.pth"
if os.path.exists(best_path):
print(f" Loading Best Saved Model from: {best_path}")
model.load_state_dict(torch.load(best_path, map_location=device))
else:
print(" Best model file not found. Using current model weights.")
model.eval()
# 2. PLOT LOSS CURVE
if 'train_loss' in history and len(history['train_loss']) > 0:
plt.figure(figsize=(12, 6))
epochs = range(1, len(history['train_loss']) + 1)
# Plot
plt.plot(epochs, history['train_loss'], 'b-', label='Training Loss', linewidth=2)
plt.plot(epochs, history['val_loss'], 'r--', label='Validation Loss', linewidth=2)
plt.title(f'Training Progress ({len(epochs)} Epochs)', fontsize=16)
plt.xlabel('Epochs', fontsize=12)
plt.ylabel('Compound Loss', fontsize=12)
plt.legend(fontsize=12)
plt.grid(True, alpha=0.3)
# Save
graph_path = f"{save_dir}{model_name}_LossCurve.png"
plt.savefig(graph_path)
plt.show()
print(f" Loss Graph saved to: {graph_path}")
else:
print(" No history data found to plot.")
# 3. VISUALIZATION (Samples)
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 3 samples
fig, axes = plt.subplots(3, 3, figsize=(12, 12))
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(3):
if i >= len(x_batch): break
# Construct RGB from bands B4(2), B3(1), B2(0)
rgb = x_batch[i, [2, 1, 0], 1, :, :].permute(1, 2, 0).cpu().numpy()
rgb = np.clip(rgb * 3.5, 0, 1) # Brighten 3.5x for visibility
gt_img = y_batch[i, 0].cpu().numpy()
pred_img = preds[i, 0].cpu().numpy()
# Overlay metrics on image
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.5))
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 Error: {e}")
# 4. DETAILED METRICS REPORT
print(" Calculating Detailed Metrics on Full Validation Set...")
all_preds, all_targets = [], []
with torch.no_grad():
for x, y in loader:
x = x.to(device)
logits = model(x)
# Binary Threshold
p_batch = (torch.sigmoid(logits) > 0.5).float().cpu().numpy().flatten()
y_batch = y.numpy().flatten()
all_preds.extend(p_batch)
all_targets.extend(y_batch)
y_p = np.array(all_preds).astype(int)
y_t = np.array(all_targets).astype(int)
metrics = {
"Model": model_name,
"Timestamp": datetime.now().strftime("%Y-%m-%d %H:%M"),
"Pixel_Accuracy": round(accuracy_score(y_t, y_p), 4),
"IoU_Score": round(jaccard_score(y_t, y_p, average='binary'), 4),
"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),
"Confusion_Matrix": {
"TN": int(confusion_matrix(y_t, y_p).ravel()[0]),
"FP": int(confusion_matrix(y_t, y_p).ravel()[1]),
"FN": int(confusion_matrix(y_t, y_p).ravel()[2]),
"TP": int(confusion_matrix(y_t, y_p).ravel()[3])
}
}
# Save JSON to Drive
json_path = f"{save_dir}{model_name}_DetailedMetrics.json"
with open(json_path, 'w') as f:
json.dump(metrics, f, indent=4)
print("\n FINAL METRICS REPORT:")
print(json.dumps(metrics, indent=4))
print(f"\n All files (PNGs + JSON) saved to: {save_dir}")
# EXECUTE
if 'history' not in globals():
try:
print(" Attempting to load history from checkpoint...")
ckpt = torch.load(SAVE_DIR + "checkpoint.pth")
history = ckpt['history']
print(" History loaded.")
except:
print(" Could not load history. Graphs will be empty.")
history = {'train_loss': [], 'val_loss': []}
generate_final_report(model, val_loader, history, SAVE_DIR)
Generating Final Report for SatMAE_Code1_Fair_500... Loading Best Saved Model from: /content/drive/MyDrive/SatMAE_LongTrain_500/SatMAE_500_BEST.pth
Loss Graph saved to: /content/drive/MyDrive/SatMAE_LongTrain_500/SatMAE_Code1_Fair_500_LossCurve.png
Visualizations saved to: /content/drive/MyDrive/SatMAE_LongTrain_500/SatMAE_Code1_Fair_500_Visuals.png
Calculating Detailed Metrics on Full Validation Set...
FINAL METRICS REPORT:
{
"Model": "SatMAE_Code1_Fair_500",
"Timestamp": "2026-01-13 06:12",
"Pixel_Accuracy": 0.8671,
"IoU_Score": 0.8008,
"F1_Score": 0.8894,
"Precision": 0.8894,
"Recall": 0.8893,
"Confusion_Matrix": {
"TN": 66806,
"FP": 13336,
"FN": 13341,
"TP": 107221
}
}
All files (PNGs + JSON) saved to: /content/drive/MyDrive/SatMAE_LongTrain_500/
In [ ]:
# CELL 1: HYPER-STACK DATA LOADING
import os
import numpy as np
import rasterio
from rasterio.windows import from_bounds
import ee
import geemap
import torch
import cv2
import shutil
import time
# CONFIG
TIME_WINDOWS = [
('2024-10-15', '2024-11-15'), # 1. Pre-Sowing
('2024-11-16', '2024-12-15'), # 2. Early Growth
('2024-12-16', '2025-01-15'), # 3. Late Growth
('2025-01-16', '2025-02-15'), # 4. Peak Greenness
('2025-02-16', '2025-03-15'), # 5. Flowering
('2025-03-16', '2025-04-15') # 6. Harvest
]
ASSET_ID = 'projects/[REDACTED_FOR_SECURITY]/assets/Punjab_Mask_2024_NEW'
PATCH_SIZE = 224
def get_hyper_satmae_data():
print(" Starting Hyper-Stack Ingestion...")
mask_img = ee.Image(ASSET_ID)
roi_geom = mask_img.geometry()
# Download Mask
mask_file = 'local_mask_hyper.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 = []
for i, (start, end) in enumerate(TIME_WINDOWS):
fname = f'hyper_time_{i}.tif'
# RETRY LOGIC for robustness
attempts = 0
while not os.path.exists(fname) and attempts < 3:
try:
print(f" Downloading Step {i+1}/{len(TIME_WINDOWS)} (Attempt {attempts+1})...")
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)
except Exception as e:
print(f" Error: {e}")
attempts += 1
time.sleep(2) # Wait 2 seconds before retry
# Fallback if download failed 3 times
if not os.path.exists(fname):
print(f" Failed to download {start}. Checking fallback...")
if i > 0:
print(" Copying previous month's data.")
shutil.copy(f'hyper_time_{i-1}.tif', fname)
else:
raise RuntimeError(" CRITICAL: First time step failed. Cannot proceed.")
# LOAD & PROCESS
with rasterio.open(fname) as src:
arr = src.read() # (8, H, W)
arr = np.transpose(arr, (1, 2, 0)) # (H, W, 8)
if arr.shape[:2] != (target_h, target_w):
arr = cv2.resize(arr, (target_w, target_h), interpolation=cv2.INTER_LINEAR)
# Normalization
s2_bands = np.clip(arr[:,:,:6] / 5000.0, 0, 1)
s1_bands = np.clip((arr[:,:,6:] - (-25.0)) / 25.0, 0, 1)
# Indices
nir = s2_bands[:, :, 3] + 1e-6
red = s2_bands[:, :, 2] + 1e-6
blue = s2_bands[:, :, 0] + 1e-6
ndvi = ((nir - red) / (nir + red) + 1) / 2.0
evi = np.clip(2.5 * ((nir - red) / (nir + 6 * red - 7.5 * blue + 1)), 0, 1)
combined = np.concatenate([s2_bands, s1_bands, ndvi[:, :, None], evi[:, :, None]], axis=2)
stack.append(combined)
full_cube = np.stack(stack, axis=2)
x_out, y_out = [], []
stride = PATCH_SIZE
print(" 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.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) # (B, C, T, H, W)
y = np.array(y_out, dtype=np.float32)[:, None, :, :]
X = np.nan_to_num(X, nan=0.0)
print(f" Hyper-Dataset Ready. Shape: {X.shape}")
return torch.tensor(X), torch.tensor(y)
X_data, y_data = get_hyper_satmae_data()
/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) 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.
Starting Hyper-Stack Ingestion...
...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:googleapiclient.http:Sleeping 0.50 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:googleapiclient.http:Sleeping 0.60 seconds before retry 1 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/satmae-2026/thumbnails?fields=name&alt=json, after 429 WARNING:googleapiclient.http:Sleeping 1.78 seconds before retry 1 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/satmae-2026/thumbnails?fields=name&alt=json, after 429 WARNING:googleapiclient.http:Sleeping 1.92 seconds before retry 1 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/satmae-2026/thumbnails?fields=name&alt=json, after 429 WARNING:googleapiclient.http:Sleeping 0.10 seconds before retry 1 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/satmae-2026/thumbnails?fields=name&alt=json, after 429 WARNING:googleapiclient.http:Sleeping 0.14 seconds before retry 1 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/satmae-2026/thumbnails?fields=name&alt=json, after 429 WARNING:googleapiclient.http:Sleeping 1.85 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
Downloading Step 1/6 (Attempt 1)...
0%| |0/48 tiles [00:00<?]
WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:googleapiclient.http:Sleeping 0.32 seconds before retry 1 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/satmae-2026/thumbnails?fields=name&alt=json, after 429 WARNING:googleapiclient.http:Sleeping 0.15 seconds before retry 1 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/satmae-2026/thumbnails?fields=name&alt=json, after 429 WARNING:googleapiclient.http:Sleeping 0.86 seconds before retry 1 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/satmae-2026/thumbnails?fields=name&alt=json, after 429
Downloading Step 2/6 (Attempt 1)...
0%| |0/48 tiles [00:00<?]
WARNING:googleapiclient.http:Sleeping 1.94 seconds before retry 1 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/satmae-2026/thumbnails?fields=name&alt=json, after 429 WARNING:googleapiclient.http:Sleeping 0.29 seconds before retry 1 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/satmae-2026/thumbnails?fields=name&alt=json, after 429 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:googleapiclient.http:Sleeping 0.64 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
Downloading Step 3/6 (Attempt 1)...
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
Downloading Step 4/6 (Attempt 1)...
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:googleapiclient.http:Sleeping 1.37 seconds before retry 1 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/satmae-2026/thumbnails?fields=name&alt=json, after 429 WARNING:googleapiclient.http:Sleeping 0.83 seconds before retry 1 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/satmae-2026/thumbnails?fields=name&alt=json, after 429 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10
Downloading Step 5/6 (Attempt 1)...
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:googleapiclient.http:Sleeping 1.80 seconds before retry 1 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/satmae-2026/thumbnails?fields=name&alt=json, after 429 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10 WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10
Downloading Step 6/6 (Attempt 1)...
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
Creating Patches... Hyper-Dataset Ready. Shape: (25, 10, 6, 224, 224)
In [ ]:
# CELL 2: MODEL DEFINITION
import torch.nn as nn
from huggingface_hub import hf_hub_download
class SatMAEPatchEmbed(nn.Module):
def __init__(self, in_chans=10, 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=6, in_chans=10, 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=6, embed_dim=768):
super().__init__()
print(f" Constructing SatMAE Hyper-Stack ({num_frames} Time Steps)...")
self.num_frames = num_frames # <--- SAVE THIS VARIABLE
self.backbone = SatMAEBackbone(num_frames=num_frames, in_chans=10, embed_dim=embed_dim)
try:
print(" Adapting Google ViT Weights to 10 Channels...")
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, 10, 16, 16)
new_w[:, :3] = w
new_w[:, 3:] = w.mean(dim=1, keepdim=True).repeat(1, 7, 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 Adapted.")
except:
print(" Weights missing, using random init.")
# FREEZING
for param in self.backbone.parameters(): param.requires_grad = False
self.backbone.time_embed.requires_grad = True
self.backbone.patch_embed.proj.weight.requires_grad = True
# 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).permute(0, 4, 1, 2, 3).flatten(1, 2)
features = self.temporal_agg(features)
return self.decoder(features)
In [ ]:
import os
# Force delete the old checkpoint
if os.path.exists('/content/drive/MyDrive/SatMAE_Hyper_500/checkpoint_hyper.pth'):
os.remove('/content/drive/MyDrive/SatMAE_Hyper_500/checkpoint_hyper.pth')
print(" Bad checkpoint deleted. Ready to start fresh from Epoch 0.")
else:
print(" No checkpoint found. Good to go!")
No checkpoint found. Good to go!
In [ ]:
# CELL 3: TRAINING
import torch
import torch.optim as optim
import torch.optim.swa_utils as swa_utils
from torch.utils.data import TensorDataset, DataLoader, random_split
from scipy.ndimage import distance_transform_edt as distance
import os
# --- CONFIGURATION ---
RESUME = False
BATCH_SIZE = 4
EPOCHS = 500
CHECKPOINT_PATH = SAVE_DIR + "checkpoint_hyper.pth"
# 1. TIME-SAFE AUGMENTATION
def apply_augmentation(x, y):
"""
IT will Flip only Spatial Dimensions (H=3, W=4 for x) | (H=2, W=3 for y).
Strictly preserves Time (Dim 2).
"""
if np.random.rand() > 0.5: # H-Flip (Width)
x = torch.flip(x, [4])
y = torch.flip(y, [3])
if np.random.rand() > 0.5: # V-Flip (Height)
x = torch.flip(x, [3])
y = torch.flip(y, [2])
k = np.random.randint(0, 4) # Rotation
x = torch.rot90(x, k, [3, 4])
y = torch.rot90(y, k, [2, 3])
return x, y
# 2. FIXED LOSS (Reshape instead of View)
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) # Fixed
targets = targets.reshape(-1) # Fixed
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)
probs = torch.sigmoid(pred)
return torch.mean((probs - gt) ** 2 * (1 + self.alpha * torch.abs(dist_map)))
class CompoundLoss(nn.Module):
def __init__(self):
super().__init__(); self.dice = DiceLoss(); self.boundary = HausdorffDTLoss(alpha=2.0)
def forward(self, p, t): return 0.7*self.dice(p, t) + 0.3*self.boundary(p, t)
# 3. SETUP MODEL & OPTIMIZER
model = SatMAESegmentation(num_frames=6, embed_dim=768).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, eta_min=1e-6)
# SWA Config
swa_model = swa_utils.AveragedModel(model)
swa_start = 350
swa_scheduler = swa_utils.SWALR(optimizer, swa_lr=5e-5)
# Loaders
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)
# 4. RESUME LOGIC (WITH FORCE DELETE)
start_epoch = 0
history = {'train_loss': [], 'val_loss': []}
best_loss = float('inf')
if RESUME and os.path.exists(CHECKPOINT_PATH):
print(" Found Checkpoint. Resuming...")
ckpt = torch.load(CHECKPOINT_PATH)
model.load_state_dict(ckpt['model_state_dict'])
optimizer.load_state_dict(ckpt['optimizer_state_dict'])
start_epoch = ckpt['epoch'] + 1
history = ckpt['history']
best_loss = ckpt['best_loss']
elif not RESUME and os.path.exists(CHECKPOINT_PATH):
print(" Force Restart: Deleting old checkpoint...")
os.remove(CHECKPOINT_PATH)
print(" Starting Fresh from Epoch 0.")
else:
print(" No checkpoint found. Starting Fresh.")
# 5. TRAINING LOOP
print(f" Starting Hyper-Stack Training ({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) # Safe Augmentation
optimizer.zero_grad()
preds = model(x)
loss = criterion(preds, y)
loss.backward()
optimizer.step()
train_loss += loss.item()
# Validation
model.eval()
val_loss = 0
with torch.no_grad():
for x, y in val_loader:
x, y = x.to(device), y.to(device)
preds = model(x)
val_loss += criterion(preds, y).item()
avg_t = train_loss / len(train_loader)
avg_v = val_loss / len(val_loader)
history['train_loss'].append(avg_t)
history['val_loss'].append(avg_v)
# --- SWA & SAVING LOGIC ---
if ep >= swa_start:
# Phase 2: SWA Mode
swa_model.update_parameters(model)
swa_scheduler.step()
lr_stat = f"SWA-LR: {swa_scheduler.get_last_lr()[0]:.1e}"
else:
# Phase 1: Standard Mode
scheduler.step()
lr_stat = f"LR: {scheduler.get_last_lr()[0]:.1e}"
# SAVE BEST MODEL
if avg_v < best_loss:
best_loss = avg_v
torch.save(model.state_dict(), SAVE_DIR + "SatMAE_Hyper_Best.pth")
print(f" Best Model Updated (Loss: {best_loss:.4f})")
# Print Stats (Every 5 Epochs)
if (ep+1) % 5 == 0:
print(f"Ep {ep+1} | T: {avg_t:.4f} | V: {avg_v:.4f} | {lr_stat}")
# SAVE CHECKPOINT (Every 5 Epochs)
if (ep+1) % 5 == 0:
torch.save({
'epoch': ep,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'history': history,
'best_loss': best_loss
}, CHECKPOINT_PATH)
print(" Checkpoint Saved.")
print(" Finalizing SWA Model...")
swa_utils.update_bn(train_loader, swa_model, device=device)
torch.save(swa_model.state_dict(), SAVE_DIR + "SatMAE_Hyper_SWA.pth")
print("Done.")
Constructing SatMAE Hyper-Stack (6 Time Steps)... Adapting Google ViT Weights to 10 Channels... Weights Adapted. Force Restart: Deleting old checkpoint... Starting Fresh from Epoch 0. Starting Hyper-Stack Training (500 Epochs)... Best Model Updated (Loss: 1.0251) Best Model Updated (Loss: 0.8995) Best Model Updated (Loss: 0.7510) Best Model Updated (Loss: 0.6757) Best Model Updated (Loss: 0.6691) Ep 5 | T: 0.4751 | V: 0.6691 | LR: 9.8e-05 Checkpoint Saved. Best Model Updated (Loss: 0.5908) Best Model Updated (Loss: 0.5308) Best Model Updated (Loss: 0.4995) Best Model Updated (Loss: 0.3803) Ep 10 | T: 0.4097 | V: 0.3803 | LR: 9.1e-05 Checkpoint Saved. Ep 15 | T: 0.3892 | V: 0.4322 | LR: 8.0e-05 Checkpoint Saved. Best Model Updated (Loss: 0.3486) Ep 20 | T: 0.3657 | V: 0.4360 | LR: 6.6e-05 Checkpoint Saved. Best Model Updated (Loss: 0.3454) Ep 25 | T: 0.3834 | V: 0.3533 | LR: 5.1e-05 Checkpoint Saved. Best Model Updated (Loss: 0.3442) Ep 30 | T: 0.3406 | V: 0.3442 | LR: 3.5e-05 Checkpoint Saved. Best Model Updated (Loss: 0.3440) Best Model Updated (Loss: 0.3201) Ep 35 | T: 0.3369 | V: 0.3541 | LR: 2.1e-05 Checkpoint Saved. Ep 40 | T: 0.3310 | V: 0.3429 | LR: 1.0e-05 Checkpoint Saved. Ep 45 | T: 0.3350 | V: 0.3480 | LR: 3.4e-06 Checkpoint Saved. Ep 50 | T: 0.3303 | V: 0.3368 | LR: 1.0e-04 Checkpoint Saved. Ep 55 | T: 0.3369 | V: 0.3291 | LR: 9.9e-05 Checkpoint Saved. Ep 60 | T: 0.3130 | V: 0.3682 | LR: 9.8e-05 Checkpoint Saved. Ep 65 | T: 0.3040 | V: 0.3327 | LR: 9.5e-05 Checkpoint Saved. Best Model Updated (Loss: 0.3022) Best Model Updated (Loss: 0.2930) Ep 70 | T: 0.2928 | V: 0.3522 | LR: 9.1e-05 Checkpoint Saved. Best Model Updated (Loss: 0.2837) Ep 75 | T: 0.2944 | V: 0.7233 | LR: 8.6e-05 Checkpoint Saved. Ep 80 | T: 0.2964 | V: 0.3968 | LR: 8.0e-05 Checkpoint Saved. Best Model Updated (Loss: 0.2519) Ep 85 | T: 0.2780 | V: 0.3613 | LR: 7.3e-05 Checkpoint Saved. Ep 90 | T: 0.2656 | V: 0.3216 | LR: 6.6e-05 Checkpoint Saved. Ep 95 | T: 0.2633 | V: 0.2880 | LR: 5.8e-05 Checkpoint Saved. Ep 100 | T: 0.2594 | V: 0.2766 | LR: 5.1e-05 Checkpoint Saved. Ep 105 | T: 0.2603 | V: 0.2719 | LR: 4.3e-05 Checkpoint Saved. Ep 110 | T: 0.2521 | V: 0.2780 | LR: 3.5e-05 Checkpoint Saved. Ep 115 | T: 0.2447 | V: 0.2759 | LR: 2.8e-05 Checkpoint Saved. Best Model Updated (Loss: 0.2486) Ep 120 | T: 0.2388 | V: 0.2549 | LR: 2.1e-05 Checkpoint Saved. Ep 125 | T: 0.2337 | V: 0.2559 | LR: 1.5e-05 Checkpoint Saved. Ep 130 | T: 0.2440 | V: 0.2505 | LR: 1.0e-05 Checkpoint Saved. Best Model Updated (Loss: 0.2432) Ep 135 | T: 0.2329 | V: 0.2471 | LR: 6.4e-06 Checkpoint Saved. Ep 140 | T: 0.2341 | V: 0.2452 | LR: 3.4e-06 Checkpoint Saved. Ep 145 | T: 0.2409 | V: 0.2552 | LR: 1.6e-06 Checkpoint Saved. Ep 150 | T: 0.2385 | V: 0.2561 | LR: 1.0e-04 Checkpoint Saved. Ep 155 | T: 0.2456 | V: 0.3684 | LR: 1.0e-04 Checkpoint Saved. Ep 160 | T: 0.2368 | V: 0.3029 | LR: 9.9e-05 Checkpoint Saved. Best Model Updated (Loss: 0.2324) Ep 165 | T: 0.2236 | V: 0.2324 | LR: 9.9e-05 Checkpoint Saved. Ep 170 | T: 0.2155 | V: 0.2603 | LR: 9.8e-05 Checkpoint Saved. Ep 175 | T: 0.2164 | V: 0.2942 | LR: 9.6e-05 Checkpoint Saved. Best Model Updated (Loss: 0.2283) Ep 180 | T: 0.2077 | V: 0.2453 | LR: 9.5e-05 Checkpoint Saved. Ep 185 | T: 0.2053 | V: 0.3523 | LR: 9.3e-05 Checkpoint Saved. Ep 190 | T: 0.1970 | V: 0.2370 | LR: 9.1e-05 Checkpoint Saved. Best Model Updated (Loss: 0.2277) Ep 195 | T: 0.1918 | V: 0.2713 | LR: 8.8e-05 Checkpoint Saved. Best Model Updated (Loss: 0.2232) Best Model Updated (Loss: 0.2218) Ep 200 | T: 0.2063 | V: 0.2218 | LR: 8.6e-05 Checkpoint Saved. Best Model Updated (Loss: 0.2206) Ep 205 | T: 0.1925 | V: 0.2233 | LR: 8.3e-05 Checkpoint Saved. Best Model Updated (Loss: 0.2204) Ep 210 | T: 0.1961 | V: 0.2380 | LR: 8.0e-05 Checkpoint Saved. Best Model Updated (Loss: 0.2092) Ep 215 | T: 0.1772 | V: 0.2190 | LR: 7.6e-05 Checkpoint Saved. Best Model Updated (Loss: 0.2034) Ep 220 | T: 0.1720 | V: 0.2034 | LR: 7.3e-05 Checkpoint Saved. Best Model Updated (Loss: 0.2031) Ep 225 | T: 0.1701 | V: 0.2031 | LR: 6.9e-05 Checkpoint Saved. Best Model Updated (Loss: 0.1997) Best Model Updated (Loss: 0.1951) Ep 230 | T: 0.1670 | V: 0.2034 | LR: 6.6e-05 Checkpoint Saved. Ep 235 | T: 0.1743 | V: 0.2078 | LR: 6.2e-05 Checkpoint Saved. Ep 240 | T: 0.1632 | V: 0.2033 | LR: 5.8e-05 Checkpoint Saved. Ep 245 | T: 0.1581 | V: 0.1979 | LR: 5.4e-05 Checkpoint Saved. Best Model Updated (Loss: 0.1937) Ep 250 | T: 0.1539 | V: 0.1997 | LR: 5.1e-05 Checkpoint Saved. Best Model Updated (Loss: 0.1904) Best Model Updated (Loss: 0.1903) Best Model Updated (Loss: 0.1823) Ep 255 | T: 0.1661 | V: 0.1823 | LR: 4.7e-05 Checkpoint Saved. Ep 260 | T: 0.1517 | V: 0.1960 | LR: 4.3e-05 Checkpoint Saved. Ep 265 | T: 0.1536 | V: 0.2066 | LR: 3.9e-05 Checkpoint Saved. Ep 270 | T: 0.1480 | V: 0.1840 | LR: 3.5e-05 Checkpoint Saved. Ep 275 | T: 0.1467 | V: 0.1850 | LR: 3.2e-05 Checkpoint Saved. Ep 280 | T: 0.1433 | V: 0.1832 | LR: 2.8e-05 Checkpoint Saved. Ep 285 | T: 0.1436 | V: 0.1888 | LR: 2.5e-05 Checkpoint Saved. Ep 290 | T: 0.1401 | V: 0.1862 | LR: 2.1e-05 Checkpoint Saved. Best Model Updated (Loss: 0.1815) Ep 295 | T: 0.1405 | V: 0.1815 | LR: 1.8e-05 Checkpoint Saved. Ep 300 | T: 0.1471 | V: 0.1877 | LR: 1.5e-05 Checkpoint Saved. Ep 305 | T: 0.1406 | V: 0.1862 | LR: 1.3e-05 Checkpoint Saved. Ep 310 | T: 0.1497 | V: 0.1862 | LR: 1.0e-05 Checkpoint Saved. Ep 315 | T: 0.1358 | V: 0.1848 | LR: 8.3e-06 Checkpoint Saved. Ep 320 | T: 0.1371 | V: 0.1828 | LR: 6.4e-06 Checkpoint Saved. Best Model Updated (Loss: 0.1814) Best Model Updated (Loss: 0.1807) Ep 325 | T: 0.1454 | V: 0.1852 | LR: 4.8e-06 Checkpoint Saved. Ep 330 | T: 0.1380 | V: 0.1826 | LR: 3.4e-06 Checkpoint Saved. Ep 335 | T: 0.1454 | V: 0.1812 | LR: 2.4e-06 Checkpoint Saved. Best Model Updated (Loss: 0.1800) Ep 340 | T: 0.1412 | V: 0.1863 | LR: 1.6e-06 Checkpoint Saved. Ep 345 | T: 0.1371 | V: 0.1807 | LR: 1.2e-06 Checkpoint Saved. Ep 350 | T: 0.1379 | V: 0.1815 | LR: 1.0e-04 Checkpoint Saved. Ep 355 | T: 0.1462 | V: 0.2154 | SWA-LR: 7.5e-05 Checkpoint Saved. Ep 360 | T: 0.1528 | V: 0.1930 | SWA-LR: 5.0e-05 Checkpoint Saved. Ep 365 | T: 0.1351 | V: 0.2021 | SWA-LR: 5.0e-05 Checkpoint Saved. Ep 370 | T: 0.1304 | V: 0.1913 | SWA-LR: 5.0e-05 Checkpoint Saved. Ep 375 | T: 0.1324 | V: 0.1931 | SWA-LR: 5.0e-05 Checkpoint Saved. Ep 380 | T: 0.1316 | V: 0.1892 | SWA-LR: 5.0e-05 Checkpoint Saved. Ep 385 | T: 0.1284 | V: 0.1924 | SWA-LR: 5.0e-05 Checkpoint Saved. Ep 390 | T: 0.1281 | V: 0.1718 | SWA-LR: 5.0e-05 Checkpoint Saved. Ep 395 | T: 0.1379 | V: 0.1814 | SWA-LR: 5.0e-05 Checkpoint Saved. Ep 400 | T: 0.1237 | V: 0.1738 | SWA-LR: 5.0e-05 Checkpoint Saved. Ep 405 | T: 0.1270 | V: 0.1808 | SWA-LR: 5.0e-05 Checkpoint Saved. Ep 410 | T: 0.1214 | V: 0.1739 | SWA-LR: 5.0e-05 Checkpoint Saved. Ep 415 | T: 0.1199 | V: 0.1797 | SWA-LR: 5.0e-05 Checkpoint Saved. Ep 420 | T: 0.1201 | V: 0.1776 | SWA-LR: 5.0e-05 Checkpoint Saved. Ep 425 | T: 0.1196 | V: 0.1850 | SWA-LR: 5.0e-05 Checkpoint Saved. Ep 430 | T: 0.1143 | V: 0.1784 | SWA-LR: 5.0e-05 Checkpoint Saved. Ep 435 | T: 0.1173 | V: 0.1882 | SWA-LR: 5.0e-05 Checkpoint Saved. Ep 440 | T: 0.1141 | V: 0.1758 | SWA-LR: 5.0e-05 Checkpoint Saved. Ep 445 | T: 0.1137 | V: 0.1784 | SWA-LR: 5.0e-05 Checkpoint Saved. Ep 450 | T: 0.1125 | V: 0.1810 | SWA-LR: 5.0e-05 Checkpoint Saved. Ep 455 | T: 0.1155 | V: 0.1704 | SWA-LR: 5.0e-05 Checkpoint Saved. Ep 460 | T: 0.1173 | V: 0.1851 | SWA-LR: 5.0e-05 Checkpoint Saved. Ep 465 | T: 0.1064 | V: 0.1739 | SWA-LR: 5.0e-05 Checkpoint Saved. Ep 470 | T: 0.1103 | V: 0.1829 | SWA-LR: 5.0e-05 Checkpoint Saved. Ep 475 | T: 0.1066 | V: 0.1824 | SWA-LR: 5.0e-05 Checkpoint Saved. Ep 480 | T: 0.1109 | V: 0.1798 | SWA-LR: 5.0e-05 Checkpoint Saved. Ep 485 | T: 0.1041 | V: 0.1758 | SWA-LR: 5.0e-05 Checkpoint Saved. Ep 490 | T: 0.1038 | V: 0.1725 | SWA-LR: 5.0e-05 Checkpoint Saved. Ep 495 | T: 0.1009 | V: 0.1698 | SWA-LR: 5.0e-05 Checkpoint Saved. Ep 500 | T: 0.1037 | V: 0.1798 | SWA-LR: 5.0e-05 Checkpoint Saved. 💾 Finalizing SWA Model... Done.
In [ ]:
# CELL 4: VISUALIZATION & SAMPLING
import matplotlib.pyplot as plt
import numpy as np
import torch
import os
# CONFIG
MODEL_NAME = "SatMAE_Hyper_500_Epochs"
GRAPH_PATH = os.path.join(SAVE_DIR, f"{MODEL_NAME}_LossCurve.png")
VISUALS_PATH = os.path.join(SAVE_DIR, f"{MODEL_NAME}_VisualSamples.png")
# Check if necessary data exists
if 'history' not in globals() or not history['train_loss']:
print(" No training history found. Cannot plot loss curve.")
# Try loading from checkpoint if available
ckpt_path = os.path.join(SAVE_DIR, "checkpoint_hyper.pth")
if os.path.exists(ckpt_path):
print(" Attempting to load history from checkpoint...")
ckpt = torch.load(ckpt_path, map_location=device)
history = ckpt['history']
print(" History loaded.")
else:
print("Checkpoint not found. Skipping loss plot.")
history = None
# 1. PLOT LOSS CURVE
if history:
print(f" Generating Loss Curve Graph for {len(history['train_loss'])} epochs...")
epochs = range(1, len(history['train_loss']) + 1)
plt.figure(figsize=(10, 6))
# Plot lines with markers for better visibility
plt.plot(epochs, history['train_loss'], 'b-o', label='Training Loss', linewidth=2, markersize=4)
plt.plot(epochs, history['val_loss'], 'r--s', label='Validation Loss', linewidth=2, markersize=4)
# Styling
plt.title(f'Training Progress: {MODEL_NAME}', fontsize=16, fontweight='bold')
plt.xlabel('Epochs', fontsize=12)
plt.ylabel('Compound Loss (Dice + Hausdorff)', fontsize=12)
plt.legend(fontsize=12, loc='upper right', frameon=True, shadow=True)
plt.grid(True, linestyle='--', alpha=0.6)
plt.xticks(fontsize=10)
plt.yticks(fontsize=10)
plt.tight_layout()
# Save and Show
plt.savefig(GRAPH_PATH, dpi=150)
print(f" Loss graph saved to: {GRAPH_PATH}")
plt.show()
# 2. GENERATE VISUAL SAMPLES
print("\n Generating Visual Sample Comparisons...")
model.eval()
try:
# Get a single batch from the validation loader
x_batch, y_batch = next(iter(val_loader))
x_batch = x_batch.to(device)
with torch.no_grad():
# Get model predictions
logits = model(x_batch)
# Apply sigmoid and threshold at 0.5 for binary mask
preds = (torch.sigmoid(logits) > 0.5).float()
# Determine how many samples to plot (max 3)
num_samples = min(3, x_batch.shape[0])
# Setup figure: rows = samples, cols = Input, GT, Pred
fig, axes = plt.subplots(num_samples, 3, figsize=(15, 5 * num_samples))
# Handle single-sample case where axes is 1D
if num_samples == 1: axes = np.expand_dims(axes, axis=0)
cols = ["Input (RGB - Peak Season)", "Ground Truth Mask", "Predicted Mask"]
for ax, col in zip(axes[0], cols):
ax.set_title(col, fontsize=14, fontweight='bold', pad=10)
for i in range(num_samples):
# --- A. Extract RGB Image for Display ---
# Data shape is (B, C, T, H, W). We need (H, W, C) for matplotlib.
# Channels [2, 1, 0] correspond to B4(Red), B3(Green), B2(Blue).
# Time Step 3 corresponds to Peak Season (Jan/Feb).
rgb_img = x_batch[i, [2, 1, 0], 3, :, :].permute(1, 2, 0).cpu().numpy()
# Brighten and clip to 0-1 range for better visibility
rgb_img = np.clip(rgb_img * 3.5, 0, 1)
# --- B. Extract Masks ---
gt_mask = y_batch[i, 0].numpy()
pred_mask = preds[i, 0].cpu().numpy()
# --- C. Calculate Sample IoU ---
intersection = np.logical_and(gt_mask, pred_mask).sum()
union = np.logical_or(gt_mask, pred_mask).sum()
iou = intersection / (union + 1e-6) # Avoid division by zero
# --- D. Plotting ---
# 1. RGB Input
axes[i, 0].imshow(rgb_img)
axes[i, 0].axis('off')
# 2. Ground Truth Mask
axes[i, 1].imshow(gt_mask, cmap='gray', vmin=0, vmax=1)
axes[i, 1].axis('off')
# 3. Predicted Mask with IoU Label
axes[i, 2].imshow(pred_mask, cmap='gray', vmin=0, vmax=1)
axes[i, 2].axis('off')
# Add a colored text box for IoU score
color = 'lime' if iou > 0.8 else ('yellow' if iou > 0.6 else 'red')
axes[i, 2].text(10, 30, f"IoU: {iou:.4f}", color=color, fontsize=12, fontweight='bold',
bbox=dict(facecolor='black', alpha=0.7, edgecolor=color))
plt.tight_layout()
plt.subplots_adjust(wspace=0.1, hspace=0.1)
plt.savefig(VISUALS_PATH, dpi=150, bbox_inches='tight')
print(f" Visual samples saved to: {VISUALS_PATH}")
plt.show()
except Exception as e:
print(f" Could not generate visual samples. Error: {e}")
Generating Loss Curve Graph for 500 epochs... Loss graph saved to: /content/drive/MyDrive/SatMAE_LongTrain_500/SatMAE_Hyper_500_Epochs_LossCurve.png
Generating Visual Sample Comparisons... Visual samples saved to: /content/drive/MyDrive/SatMAE_LongTrain_500/SatMAE_Hyper_500_Epochs_VisualSamples.png
In [ ]:
# CELL 5: ADVANCED METRICS REPORT
import json
import time
import numpy as np
import torch
import cv2
from scipy.spatial.distance import directed_hausdorff
from sklearn.metrics import accuracy_score, f1_score, jaccard_score, precision_score, recall_score
from datetime import datetime
# CONFIG
MODEL_NAME = "SatMAE_Full_FineTune"
SAVE_PATH = SAVE_DIR + "SatMAE_Advanced_Metrics.json"
def get_boundary(mask, thickness=1):
mask = mask.astype(np.uint8)
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
boundary = cv2.morphologyEx(mask, cv2.MORPH_GRADIENT, kernel)
return boundary > 0
def calculate_boundary_iou(pred_mask, gt_mask):
pred_boundary = get_boundary(pred_mask)
gt_boundary = get_boundary(gt_mask)
intersection = np.logical_and(pred_boundary, gt_boundary).sum()
union = np.logical_or(pred_boundary, gt_boundary).sum()
if union == 0: return 1.0 # Perfect match (both empty)
return intersection / union
def calculate_hausdorff(pred_mask, gt_mask):
# Get coordinates of all non-zero pixels
pred_pts = np.argwhere(pred_mask > 0)
gt_pts = np.argwhere(gt_mask > 0)
# Handle empty cases
if len(pred_pts) == 0 or len(gt_pts) == 0:
return 0.0 if (len(pred_pts) == 0 and len(gt_pts) == 0) else 100.0 # Penalty
# Directed Hausdorff (A->B and B->A), take the max
d1 = directed_hausdorff(pred_pts, gt_pts)[0]
d2 = directed_hausdorff(gt_pts, pred_pts)[0]
return max(d1, d2)
def generate_full_report(model, loader, device):
print(f" Starting Advanced Evaluation for {MODEL_NAME}...")
model.eval()
all_preds = []
all_targets = []
boundary_ious = []
hausdorff_dists = []
# 1. INFERENCE SPEED TEST
print(" Measuring Inference Speed...")
start_time = time.time()
total_frames = 0
with torch.no_grad():
for x, y in loader:
x = x.to(device)
# Predict
logits = model(x)
preds = (torch.sigmoid(logits) > 0.5).float().cpu().numpy()
targets = y.numpy()
total_frames += x.shape[0]
# Loop for Pixel-Wise Metrics
for i in range(len(preds)):
p_img = preds[i, 0]
t_img = targets[i, 0]
# Flatten for Standard Metrics
all_preds.extend(p_img.flatten())
all_targets.extend(t_img.flatten())
# Calculate Advanced Metrics (Per Image)
# Only calculate if ground truth has wheat (otherwise boundary is undefined)
if t_img.sum() > 0:
b_iou = calculate_boundary_iou(p_img, t_img)
h_dist = calculate_hausdorff(p_img, t_img)
boundary_ious.append(b_iou)
hausdorff_dists.append(h_dist)
end_time = time.time()
total_time = end_time - start_time
# 2. CALCULATE METRICS
print(" Calculating Aggregate Metrics...")
y_p = np.array(all_preds).astype(int)
y_t = np.array(all_targets).astype(int)
fps = total_frames / total_time
inf_time_per_batch = total_time / len(loader) # Average batch time
# Approximate inference per image (Batch Size 4)
inf_time_per_img = total_time / total_frames
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(inf_time_per_img, 4),
"FPS": round(fps, 2)
}
# 3. SAVE & PRINT
with open(SAVE_PATH, 'w') as f:
json.dump(metrics, f, indent=4)
print("\n" + "="*40)
print(" FINAL METRICS REPORT")
print("="*40)
print(json.dumps(metrics, indent=4))
print("="*40)
print(f" Report saved to: {SAVE_PATH}")
# EXECUTE
generate_full_report(model, val_loader, device)
Starting Advanced Evaluation for SatMAE_Full_FineTune...
⏱️ Measuring Inference Speed...
Calculating Aggregate Metrics...
========================================
FINAL METRICS REPORT
========================================
{
"Model": "SatMAE_Full_FineTune",
"Date": "2026-01-13 08:17",
"Standard_IoU": 0.8551,
"Boundary_IoU": 0.29,
"Hausdorff_Dist_px": 15.97,
"F1_Score": 0.9219,
"Precision": 0.9291,
"Recall": 0.9148,
"Accuracy": 0.9034,
"Inference_Time_Sec": 0.2862,
"FPS": 3.49
}
========================================
Report saved to: /content/drive/MyDrive/SatMAE_LongTrain_500/SatMAE_Advanced_Metrics.json
In [ ]:
import json
import os
# Define the file path
metrics_path = SAVE_DIR + "SatMAE_Advanced_Metrics.json"
# Define the hyperparameters dictionary
hyperparameters = {
"Architecture": "SatMAE (ViT-Base) - Frozen Backbone",
"Input_Specs": "10 Channels x 6 Time Steps",
"Epochs": 500,
"Batch_Size": 4,
"Optimizer": "AdamW (lr=1e-4)",
"Scheduler": "CosineAnnealingWarmRestarts (T0=50)",
"Loss_Function": "0.7*Dice + 0.3*Hausdorff",
"Technique_SWA": "Enabled (Start Ep 350, lr=5e-5)",
"Technique_Augmentation": "Spatial Flip/Rotate (Time Preserved)"
}
# Load, Update, and Save
if os.path.exists(metrics_path):
with open(metrics_path, 'r') as f:
data = json.load(f)
# Add hyperparameters to the existing data
data["Hyperparameters"] = hyperparameters
with open(metrics_path, 'w') as f:
json.dump(data, f, indent=4)
print(f" Hyperparameters added to {metrics_path}")
print(json.dumps(data, indent=4))
else:
print(" Metrics file not found. Run Cell 5 first.")
Hyperparameters added to /content/drive/MyDrive/SatMAE_LongTrain_500/SatMAE_Advanced_Metrics.json
{
"Model": "SatMAE_Full_FineTune",
"Date": "2026-01-13 08:17",
"Standard_IoU": 0.8551,
"Boundary_IoU": 0.29,
"Hausdorff_Dist_px": 15.97,
"F1_Score": 0.9219,
"Precision": 0.9291,
"Recall": 0.9148,
"Accuracy": 0.9034,
"Inference_Time_Sec": 0.2862,
"FPS": 3.49,
"Hyperparameters": {
"Architecture": "SatMAE (ViT-Base) - Frozen Backbone",
"Input_Specs": "10 Channels x 6 Time Steps",
"Epochs": 500,
"Batch_Size": 4,
"Optimizer": "AdamW (lr=1e-4)",
"Scheduler": "CosineAnnealingWarmRestarts (T0=50)",
"Loss_Function": "0.7*Dice + 0.3*Hausdorff",
"Technique_SWA": "Enabled (Start Ep 350, lr=5e-5)",
"Technique_Augmentation": "Spatial Flip/Rotate (Time Preserved)"
}
}
In [ ]:
# Save the model weights to a file
torch.save(model.state_dict(), "my_best_punjab_model_weights.pth")
print("Weights saved successfully!")
Weights saved successfully!
In [ ]:
import os
import shutil
from google.colab import drive
# 1. Mount Google Drive
e
drive.mount('/content/drive')
# --- CONFIGURATION
local_filename = "my_best_punjab_model_weights.pth"
drive_folder = "/content/drive/MyDrive/Satame_best_500_swa_Project"
# -------------------------------------------------------------
# 2. Check if the file exists locally
if os.path.exists(local_filename):
print(f" Found '{local_filename}' in the current session.")
# Create the destination folder in Drive if it doesn't exist
if not os.path.exists(drive_folder):
os.makedirs(drive_folder)
print(f"Created new folder in Drive: {drive_folder}")
# 3. Copy the file to Drive
destination_path = os.path.join(drive_folder, local_filename)
try:
shutil.copy(local_filename, destination_path)
print(f" Success! Weights saved to: {destination_path}")
except Exception as e:
print(f" Error copying file: {e}")
else:
print(f" Could not find '{local_filename}' in the current folder.")
print("List of files actually present here:")
print(os.listdir('.')) # Lists all files so you can see the correct name
Drive already mounted at /content/drive; to attempt to forcibly remount, call drive.mount("/content/drive", force_remount=True).
Found 'my_best_punjab_model_weights.pth' in the current session.
Created new folder in Drive: /content/drive/MyDrive/Satame_best_500_swa_Project
Success! Weights saved to: /content/drive/MyDrive/Satame_best_500_swa_Project/my_best_punjab_model_weights.pth
In [ ]:
import torch
import numpy as np
import os
import shutil
from google.colab import drive
drive.mount('/content/drive')
drive_folder = "/content/drive/MyDrive/SatMAE_Project/Final_SWA_Save"
if not os.path.exists(drive_folder):
os.makedirs(drive_folder)
model_filename = "SatMAE_SWA_500Epochs.pth"
embeddings_filename = "SatMAE_SWA_Embeddings.npy"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 1. UNWRAP MODEL
if hasattr(swa_model, 'module'):
segmentation_model = swa_model.module
else:
segmentation_model = swa_model
segmentation_model.eval()
# 2. FIND BACKBONE
if hasattr(segmentation_model, 'backbone'):
encoder = segmentation_model.backbone
elif hasattr(segmentation_model, 'encoder'):
encoder = segmentation_model.encoder
elif hasattr(segmentation_model, 'body'):
encoder = segmentation_model.body
else:
# Fallback: try to use the model itself if no backbone attribute found
encoder = segmentation_model
# 3. SAVE CONFIG & WEIGHTS
try:
# Try to access attributes safely
img_size = getattr(encoder.patch_embed, 'img_size', 224) if hasattr(encoder, 'patch_embed') else 224
patch_size = getattr(encoder.patch_embed, 'patch_size', 16) if hasattr(encoder, 'patch_embed') else 16
in_chans = getattr(encoder.patch_embed.proj, 'in_channels', 10) if hasattr(encoder, 'patch_embed') else 10
embed_dim = getattr(encoder, 'embed_dim', 768)
config_to_save = {
"model_class": "SatMAEBackbone",
"img_size": img_size,
"patch_size": patch_size,
"in_chans": in_chans,
"embed_dim": embed_dim,
"notes": "Extracted from SatMAESegmentation"
}
except Exception as e:
config_to_save = {"error": str(e)}
full_package = {
"config": config_to_save,
"model_state_dict": segmentation_model.state_dict()
}
torch.save(full_package, model_filename)
shutil.copy(model_filename, os.path.join(drive_folder, model_filename))
# 4. EXTRACT EMBEDDINGS
extraction_loader = val_loader
all_embeddings = []
with torch.no_grad():
for i, batch in enumerate(extraction_loader):
if isinstance(batch, (list, tuple)):
images = batch[0]
elif isinstance(batch, dict):
images = batch['image']
else:
images = batch
images = images.to(device)
# 5. ROBUST FORWARD PASS
# Try standard methods since forward_encoder is missing
if hasattr(encoder, 'forward_features'):
features = encoder.forward_features(images)
else:
features = encoder(images)
# Handle List output (common in Segmentation backbones)
if isinstance(features, (list, tuple)):
features = features[-1] # Take the deepest layer
# Handle Dimensions
# Case A: (Batch, Channels, H, W) -> Global Average Pool
if features.dim() == 4:
features = features.mean(dim=[2, 3])
# Case B: (Batch, Tokens, Channels) -> Take CLS Token
elif features.dim() == 3:
features = features[:, 0, :]
all_embeddings.append(features.cpu().numpy())
if all_embeddings:
final_embeddings = np.concatenate(all_embeddings, axis=0)
np.save(embeddings_filename, final_embeddings)
shutil.copy(embeddings_filename, os.path.join(drive_folder, embeddings_filename))
print(f"Saved {final_embeddings.shape} embeddings to {drive_folder}")
else:
print("No embeddings extracted")
Drive already mounted at /content/drive; to attempt to forcibly remount, call drive.mount("/content/drive", force_remount=True).
Saved (4, 768) embeddings to /content/drive/MyDrive/SatMAE_Project/Final_SWA_Save