import ee
import geemap
import datetime
import os
import numpy as np
import rasterio
from rasterio.windows import from_bounds
from sklearn.model_selection import train_test_split
!pip install geedim
!pip install geemap

# --- 1. INITIALIZATION & CONFIGURATION ---
MY_PROJECT_ID = '[REDACTED_FOR_SECURITY]'
try:
    ee.Initialize(project=MY_PROJECT_ID)
    print(f"Earth Engine Initialized with project: {MY_PROJECT_ID}")
except:
    ee.Authenticate()
    ee.Initialize(project=MY_PROJECT_ID)

# Configuration
year = 2021
START_DATE = f'{year}-10-16'
END_DATE = f'{year + 1}-04-16'
ASSET_ID = '[REDACTED_FOR_SECURITY]'
SAVE_DIR = '/content/drive/MyDrive/SatMAE_PartialFT_Results_4/'

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

# Assets
wheat_mask = ee.Image(ASSET_ID)
roi = wheat_mask.geometry() # roi is ee.Geometry
s2Bands = ['B2', 'B3', 'B4', 'B5', 'B6', 'B7', 'B8', 'B8A', 'B11', 'B12']
MAX_CLOUD_PROB = 70

# ESRI Crop Mask
TARGET_INDICES = [f'43R_{year}', f'43S_{year}']
esri_lulc = ee.ImageCollection("projects/sat-io/open-datasets/landcover/ESRI_Global-LULC_10m_TS") \
    .filter(ee.Filter.inList('system:index', TARGET_INDICES)).mosaic()
esri_crop_mask = esri_lulc.remap([5], [1], 0).clip(roi)
final_crop_mask = esri_crop_mask.updateMask(esri_crop_mask)

# Sentinel-2 Processing
s2Raw = (ee.ImageCollection('COPERNICUS/S2_SR_HARMONIZED')
         .filterDate(START_DATE, END_DATE).filterBounds(roi)
         .filter(ee.Filter.lt('CLOUDY_PIXEL_PERCENTAGE', 100))
         .map(lambda img: img.updateMask(
             img.select('MSK_SNWPRB').lt(1)
             .add(img.select('SCL').gte(4).And(img.select('SCL').lte(6)))
         )))

s2Clouds = ee.ImageCollection('COPERNICUS/S2_CLOUD_PROBABILITY').filterDate(START_DATE, END_DATE).filterBounds(roi)
s2Joined = ee.ImageCollection(ee.Join.saveFirst('cloud_mask_img').apply(
    s2Raw, s2Clouds, ee.Filter.equals(leftField='system:index', rightField='system:index')))

def processS2(img):
    isCloud = ee.Image(img.get('cloud_mask_img')).select('probability').gt(MAX_CLOUD_PROB)
    optical = img.select(s2Bands).multiply(0.025)
    ndvi = optical.normalizedDifference(['B8', 'B4']).rename('NDVI')
    combined_mask = isCloud.Not().And(final_crop_mask.gt(0))
    return (optical.addBands(ndvi).updateMask(combined_mask).copyProperties(img, ['system:time_start']))

s2Clean = s2Joined.map(processS2)

# Fortnightly Composites
intervals = []
cur = datetime.datetime.fromisoformat(START_DATE)
end = datetime.datetime.fromisoformat(END_DATE.replace('-01-31', '-02-01'))
while cur < end:
    nxt = datetime.datetime(cur.year, cur.month, 16) if cur.day <= 15 else (datetime.datetime(cur.year+1, 1, 1) if cur.month==12 else datetime.datetime(cur.year, cur.month+1, 1))
    if nxt > end: nxt = end
    if cur < nxt: intervals.append([cur.strftime('%Y-%m-%d'), (nxt-datetime.timedelta(days=1)).strftime('%Y-%m-%d'), f"{cur.year}_{cur.month:02d}_{'1' if cur.day<=15 else '2'}"])
    cur = nxt

def makeComposite(item):
    start, end = ee.Date(ee.List(item).get(0)), ee.Date(ee.List(item).get(1)).advance(1,'day')
    return s2Clean.filterDate(start, end).qualityMosaic('NDVI').select(s2Bands).clamp(0, 250).toByte().unmask(255) \
        .set('fortnight_label', ee.List(item).get(2))

fortnightlyCol = ee.ImageCollection.fromImages(ee.List(intervals).map(makeComposite))
print(f"Composites generated: {fortnightlyCol.size().getInfo()}")

# --- 5. REAL LOCAL EXPORT LOGIC ---
def generate_local_dataset():
    print("Starting Local Data Generation (Real Sampling)...")

    # 1. Define Sample Area (Centroid of ROI)
    centroid = roi.centroid()
    buffer_poly = centroid.buffer(2500).bounds() # 2.5km box

    mask_file = 'local_mask.tif'
    print("Downloading Mask...")
    geemap.download_ee_image(wheat_mask, mask_file, region=buffer_poly, scale=10, crs='EPSG:4326', overwrite=True)

    print("Downloading S2 Stack...")
    # Convert Collection to Multi-Band Image (limit 10 time steps)
    stack = fortnightlyCol.limit(10).toBands()
    stack_file = 'local_stack.tif'
    geemap.download_ee_image(stack, stack_file, region=buffer_poly, scale=10, crs='EPSG:4326', overwrite=True)

    # 2. Process into Numpy
    print("Processing Tiles...")
    with rasterio.open(stack_file) as src_stack, rasterio.open(mask_file) as src_mask:
        img_data = src_stack.read()
        mask_data = src_mask.read(1)

        H, W = mask_data.shape
        # Strict Check: Must have enough data for 10 frames
        if img_data.shape[0] < 100:
             raise RuntimeError(f"Insufficient S2 data found. Expected 100 bands (10 frames * 10 bands), got {img_data.shape[0]}. Check Date Range or Cloud Masking.")

        img_reshaped = img_data[:100].reshape(10, 10, H, W)

        # Tiling Loop
        X_list, y_list = [], []
        if H < 224 or W < 224:
             raise RuntimeError(f"Sample area too small ({H}x{W}). Must be > 224x224.")

        for r in range(0, H-224, 112): # Stride 112
            for c in range(0, W-224, 112):
                chip_x = img_reshaped[:, :, r:r+224, c:c+224]
                chip_y = mask_data[r:r+224, c:c+224]

                # Strict Filter: Only keep tiles with >1% Wheat pixels
                if np.mean(chip_y > 0) > 0.01:
                    X_list.append(chip_x)
                    y_list.append(chip_y)

        if len(X_list) == 0:
            raise RuntimeError("No valid tiles found! The sampled 2.5km area contains ZERO wheat pixels. Try a different ROI or check the Mask Asset.")

        X = np.array(X_list).astype(np.float32)
        y = np.array(y_list).astype(np.float32)[:, None, :, :]

        # Save Real Data
        X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=42)
        np.save(os.path.join(SAVE_DIR, 'train_x.npy'), X_train)
        np.save(os.path.join(SAVE_DIR, 'train_y.npy'), y_train)
        np.save(os.path.join(SAVE_DIR, 'val_x.npy'), X_val)
        np.save(os.path.join(SAVE_DIR, 'val_y.npy'), y_val)
        print(f"SUCCESS: Real data saved to {SAVE_DIR}. Train Shapes: {X_train.shape}, {y_train.shape}")

generate_local_dataset()

Requirement already satisfied: geedim in /usr/local/lib/python3.12/dist-packages (2.0.0)
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)
Requirement already satisfied: geemap in /usr/local/lib/python3.12/dist-packages (0.35.3)
Requirement already satisfied: bqplot in /usr/local/lib/python3.12/dist-packages (from geemap) (0.12.45)
Requirement already satisfied: colour in /usr/local/lib/python3.12/dist-packages (from geemap) (0.1.5)
Requirement already satisfied: earthengine-api>=1.0.0 in /usr/local/lib/python3.12/dist-packages (from geemap) (1.5.24)
Requirement already satisfied: eerepr>=0.1.0 in /usr/local/lib/python3.12/dist-packages (from geemap) (0.1.2)
Requirement already satisfied: folium>=0.17.0 in /usr/local/lib/python3.12/dist-packages (from geemap) (0.20.0)
Requirement already satisfied: geocoder in /usr/local/lib/python3.12/dist-packages (from geemap) (1.38.1)
Requirement already satisfied: ipyevents in /usr/local/lib/python3.12/dist-packages (from geemap) (2.0.4)
Requirement already satisfied: ipyfilechooser>=0.6.0 in /usr/local/lib/python3.12/dist-packages (from geemap) (0.6.0)
Requirement already satisfied: ipyleaflet>=0.19.2 in /usr/local/lib/python3.12/dist-packages (from geemap) (0.20.0)
Requirement already satisfied: ipytree in /usr/local/lib/python3.12/dist-packages (from geemap) (0.2.2)
Requirement already satisfied: matplotlib in /usr/local/lib/python3.12/dist-packages (from geemap) (3.10.0)
Requirement already satisfied: numpy in /usr/local/lib/python3.12/dist-packages (from geemap) (2.0.2)
Requirement already satisfied: pandas in /usr/local/lib/python3.12/dist-packages (from geemap) (2.2.2)
Requirement already satisfied: plotly in /usr/local/lib/python3.12/dist-packages (from geemap) (5.24.1)
Requirement already satisfied: pyperclip in /usr/local/lib/python3.12/dist-packages (from geemap) (1.11.0)
Requirement already satisfied: pyshp>=2.3.1 in /usr/local/lib/python3.12/dist-packages (from geemap) (3.0.3)
Requirement already satisfied: python-box in /usr/local/lib/python3.12/dist-packages (from geemap) (7.3.2)
Requirement already satisfied: scooby in /usr/local/lib/python3.12/dist-packages (from geemap) (0.11.0)
Requirement already satisfied: google-cloud-storage in /usr/local/lib/python3.12/dist-packages (from earthengine-api>=1.0.0->geemap) (3.7.0)
Requirement already satisfied: google-api-python-client>=1.12.1 in /usr/local/lib/python3.12/dist-packages (from earthengine-api>=1.0.0->geemap) (2.187.0)
Requirement already satisfied: google-auth>=1.4.1 in /usr/local/lib/python3.12/dist-packages (from earthengine-api>=1.0.0->geemap) (2.43.0)
Requirement already satisfied: google-auth-httplib2>=0.0.3 in /usr/local/lib/python3.12/dist-packages (from earthengine-api>=1.0.0->geemap) (0.3.0)
Requirement already satisfied: httplib2<1dev,>=0.9.2 in /usr/local/lib/python3.12/dist-packages (from earthengine-api>=1.0.0->geemap) (0.31.0)
Requirement already satisfied: requests in /usr/local/lib/python3.12/dist-packages (from earthengine-api>=1.0.0->geemap) (2.32.4)
Requirement already satisfied: branca>=0.6.0 in /usr/local/lib/python3.12/dist-packages (from folium>=0.17.0->geemap) (0.8.2)
Requirement already satisfied: jinja2>=2.9 in /usr/local/lib/python3.12/dist-packages (from folium>=0.17.0->geemap) (3.1.6)
Requirement already satisfied: xyzservices in /usr/local/lib/python3.12/dist-packages (from folium>=0.17.0->geemap) (2025.11.0)
Requirement already satisfied: ipywidgets in /usr/local/lib/python3.12/dist-packages (from ipyfilechooser>=0.6.0->geemap) (7.7.1)
Requirement already satisfied: jupyter-leaflet<0.21,>=0.20 in /usr/local/lib/python3.12/dist-packages (from ipyleaflet>=0.19.2->geemap) (0.20.0)
Requirement already satisfied: traittypes<3,>=0.2.1 in /usr/local/lib/python3.12/dist-packages (from ipyleaflet>=0.19.2->geemap) (0.2.3)
Requirement already satisfied: traitlets>=4.3.0 in /usr/local/lib/python3.12/dist-packages (from bqplot->geemap) (5.7.1)
Requirement already satisfied: python-dateutil>=2.8.2 in /usr/local/lib/python3.12/dist-packages (from pandas->geemap) (2.9.0.post0)
Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.12/dist-packages (from pandas->geemap) (2025.2)
Requirement already satisfied: tzdata>=2022.7 in /usr/local/lib/python3.12/dist-packages (from pandas->geemap) (2025.3)
Requirement already satisfied: click in /usr/local/lib/python3.12/dist-packages (from geocoder->geemap) (8.3.1)
Requirement already satisfied: future in /usr/local/lib/python3.12/dist-packages (from geocoder->geemap) (1.0.0)
Requirement already satisfied: ratelim in /usr/local/lib/python3.12/dist-packages (from geocoder->geemap) (0.1.6)
Requirement already satisfied: six in /usr/local/lib/python3.12/dist-packages (from geocoder->geemap) (1.17.0)
Requirement already satisfied: contourpy>=1.0.1 in /usr/local/lib/python3.12/dist-packages (from matplotlib->geemap) (1.3.3)
Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.12/dist-packages (from matplotlib->geemap) (0.12.1)
Requirement already satisfied: fonttools>=4.22.0 in /usr/local/lib/python3.12/dist-packages (from matplotlib->geemap) (4.61.1)
Requirement already satisfied: kiwisolver>=1.3.1 in /usr/local/lib/python3.12/dist-packages (from matplotlib->geemap) (1.4.9)
Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.12/dist-packages (from matplotlib->geemap) (25.0)
Requirement already satisfied: pillow>=8 in /usr/local/lib/python3.12/dist-packages (from matplotlib->geemap) (11.3.0)
Requirement already satisfied: pyparsing>=2.3.1 in /usr/local/lib/python3.12/dist-packages (from matplotlib->geemap) (3.3.1)
Requirement already satisfied: tenacity>=6.2.0 in /usr/local/lib/python3.12/dist-packages (from plotly->geemap) (9.1.2)
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>=1.0.0->geemap) (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>=1.0.0->geemap) (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>=1.0.0->geemap) (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>=1.0.0->geemap) (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>=1.0.0->geemap) (4.9.1)
Requirement already satisfied: ipykernel>=4.5.1 in /usr/local/lib/python3.12/dist-packages (from ipywidgets->ipyfilechooser>=0.6.0->geemap) (6.17.1)
Requirement already satisfied: ipython-genutils~=0.2.0 in /usr/local/lib/python3.12/dist-packages (from ipywidgets->ipyfilechooser>=0.6.0->geemap) (0.2.0)
Requirement already satisfied: widgetsnbextension~=3.6.0 in /usr/local/lib/python3.12/dist-packages (from ipywidgets->ipyfilechooser>=0.6.0->geemap) (3.6.10)
Requirement already satisfied: ipython>=4.0.0 in /usr/local/lib/python3.12/dist-packages (from ipywidgets->ipyfilechooser>=0.6.0->geemap) (7.34.0)
Requirement already satisfied: jupyterlab-widgets>=1.0.0 in /usr/local/lib/python3.12/dist-packages (from ipywidgets->ipyfilechooser>=0.6.0->geemap) (3.0.16)
Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.12/dist-packages (from jinja2>=2.9->folium>=0.17.0->geemap) (3.0.3)
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>=1.0.0->geemap) (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>=1.0.0->geemap) (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>=1.0.0->geemap) (1.8.0)
Requirement already satisfied: charset_normalizer<4,>=2 in /usr/local/lib/python3.12/dist-packages (from requests->earthengine-api>=1.0.0->geemap) (3.4.4)
Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.12/dist-packages (from requests->earthengine-api>=1.0.0->geemap) (3.11)
Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.12/dist-packages (from requests->earthengine-api>=1.0.0->geemap) (2.5.0)
Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.12/dist-packages (from requests->earthengine-api>=1.0.0->geemap) (2026.1.4)
Requirement already satisfied: decorator in /usr/local/lib/python3.12/dist-packages (from ratelim->geocoder->geemap) (4.4.2)
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>=1.0.0->geemap) (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>=1.0.0->geemap) (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>=1.0.0->geemap) (1.27.0)
Requirement already satisfied: debugpy>=1.0 in /usr/local/lib/python3.12/dist-packages (from ipykernel>=4.5.1->ipywidgets->ipyfilechooser>=0.6.0->geemap) (1.8.15)
Requirement already satisfied: jupyter-client>=6.1.12 in /usr/local/lib/python3.12/dist-packages (from ipykernel>=4.5.1->ipywidgets->ipyfilechooser>=0.6.0->geemap) (7.4.9)
Requirement already satisfied: matplotlib-inline>=0.1 in /usr/local/lib/python3.12/dist-packages (from ipykernel>=4.5.1->ipywidgets->ipyfilechooser>=0.6.0->geemap) (0.2.1)
Requirement already satisfied: nest-asyncio in /usr/local/lib/python3.12/dist-packages (from ipykernel>=4.5.1->ipywidgets->ipyfilechooser>=0.6.0->geemap) (1.6.0)
Requirement already satisfied: psutil in /usr/local/lib/python3.12/dist-packages (from ipykernel>=4.5.1->ipywidgets->ipyfilechooser>=0.6.0->geemap) (5.9.5)
Requirement already satisfied: pyzmq>=17 in /usr/local/lib/python3.12/dist-packages (from ipykernel>=4.5.1->ipywidgets->ipyfilechooser>=0.6.0->geemap) (26.2.1)
Requirement already satisfied: tornado>=6.1 in /usr/local/lib/python3.12/dist-packages (from ipykernel>=4.5.1->ipywidgets->ipyfilechooser>=0.6.0->geemap) (6.5.1)
Requirement already satisfied: setuptools>=18.5 in /usr/local/lib/python3.12/dist-packages (from ipython>=4.0.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (75.2.0)
Requirement already satisfied: jedi>=0.16 in /usr/local/lib/python3.12/dist-packages (from ipython>=4.0.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (0.19.2)
Requirement already satisfied: pickleshare in /usr/local/lib/python3.12/dist-packages (from ipython>=4.0.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (0.7.5)
Requirement already satisfied: prompt-toolkit!=3.0.0,!=3.0.1,<3.1.0,>=2.0.0 in /usr/local/lib/python3.12/dist-packages (from ipython>=4.0.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (3.0.52)
Requirement already satisfied: pygments in /usr/local/lib/python3.12/dist-packages (from ipython>=4.0.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (2.19.2)
Requirement already satisfied: backcall in /usr/local/lib/python3.12/dist-packages (from ipython>=4.0.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (0.2.0)
Requirement already satisfied: pexpect>4.3 in /usr/local/lib/python3.12/dist-packages (from ipython>=4.0.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (4.9.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>=1.0.0->geemap) (0.6.1)
Requirement already satisfied: notebook>=4.4.1 in /usr/local/lib/python3.12/dist-packages (from widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (6.5.7)
Requirement already satisfied: parso<0.9.0,>=0.8.4 in /usr/local/lib/python3.12/dist-packages (from jedi>=0.16->ipython>=4.0.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (0.8.5)
Requirement already satisfied: entrypoints in /usr/local/lib/python3.12/dist-packages (from jupyter-client>=6.1.12->ipykernel>=4.5.1->ipywidgets->ipyfilechooser>=0.6.0->geemap) (0.4)
Requirement already satisfied: jupyter-core>=4.9.2 in /usr/local/lib/python3.12/dist-packages (from jupyter-client>=6.1.12->ipykernel>=4.5.1->ipywidgets->ipyfilechooser>=0.6.0->geemap) (5.9.1)
Requirement already satisfied: argon2-cffi in /usr/local/lib/python3.12/dist-packages (from notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (25.1.0)
Requirement already satisfied: nbformat in /usr/local/lib/python3.12/dist-packages (from notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (5.10.4)
Requirement already satisfied: nbconvert>=5 in /usr/local/lib/python3.12/dist-packages (from notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (7.16.6)
Requirement already satisfied: Send2Trash>=1.8.0 in /usr/local/lib/python3.12/dist-packages (from notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (2.0.0)
Requirement already satisfied: terminado>=0.8.3 in /usr/local/lib/python3.12/dist-packages (from notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (0.18.1)
Requirement already satisfied: prometheus-client in /usr/local/lib/python3.12/dist-packages (from notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (0.23.1)
Requirement already satisfied: nbclassic>=0.4.7 in /usr/local/lib/python3.12/dist-packages (from notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (1.3.3)
Requirement already satisfied: ptyprocess>=0.5 in /usr/local/lib/python3.12/dist-packages (from pexpect>4.3->ipython>=4.0.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (0.7.0)
Requirement already satisfied: wcwidth in /usr/local/lib/python3.12/dist-packages (from prompt-toolkit!=3.0.0,!=3.0.1,<3.1.0,>=2.0.0->ipython>=4.0.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (0.2.14)
Requirement already satisfied: platformdirs>=2.5 in /usr/local/lib/python3.12/dist-packages (from jupyter-core>=4.9.2->jupyter-client>=6.1.12->ipykernel>=4.5.1->ipywidgets->ipyfilechooser>=0.6.0->geemap) (4.5.1)
Requirement already satisfied: notebook-shim>=0.2.3 in /usr/local/lib/python3.12/dist-packages (from nbclassic>=0.4.7->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (0.2.4)
Requirement already satisfied: beautifulsoup4 in /usr/local/lib/python3.12/dist-packages (from nbconvert>=5->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (4.13.5)
Requirement already satisfied: bleach!=5.0.0 in /usr/local/lib/python3.12/dist-packages (from bleach[css]!=5.0.0->nbconvert>=5->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (6.3.0)
Requirement already satisfied: defusedxml in /usr/local/lib/python3.12/dist-packages (from nbconvert>=5->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (0.7.1)
Requirement already satisfied: jupyterlab-pygments in /usr/local/lib/python3.12/dist-packages (from nbconvert>=5->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (0.3.0)
Requirement already satisfied: mistune<4,>=2.0.3 in /usr/local/lib/python3.12/dist-packages (from nbconvert>=5->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (3.2.0)
Requirement already satisfied: nbclient>=0.5.0 in /usr/local/lib/python3.12/dist-packages (from nbconvert>=5->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (0.10.4)
Requirement already satisfied: pandocfilters>=1.4.1 in /usr/local/lib/python3.12/dist-packages (from nbconvert>=5->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (1.5.1)
Requirement already satisfied: fastjsonschema>=2.15 in /usr/local/lib/python3.12/dist-packages (from nbformat->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (2.21.2)
Requirement already satisfied: jsonschema>=2.6 in /usr/local/lib/python3.12/dist-packages (from nbformat->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (4.26.0)
Requirement already satisfied: argon2-cffi-bindings in /usr/local/lib/python3.12/dist-packages (from argon2-cffi->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (25.1.0)
Requirement already satisfied: webencodings in /usr/local/lib/python3.12/dist-packages (from bleach!=5.0.0->bleach[css]!=5.0.0->nbconvert>=5->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (0.5.1)
Requirement already satisfied: tinycss2<1.5,>=1.1.0 in /usr/local/lib/python3.12/dist-packages (from bleach[css]!=5.0.0->nbconvert>=5->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (1.4.0)
Requirement already satisfied: attrs>=22.2.0 in /usr/local/lib/python3.12/dist-packages (from jsonschema>=2.6->nbformat->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (25.4.0)
Requirement already satisfied: jsonschema-specifications>=2023.03.6 in /usr/local/lib/python3.12/dist-packages (from jsonschema>=2.6->nbformat->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (2025.9.1)
Requirement already satisfied: referencing>=0.28.4 in /usr/local/lib/python3.12/dist-packages (from jsonschema>=2.6->nbformat->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (0.37.0)
Requirement already satisfied: rpds-py>=0.25.0 in /usr/local/lib/python3.12/dist-packages (from jsonschema>=2.6->nbformat->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (0.30.0)
Requirement already satisfied: jupyter-server<3,>=1.8 in /usr/local/lib/python3.12/dist-packages (from notebook-shim>=0.2.3->nbclassic>=0.4.7->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (2.14.0)
Requirement already satisfied: cffi>=1.0.1 in /usr/local/lib/python3.12/dist-packages (from argon2-cffi-bindings->argon2-cffi->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (2.0.0)
Requirement already satisfied: soupsieve>1.2 in /usr/local/lib/python3.12/dist-packages (from beautifulsoup4->nbconvert>=5->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (2.8.1)
Requirement already satisfied: typing-extensions>=4.0.0 in /usr/local/lib/python3.12/dist-packages (from beautifulsoup4->nbconvert>=5->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (4.15.0)
Requirement already satisfied: pycparser in /usr/local/lib/python3.12/dist-packages (from cffi>=1.0.1->argon2-cffi-bindings->argon2-cffi->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (2.23)
Requirement already satisfied: anyio>=3.1.0 in /usr/local/lib/python3.12/dist-packages (from jupyter-server<3,>=1.8->notebook-shim>=0.2.3->nbclassic>=0.4.7->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (4.12.1)
Requirement already satisfied: jupyter-events>=0.9.0 in /usr/local/lib/python3.12/dist-packages (from jupyter-server<3,>=1.8->notebook-shim>=0.2.3->nbclassic>=0.4.7->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (0.12.0)
Requirement already satisfied: jupyter-server-terminals>=0.4.4 in /usr/local/lib/python3.12/dist-packages (from jupyter-server<3,>=1.8->notebook-shim>=0.2.3->nbclassic>=0.4.7->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (0.5.3)
Requirement already satisfied: overrides>=5.0 in /usr/local/lib/python3.12/dist-packages (from jupyter-server<3,>=1.8->notebook-shim>=0.2.3->nbclassic>=0.4.7->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (7.7.0)
Requirement already satisfied: websocket-client>=1.7 in /usr/local/lib/python3.12/dist-packages (from jupyter-server<3,>=1.8->notebook-shim>=0.2.3->nbclassic>=0.4.7->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (1.9.0)
Requirement already satisfied: python-json-logger>=2.0.4 in /usr/local/lib/python3.12/dist-packages (from jupyter-events>=0.9.0->jupyter-server<3,>=1.8->notebook-shim>=0.2.3->nbclassic>=0.4.7->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (4.0.0)
Requirement already satisfied: pyyaml>=5.3 in /usr/local/lib/python3.12/dist-packages (from jupyter-events>=0.9.0->jupyter-server<3,>=1.8->notebook-shim>=0.2.3->nbclassic>=0.4.7->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (6.0.3)
Requirement already satisfied: rfc3339-validator in /usr/local/lib/python3.12/dist-packages (from jupyter-events>=0.9.0->jupyter-server<3,>=1.8->notebook-shim>=0.2.3->nbclassic>=0.4.7->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (0.1.4)
Requirement already satisfied: rfc3986-validator>=0.1.1 in /usr/local/lib/python3.12/dist-packages (from jupyter-events>=0.9.0->jupyter-server<3,>=1.8->notebook-shim>=0.2.3->nbclassic>=0.4.7->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (0.1.1)
Requirement already satisfied: fqdn in /usr/local/lib/python3.12/dist-packages (from jsonschema[format-nongpl]>=4.18.0->jupyter-events>=0.9.0->jupyter-server<3,>=1.8->notebook-shim>=0.2.3->nbclassic>=0.4.7->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (1.5.1)
Requirement already satisfied: isoduration in /usr/local/lib/python3.12/dist-packages (from jsonschema[format-nongpl]>=4.18.0->jupyter-events>=0.9.0->jupyter-server<3,>=1.8->notebook-shim>=0.2.3->nbclassic>=0.4.7->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (20.11.0)
Requirement already satisfied: jsonpointer>1.13 in /usr/local/lib/python3.12/dist-packages (from jsonschema[format-nongpl]>=4.18.0->jupyter-events>=0.9.0->jupyter-server<3,>=1.8->notebook-shim>=0.2.3->nbclassic>=0.4.7->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (3.0.0)
Requirement already satisfied: rfc3987-syntax>=1.1.0 in /usr/local/lib/python3.12/dist-packages (from jsonschema[format-nongpl]>=4.18.0->jupyter-events>=0.9.0->jupyter-server<3,>=1.8->notebook-shim>=0.2.3->nbclassic>=0.4.7->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (1.1.0)
Requirement already satisfied: uri-template in /usr/local/lib/python3.12/dist-packages (from jsonschema[format-nongpl]>=4.18.0->jupyter-events>=0.9.0->jupyter-server<3,>=1.8->notebook-shim>=0.2.3->nbclassic>=0.4.7->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (1.3.0)
Requirement already satisfied: webcolors>=24.6.0 in /usr/local/lib/python3.12/dist-packages (from jsonschema[format-nongpl]>=4.18.0->jupyter-events>=0.9.0->jupyter-server<3,>=1.8->notebook-shim>=0.2.3->nbclassic>=0.4.7->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (25.10.0)
Requirement already satisfied: lark>=1.2.2 in /usr/local/lib/python3.12/dist-packages (from rfc3987-syntax>=1.1.0->jsonschema[format-nongpl]>=4.18.0->jupyter-events>=0.9.0->jupyter-server<3,>=1.8->notebook-shim>=0.2.3->nbclassic>=0.4.7->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (1.3.1)
Requirement already satisfied: arrow>=0.15.0 in /usr/local/lib/python3.12/dist-packages (from isoduration->jsonschema[format-nongpl]>=4.18.0->jupyter-events>=0.9.0->jupyter-server<3,>=1.8->notebook-shim>=0.2.3->nbclassic>=0.4.7->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets->ipyfilechooser>=0.6.0->geemap) (1.4.0)
Earth Engine Initialized with project: local-dialect-484618-b9
Composites generated: 12
Starting Local Data Generation (Real Sampling)...
Downloading Mask...

/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/1 tiles [00:00<?]

Downloading S2 Stack...

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

WARNING:googleapiclient.http:Sleeping 1.61 seconds before retry 1 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/local-dialect-484618-b9/thumbnails?fields=name&alt=json, after 429
WARNING:googleapiclient.http:Sleeping 0.92 seconds before retry 1 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/local-dialect-484618-b9/thumbnails?fields=name&alt=json, after 429
WARNING:googleapiclient.http:Sleeping 0.55 seconds before retry 1 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/local-dialect-484618-b9/thumbnails?fields=name&alt=json, after 429
WARNING:googleapiclient.http:Sleeping 1.72 seconds before retry 1 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/local-dialect-484618-b9/thumbnails?fields=name&alt=json, after 429
WARNING:googleapiclient.http:Sleeping 1.02 seconds before retry 1 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/local-dialect-484618-b9/thumbnails?fields=name&alt=json, after 429
WARNING:googleapiclient.http:Sleeping 2.87 seconds before retry 2 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/local-dialect-484618-b9/thumbnails?fields=name&alt=json, after 429
WARNING:googleapiclient.http:Sleeping 1.36 seconds before retry 1 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/local-dialect-484618-b9/thumbnails?fields=name&alt=json, after 429
WARNING:googleapiclient.http:Sleeping 0.27 seconds before retry 1 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/local-dialect-484618-b9/thumbnails?fields=name&alt=json, after 429
WARNING:googleapiclient.http:Sleeping 0.67 seconds before retry 1 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/local-dialect-484618-b9/thumbnails?fields=name&alt=json, after 429
WARNING:googleapiclient.http:Sleeping 1.51 seconds before retry 1 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/local-dialect-484618-b9/thumbnails?fields=name&alt=json, after 429
WARNING:googleapiclient.http:Sleeping 1.99 seconds before retry 1 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/local-dialect-484618-b9/thumbnails?fields=name&alt=json, after 429
WARNING:googleapiclient.http:Sleeping 2.05 seconds before retry 2 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/local-dialect-484618-b9/thumbnails?fields=name&alt=json, after 429
WARNING:googleapiclient.http:Sleeping 3.21 seconds before retry 2 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/local-dialect-484618-b9/thumbnails?fields=name&alt=json, after 429
WARNING:googleapiclient.http:Sleeping 3.90 seconds before retry 2 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/local-dialect-484618-b9/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:googleapiclient.http:Sleeping 0.52 seconds before retry 2 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/local-dialect-484618-b9/thumbnails?fields=name&alt=json, after 429
WARNING:googleapiclient.http:Sleeping 1.65 seconds before retry 2 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/local-dialect-484618-b9/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:rasterio._env:CPLE_AppDefined in 0c5442f3-22c7-497e-ad32-6d879739a2af.tif: TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10
WARNING:googleapiclient.http:Sleeping 2.90 seconds before retry 2 of 5 for request: POST https://earthengine.googleapis.com/v1/projects/local-dialect-484618-b9/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:rasterio._env:CPLE_AppDefined in eaa469a5-f6fd-4a16-a2c7-b14874617c54.tif: TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in dff638fe-da9c-4230-ac61-d7bb8dd1c7b2.tif: TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in ab7628b3-d1de-44b6-a28f-9b55e31e8400.tif: TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in df68cf07-e352-476d-b8ff-cdc0c5bfd323.tif: TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in e23c6b60-5ab6-4281-ace9-f78c71cea481.tif: TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in 809d5bda-95e2-4b74-85d3-99ab2861e8a7.tif: TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in 65c29bb5-989e-4082-9ee3-44d2ff6998aa.tif: TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in 16c9d1c3-1b16-40e8-9d1c-8953fc5dec2c.tif: TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in 1313f3fb-4e37-4bda-bab7-0b6a6c048c1c.tif: TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in 525efda4-a096-40bc-9b2f-166586dc70e0.tif: TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in d4af41ef-9fc7-4860-bbe4-8fd9e36b7321.tif: TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in ec51d8c4-8f09-4f3c-a36c-7859870e2dd4.tif: TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in 45e0f845-8fbe-4487-9719-5e1e5baf775f.tif: TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in a6759be9-0844-42c0-b5ef-5b152fa7d966.tif: TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in f04c5b60-05e3-4bac-b837-62620561c9c2.tif: TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10
WARNING:rasterio._env:CPLE_AppDefined in e273072d-896a-42c8-a8b6-8bb9db2d412f.tif: TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in f411eb3a-3d6f-4c33-9ce4-1e56e15a91f3.tif: TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in 260d9ed2-1ed5-4fa7-b6ae-d2d7cb711a80.tif: TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10
WARNING:rasterio._env:CPLE_AppDefined in b9b86885-e022-41a8-a511-9db08eb5e538.tif: TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10
WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10
WARNING:rasterio._env:CPLE_AppDefined in 71a22aed-52f0-4873-8206-1facd98dc41a.tif: TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10
WARNING:urllib3.connectionpool:Connection pool is full, discarding connection: earthengine.googleapis.com. Connection pool size: 10
WARNING:rasterio._env:CPLE_AppDefined in c51a0e3d-530d-4cf8-a576-d6de0ac2a27d.tif: TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in 49678e55-b806-4c26-9f47-e899bc8d3c3f.tif: TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in 15cfd71d-e31e-48f8-bce3-93f342a9209d.tif: TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in aad13e55-d92d-4720-b83b-66601e0f3f30.tif: TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in bd7a7b62-678b-4828-9c93-52faf562efee.tif: TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in 4a0e8b83-cd53-4b27-a001-de0914d7c5f9.tif: TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in 389d5971-78e1-4aa5-a9bc-01adb206b711.tif: TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in 1edb7c66-e727-46a3-9000-774f1a724771.tif: TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in 3402af5c-e1ab-4e99-a437-3d9cd0aaa190.tif: TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.
WARNING:rasterio._env:CPLE_AppDefined in TIFFReadDirectory:Sum of Photometric type-related color channels and ExtraSamples doesn't match SamplesPerPixel. Defining non-color channels as ExtraSamples.

Processing Tiles...
SUCCESS: Real data saved to /content/drive/MyDrive/SatMAE_PartialFT_Results_4/. Train Shapes: (9, 10, 10, 224, 224), (9, 1, 224, 224)

import os
import torch
import torch.nn as nn
import numpy as np
from torch.utils.data import TensorDataset, DataLoader

# --- 1. DOWNLOAD WEIGHTS (Standard ViT-MAE-Base) ---
def get_satmae_checkpoint():
    print("Downloading Standard ViT-MAE-Base Weights...")
    # Since SatMAE/ScaleMAE mirrors are down, we use the OFFICIAL ViT-MAE-Base
    # from the original MAE paper (Facebook/Meta).
    # This matches the architecture EXACTLY (ViT-Base + MAE).
    # It is trained on ImageNet, but acts as a perfect architectural initialization.

    url = "https://dl.fbaipublicfiles.com/mae/pretrain/mae_pretrain_vit_base.pth"
    checkpoint_path = 'mae_pretrain_vit_base.pth'

    if not os.path.exists(checkpoint_path):
        try:
            torch.hub.download_url_to_file(url, checkpoint_path)
            print(f"Weights downloaded: {checkpoint_path}")
        except Exception as e:
            raise RuntimeError(f"Download failed: {e}")

    return checkpoint_path

# --- 2. MODEL ARCHITECTURE (Standard ViT-Base) ---
class SatMAEPatchEmbed(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_chans=10, embed_dim=768):
        super().__init__()
        self.num_patches = (img_size // patch_size) ** 2
        # Projects 10 bands -> 768 dim
        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
        self.grid_size = img_size // patch_size

    def forward(self, x):
        # Expects 5D: (B, T, C, H, W)
        B, T, C, H, W = x.shape
        x = x.reshape(B * T, C, H, W)
        x = self.proj(x).flatten(2).transpose(1, 2)
        return x.reshape(B, T, -1, 768)

class SatMAEEncoder(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_chans=10, num_frames=10, embed_dim=768, depth=12, num_heads=12):
        super().__init__()
        self.patch_embed = SatMAEPatchEmbed(img_size, patch_size, in_chans, embed_dim)

        # Positional Embeddings
        self.pos_embed_spatial = nn.Parameter(torch.zeros(1, 1, self.patch_embed.num_patches, embed_dim))
        self.pos_embed_temporal = nn.Parameter(torch.zeros(1, num_frames, 1, embed_dim))

        self.blocks = nn.TransformerEncoder(nn.TransformerEncoderLayer(embed_dim, num_heads, int(embed_dim*4), 0.1, 'gelu', batch_first=True), depth)
        self.norm = nn.LayerNorm(embed_dim)

    def forward(self, x):
        # Expects 5D: (B, T, C, H, W)
        B, T, C, H, W = x.shape
        x = self.patch_embed(x)
        x = x + self.pos_embed_spatial + self.pos_embed_temporal
        x = x.reshape(B, T * x.shape[2], -1)
        x = self.norm(self.blocks(x))
        return x.reshape(B, T, 14, 14, -1)

class SatMAESegmentation(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_chans=10, num_frames=10, embed_dim=768):
        super().__init__()
        self.num_frames = num_frames; self.chans = in_chans
        self.encoder = SatMAEEncoder(img_size, patch_size, in_chans, num_frames, embed_dim)

        # Decoder
        self.up1 = nn.ConvTranspose2d(embed_dim, 256, 2, 2)
        self.conv1 = nn.Sequential(nn.Conv2d(256, 256, 3, 1, 1), nn.BatchNorm2d(256), nn.ReLU())
        self.up2 = nn.ConvTranspose2d(256, 128, 2, 2)
        self.conv2 = nn.Sequential(nn.Conv2d(128, 128, 3, 1, 1), nn.BatchNorm2d(128), nn.ReLU())
        self.up3 = nn.ConvTranspose2d(128, 64, 2, 2)
        self.conv3 = nn.Sequential(nn.Conv2d(64, 64, 3, 1, 1), nn.BatchNorm2d(64), nn.ReLU())
        self.up4 = nn.ConvTranspose2d(64, 32, 2, 2)
        self.conv4 = nn.Sequential(nn.Conv2d(32, 32, 3, 1, 1), nn.BatchNorm2d(32), nn.ReLU())
        self.final = nn.Conv2d(32, 1, 1)

    def forward(self, x):
        # --- FIX: Handle 5D Input (B, T, C, H, W) ---
        if x.ndim == 5:
            # Input is already (B, T, C, H, W), use as is
            pass
        elif x.ndim == 4:
            # Fallback if input is flattened (B, C_total, H, W)
            B, _, H, W = x.shape
            x = x.reshape(B, self.num_frames, self.chans, H, W)

        features = self.encoder(x)
        # Collapse Time
        x = features.mean(dim=1).permute(0, 3, 1, 2)
        x = self.conv1(self.up1(x))
        x = self.conv2(self.up2(x))
        x = self.conv3(self.up3(x))
        x = self.conv4(self.up4(x))
        return self.final(x)

def load_satmae_weights(model, checkpoint_path):
    print(f"Loading REAL weights from {checkpoint_path}...")
    state_dict = torch.load(checkpoint_path, map_location='cpu')
    if 'model' in state_dict: state_dict = state_dict['model']

    model_dict = model.state_dict()
    new_state_dict = {}

    for k, v in state_dict.items():
        k_new = k
        if 'blocks' in k and 'layers' not in k: k_new = k.replace('blocks', 'encoder.blocks.layers')
        elif 'patch_embed' in k: k_new = k.replace('patch_embed', 'encoder.patch_embed')
        elif 'norm' in k: k_new = k.replace('norm', 'encoder.norm')
        elif 'pos_embed' in k: k_new = k.replace('pos_embed', 'encoder.pos_embed')

        if not k_new.startswith('encoder.') and 'encoder' in model_dict.keys():
             k_new = 'encoder.' + k_new

        if k_new in model_dict:
            if v.shape == model_dict[k_new].shape:
                new_state_dict[k_new] = v

    model_dict.update(new_state_dict)
    model.load_state_dict(model_dict)
    print(f"Loaded compatible pretrained weights.")

def setup_training_logic(model):
    for p in model.parameters(): p.requires_grad = True
    model.encoder.pos_embed_spatial.requires_grad = False
    model.encoder.pos_embed_temporal.requires_grad = True
    for p in model.encoder.patch_embed.parameters(): p.requires_grad = True
    for i in range(10):
        for p in model.encoder.blocks.layers[i].parameters(): p.requires_grad = False
    for i in range(10, 12):
        for p in model.encoder.blocks.layers[i].parameters(): p.requires_grad = True

    print("Fine-tuning: Spectral=Grad, Spatial=Frozen, Temporal=Grad, Blocks 0-9=Frozen, 10-11=Grad.")

from torch.optim.swa_utils import AveragedModel, SWALR, update_bn
import torch.optim as optim
import torch.nn as nn
import torch.nn.functional as F
import torch
import numpy as np
from scipy.ndimage import distance_transform_edt as distance
from tqdm.auto import tqdm

# --- METRICS ---
def compute_metrics(pred_probs, targets, threshold=0.5):
    pred_mask = (pred_probs > threshold).float()
    intersection = (pred_mask * targets).sum()
    union = pred_mask.sum() + targets.sum() - intersection
    iou = (intersection + 1e-6) / (union + 1e-6)
    dice = (2 * intersection + 1e-6) / (pred_mask.sum() + targets.sum() + 1e-6)
    return iou.item(), dice.item()

class HausdorffDTLoss(nn.Module):
    def __init__(self, alpha=2.0):
        super(HausdorffDTLoss, self).__init__()
        self.alpha = alpha
    def forward(self, pred, gt):
        with torch.no_grad():
            gt_np = gt.detach().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: dist_map[i, 0] = np.ones_like(mask) * 100
                elif mask.sum() == mask.size: dist_map[i, 0] = np.zeros_like(mask)
                else:
                    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)
        error = (probs - gt) ** 2
        return (error * (1 + self.alpha * dist_map)).mean()

class CompoundLoss(nn.Module):
    def __init__(self, dice_weight=0.7, hausdorff_weight=0.3):
        super(CompoundLoss, self).__init__()
        self.hd_loss = HausdorffDTLoss(alpha=2.0)
        self.dice_w = dice_weight; self.hd_w = hausdorff_weight
    def forward(self, inputs, targets):
        probs = torch.sigmoid(inputs)
        inter = (probs * targets).sum()
        dice = 1 - (2. * inter / (probs.sum() + targets.sum() + 1e-6))
        hd = self.hd_loss(inputs, targets)
        return self.dice_w * dice + self.hd_w * hd

class TestTimeAugmentation:
    def __init__(self, model): self.model = model
    def apply(self, x):
        preds = []
        with torch.no_grad():
            preds.append(torch.sigmoid(self.model(x)))
            preds.append(torch.flip(torch.sigmoid(self.model(torch.flip(x, [-1]))), [-1]))
            preds.append(torch.flip(torch.sigmoid(self.model(torch.flip(x, [-2]))), [-2]))
            preds.append(torch.flip(torch.sigmoid(self.model(torch.flip(x, [-2,-1]))), [-2,-1]))
        return torch.stack(preds).mean(dim=0)

class SatMAETrainer:
    def __init__(self, model, loaders, device, lr=1e-4):
        self.model = model; self.loaders = loaders; self.device = device
        self.optimizer = optim.AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=lr, weight_decay=1e-4)
        self.scheduler = optim.lr_scheduler.CosineAnnealingLR(self.optimizer, T_max=150)
        self.swa_model = AveragedModel(model)
        self.swa_scheduler = SWALR(self.optimizer, swa_lr=lr)
        self.criterion = CompoundLoss(dice_weight=0.7, hausdorff_weight=0.3)
        self.tta = TestTimeAugmentation(model)
        # UPDATED: Added val_loss to history
        self.history = {'train_loss': [], 'val_loss': [], 'val_iou': []}

    def fit(self, epochs=200):
        print(f"Starting Training for {epochs} epochs...")
        use_swa = False
        swa_start_epoch = 151
        patience = 30
        no_improv = 0
        best_iou = 0

        for ep in range(epochs):
            # --- TRAINING ---
            self.model.train()
            t_loss = 0
            pbar = tqdm(self.loaders['train'], desc=f"Epoch {ep+1}/{epochs}", leave=False)

            for batch in pbar:
                if isinstance(batch, (list, tuple)):
                    x, y = batch[0], batch[1]
                x, y = x.to(self.device), y.to(self.device)

                self.optimizer.zero_grad()
                logits = self.model(x)
                loss = self.criterion(logits, y)
                loss.backward()
                self.optimizer.step()
                t_loss += loss.item()
                pbar.set_postfix({'loss': f"{loss.item():.4f}"})

            # --- VALIDATION ---
            eval_model = self.swa_model if use_swa else self.model
            eval_model.eval()
            v_loss = 0  # NEW: Track Validation Loss
            v_iou_sum = 0

            with torch.no_grad():
                for batch in self.loaders['val']:
                    if isinstance(batch, (list, tuple)):
                        x, y = batch[0], batch[1]
                    x, y = x.to(self.device), y.to(self.device)

                    # 1. Calc Val Loss (Raw Logits) for stability
                    logits = eval_model(x)
                    v_loss += self.criterion(logits, y).item()

                    # 2. Calc Metrics (TTA Probs) for performance
                    if use_swa:
                         probs = torch.sigmoid(logits)
                    else:
                         probs = self.tta.apply(x)

                    iou, _ = compute_metrics(probs, y)
                    v_iou_sum += iou

            # Averages
            avg_t_loss = t_loss / len(self.loaders['train'])
            if len(self.loaders['val']) > 0:
                avg_v_loss = v_loss / len(self.loaders['val'])
                avg_iou = v_iou_sum / len(self.loaders['val'])
            else:
                avg_v_loss = 0.0
                avg_iou = 0.0

            self.history['train_loss'].append(avg_t_loss)
            self.history['val_loss'].append(avg_v_loss)
            self.history['val_iou'].append(avg_iou)

            # UPDATED PRINT: Now shows Val Loss
            print(f"Ep {ep+1} | T_Loss: {avg_t_loss:.4f} | V_Loss: {avg_v_loss:.4f} | IoU: {avg_iou:.4f} | SWA: {use_swa}")

            # --- CHECKPOINTING ---
            if (ep + 1) % 5 == 0:
                torch.save(self.model.state_dict(), f"checkpoint_ep{ep+1}.pth")

            if not use_swa:
                self.scheduler.step()
                if avg_iou > best_iou:
                    best_iou = avg_iou
                    no_improv = 0
                    torch.save(self.model.state_dict(), 'best_model.pth')
                    print(f"   >>> New Best IoU: {best_iou:.4f}")
                else:
                    no_improv += 1
                    if no_improv >= patience:
                        print(f"   !! No Improvement for {patience} eps. Triggering SWA.")
                        use_swa = True
                        swa_start_epoch = ep + 1
                if ep + 1 >= swa_start_epoch: use_swa = True
            else:
                self.swa_model.update_parameters(self.model)
                self.swa_scheduler.step()

        if use_swa:
            print("Finalizing SWA...")
            update_bn(self.loaders['train'], self.swa_model, device=self.device)
            torch.save(self.swa_model.module.state_dict(), 'final_swa_model.pth')
            print("Training Complete. SWA Model Saved.")

# --- EXECUTION ---
import torch
import numpy as np
import os
import traceback
from torch.utils.data import TensorDataset, DataLoader

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 1. Initialize Model & Load Weights
model = SatMAESegmentation(img_size=224, patch_size=16, in_chans=10, num_frames=10).to(device)

try:
    checkpoint = get_satmae_checkpoint()
    load_satmae_weights(model, checkpoint)
    setup_training_logic(model)
except NameError:
    print("Warning: Helper functions not found. Skipping weight loading.")

# 2. Data Loading
SAVE_DIR = '/content/drive/MyDrive/SatMAE_PartialFT_Results_4/'

try:
    print("Loading data from disk...")
    # Load with mmap to save RAM
    X_train_raw = np.load(os.path.join(SAVE_DIR, 'train_x.npy'), mmap_mode='r')
    y_train_raw = np.load(os.path.join(SAVE_DIR, 'train_y.npy'), mmap_mode='r')
    X_val_raw = np.load(os.path.join(SAVE_DIR, 'val_x.npy'), mmap_mode='r')
    y_val_raw = np.load(os.path.join(SAVE_DIR, 'val_y.npy'), mmap_mode='r')

    print("Processing Data (Scaling & Masking 255)...")

    # FIX 1: Mask INPUTS
    X_train_proc = (np.where(X_train_raw == 255, 0, X_train_raw) / 250.0).copy()
    X_val_proc = (np.where(X_val_raw == 255, 0, X_val_raw) / 250.0).copy()

    # FIX 2: Mask TARGETS
    # Ensure NoData (255) in targets is converted to 0 so Loss doesn't explode
    y_train_proc = np.where(y_train_raw == 255, 0, y_train_raw).copy()
    y_val_proc = np.where(y_val_raw == 255, 0, y_val_raw).copy()

    # Create Datasets
    print("Creating Datasets...")
    train_ds = TensorDataset(torch.from_numpy(X_train_proc).float(), torch.from_numpy(y_train_proc).float())
    val_ds = TensorDataset(torch.from_numpy(X_val_proc).float(), torch.from_numpy(y_val_proc).float())

    # Create Loaders
    loaders = {
        'train': DataLoader(train_ds, batch_size=16, shuffle=True),
        'val': DataLoader(val_ds, batch_size=16, shuffle=False)
    }

    # 3. Start Training
    print("Starting Training...")
    trainer = SatMAETrainer(model=model, loaders=loaders, device=device, lr=1e-4)
    trainer.fit(epochs=200)

except Exception as e:
    traceback.print_exc()
    print(f"An error occurred: {e}")

Downloading Standard ViT-MAE-Base Weights...
Loading REAL weights from mae_pretrain_vit_base.pth...
Loaded compatible pretrained weights.
Fine-tuning: Spectral=Grad, Spatial=Frozen, Temporal=Grad, Blocks 0-9=Frozen, 10-11=Grad.
Loading data from disk...
Processing Data (Scaling & Masking 255)...
Creating Datasets...
Starting Training...
Starting Training for 200 epochs...

Epoch 1/200:   0%|          | 0/1 [00:00<?, ?it/s]

Ep 1 | T_Loss: 1.2444 | V_Loss: 1.1782 | IoU: 0.0000 | SWA: False
   >>> New Best IoU: 0.0000

Epoch 2/200:   0%|          | 0/1 [00:00<?, ?it/s]

Ep 2 | T_Loss: 1.2327 | V_Loss: 1.1792 | IoU: 0.0000 | SWA: False

Epoch 3/200:   0%|          | 0/1 [00:00<?, ?it/s]

Ep 3 | T_Loss: 1.2211 | V_Loss: 1.1801 | IoU: 0.0000 | SWA: False

Epoch 4/200:   0%|          | 0/1 [00:00<?, ?it/s]

Ep 4 | T_Loss: 1.2058 | V_Loss: 1.1808 | IoU: 0.0000 | SWA: False

Epoch 5/200:   0%|          | 0/1 [00:00<?, ?it/s]

import torch
import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import precision_score, recall_score, f1_score, accuracy_score, jaccard_score
from tqdm.auto import tqdm

# --- 1. CONFIGURATION ---
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
MODEL_PATH = 'best_model.pth' # Or 'satmae_swa_final.pth' if SWA finished
THRESHOLD = 0.5 # Threshold to convert probabilities to binary mask

# --- 2. LOAD BEST MODEL ---
# Re-initialize the architecture
model = SatMAESegmentation(img_size=224, patch_size=16, in_chans=10, num_frames=10).to(device)

# Load weights
try:
    state = torch.load(MODEL_PATH, map_location=device)
    model.load_state_dict(state)
    print(f" Loaded weights from {MODEL_PATH}")
except FileNotFoundError:
    print(f" Could not find {MODEL_PATH}. Using random weights (Metrics will be bad).")

model.eval()

# --- 3. METRICS FUNCTION ---
def calculate_comprehensive_metrics(model, loader, device, threshold=0.5):
    print(f"\nComputing In-Depth Metrics (Threshold: {threshold})...")

    # Accumulators for Global Metrics (Pixel-wise)
    all_preds = []
    all_targets = []

    with torch.no_grad():
        for batch in tqdm(loader, desc="Evaluating"):
            if isinstance(batch, (list, tuple)):
                x, y = batch[0], batch[1]
            x, y = x.to(device), y.to(device)

            # Inference
            logits = model(x)
            probs = torch.sigmoid(logits)

            # Flatten and Binarize
            preds = (probs > threshold).cpu().numpy().flatten().astype(np.uint8)
            targets = y.cpu().numpy().flatten().astype(np.uint8)

            all_preds.append(preds)
            all_targets.append(targets)

    # Concatenate all pixels (Huge array, be careful with RAM.
    # If RAM crashes, we need to calculate rolling metrics)
    y_pred = np.concatenate(all_preds)
    y_true = np.concatenate(all_targets)

    # Calculate Metrics
    accuracy = accuracy_score(y_true, y_pred)
    precision = precision_score(y_true, y_pred, zero_division=0)
    recall = recall_score(y_true, y_pred, zero_division=0)
    f1 = f1_score(y_true, y_pred, zero_division=0)
    iou = jaccard_score(y_true, y_pred, zero_division=0)

    print("\n" + "="*40)
    print("       MODEL PERFORMANCE REPORT       ")
    print("="*40)
    print(f"IoU (Jaccard):    {iou:.4f}  <-- KEY METRIC")
    print(f"F1 Score:         {f1:.4f}")
    print(f"Precision:        {precision:.4f}")
    print(f"Recall:           {recall:.4f}")
    print(f"Pixel Accuracy:   {accuracy:.4f}")
    print("="*40)

    return iou, f1, precision, recall

# --- 4. VISUALIZATION FUNCTION ---
def visualize_predictions(model, loader, device, num_samples=3):
    print(f"\nVisualizing {num_samples} Samples...")

    # Get a batch
    batch = next(iter(loader))
    if isinstance(batch, (list, tuple)):
        x, y = batch[0], batch[1]

    x = x.to(device)
    y = y.cpu().numpy() # Ground Truth

    with torch.no_grad():
        logits = model(x)
        probs = torch.sigmoid(logits).cpu().numpy()
        preds = (probs > THRESHOLD).astype(np.uint8)

    # Visualization Loop
    for i in range(min(num_samples, x.shape[0])):
        fig, ax = plt.subplots(1, 4, figsize=(20, 5))


        img_tensor = x[i, 0].cpu().numpy() # (C, H, W)

        # Create RGB Composite
        # Note: Data is 0-1 (scaled).
        # Sometimes S2 is dark, so we multiply by 3 for brightness
        rgb = np.stack([img_tensor[2], img_tensor[1], img_tensor[0]], axis=-1)
        rgb = np.clip(rgb * 3.0, 0, 1)

        ax[0].imshow(rgb)
        ax[0].set_title("Input (RGB - Frame 0)")
        ax[0].axis('off')

        # 2. Ground Truth
        ax[1].imshow(y[i, 0], cmap='gray')
        ax[1].set_title("Ground Truth Mask")
        ax[1].axis('off')

        # 3. Prediction (Probability Map)
        # Use 'jet' or 'plasma' to see confidence (Blue=0, Red=1)
        im = ax[2].imshow(probs[i, 0], cmap='plasma', vmin=0, vmax=1)
        ax[2].set_title("Predicted Probabilities")
        ax[2].axis('off')
        plt.colorbar(im, ax=ax[2], fraction=0.046, pad=0.04)

        # 4. Final Binary Prediction
        ax[3].imshow(preds[i, 0], cmap='gray')
        ax[3].set_title(f"Binary Prediction (>{THRESHOLD})")
        ax[3].axis('off')

        plt.tight_layout()
        plt.show()

# --- 5. RUN ---
# We use loaders['val'] defined in previous cells
if 'loaders' in locals():
    # 1. Calculate Metrics
    calculate_comprehensive_metrics(model, loaders['val'], device, threshold=0.5)

    # 2. Visualize
    visualize_predictions(model, loaders['val'], device, num_samples=4)
else:
    print(" 'loaders' dictionary not found. Run Data Loading cell first.")

Output hidden; open in https://colab.research.google.com to view.

import torch
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from tqdm.auto import tqdm
from sklearn.metrics import confusion_matrix

# --- CONFIGURATION ---
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
MODEL_PATH = 'best_model.pth'  # Checks for your best saved weights
THRESHOLD = 0.5                # Probability threshold for binary mask

# --- 1. MEMORY-SAFE METRIC CALCULATOR ---
def evaluate_in_depth(model, loader, device, threshold=0.5):
    print(f"Starting In-Depth Evaluation (Threshold: {threshold})...")
    model.eval()

    # Initialize Counters (Long Integers to prevent overflow)
    # TP: True Positive (Wheat predicted as Wheat)
    # FP: False Positive (Background predicted as Wheat)
    # TN: True Negative (Background predicted as Background)
    # FN: False Negative (Wheat predicted as Background)
    total_tp = 0
    total_fp = 0
    total_tn = 0
    total_fn = 0

    with torch.no_grad():
        for batch in tqdm(loader, desc="Evaluating Batches"):
            # Robust Unpacking
            if isinstance(batch, (list, tuple)):
                x, y = batch[0], batch[1]
            x, y = x.to(device), y.to(device)

            # Forward Pass
            logits = model(x)
            probs = torch.sigmoid(logits)
            preds = (probs > threshold).float()

            # --- FAST GPU CALCULATION ---
            # We calculate metrics on the GPU to save CPU transfer time

            # Flatten tensors for metric calculation
            preds_flat = preds.view(-1)
            targets_flat = y.view(-1)

            # Calculate basics
            tp = (preds_flat * targets_flat).sum().item()
            tn = ((1 - preds_flat) * (1 - targets_flat)).sum().item()
            fp = (preds_flat * (1 - targets_flat)).sum().item()
            fn = ((1 - preds_flat) * targets_flat).sum().item()

            # Accumulate
            total_tp += tp
            total_fp += fp
            total_tn += tn
            total_fn += fn

    # --- FINAL METRIC CALCULATION ---
    # Epsilon to prevent division by zero
    eps = 1e-6

    # 1. IoU (Jaccard Index) = TP / (TP + FP + FN)
    iou = total_tp / (total_tp + total_fp + total_fn + eps)

    # 2. Precision = TP / (TP + FP)  (How trustworthy is a 'Wheat' prediction?)
    precision = total_tp / (total_tp + total_fp + eps)

    # 3. Recall (Sensitivity) = TP / (TP + FN)  (How much of the Wheat did we find?)
    recall = total_tp / (total_tp + total_fn + eps)

    # 4. F1 Score = 2 * (Precision * Recall) / (Precision + Recall)
    f1 = 2 * (precision * recall) / (precision + recall + eps)

    # 5. Accuracy = (TP + TN) / Total
    accuracy = (total_tp + total_tn) / (total_tp + total_tn + total_fp + total_fn + eps)

    # 6. Specificity = TN / (TN + FP) (How good is it at ignoring background?)
    specificity = total_tn / (total_tn + total_fp + eps)

    print("\n" + "="*45)
    print("        FINAL MODEL PERFORMANCE REPORT  ")
    print("="*45)
    print(f"IoU (Jaccard Index):  {iou:.4f}  <-- MAIN METRIC")
    print(f"F1 Score (Dice):      {f1:.4f}")
    print("-" * 45)
    print(f"Precision:            {precision:.4f}")
    print(f"Recall:               {recall:.4f}")
    print(f"Specificity:          {specificity:.4f}")
    print(f"Global Accuracy:      {accuracy:.4f}")
    print("="*45)
    print(f"Confusion Matrix Counts:")
    print(f"  True Positives (Wheat):  {int(total_tp):,}")
    print(f"  False Positives (Noise): {int(total_fp):,}")
    print(f"  False Negatives (Miss):  {int(total_fn):,}")
    print(f"  True Negatives (Bg):     {int(total_tn):,}")
    print("="*45 + "\n")

    return iou, f1, precision, recall

# --- 2. VISUALIZATION FUNCTION ---
def visualize_sample_results(model, loader, device, num_samples=3):
    print(f"Visualizing {num_samples} random samples...")
    model.eval()

    # Get a batch
    batch = next(iter(loader))
    if isinstance(batch, (list, tuple)):
        x, y = batch[0], batch[1]
    x, y = x.to(device), y.to(device)

    with torch.no_grad():
        logits = model(x)
        probs = torch.sigmoid(logits)
        preds = (probs > THRESHOLD).float()

    # Move to CPU for plotting
    x_np = x.cpu().numpy()
    y_np = y.cpu().numpy()
    probs_np = probs.cpu().numpy()
    preds_np = preds.cpu().numpy()

    for i in range(min(num_samples, x.shape[0])):
        fig, ax = plt.subplots(1, 4, figsize=(20, 5))

        # A. Input RGB (Frame 0, Bands B4, B3, B2 -> Indices 2, 1, 0)
        # Assuming typical S2 scaling. If image is dark, we multiply by 3.
        rgb = np.stack([x_np[i, 0, 2], x_np[i, 0, 1], x_np[i, 0, 0]], axis=-1)
        rgb = np.clip(rgb * 3.5, 0, 1) # Brightness boost

        ax[0].imshow(rgb)
        ax[0].set_title("Input (RGB Composite)")
        ax[0].axis('off')

        # B. Ground Truth
        ax[1].imshow(y_np[i, 0], cmap='gray', interpolation='nearest')
        ax[1].set_title("Ground Truth Mask")
        ax[1].axis('off')

        # C. Probability Heatmap
        im = ax[2].imshow(probs_np[i, 0], cmap='turbo', vmin=0, vmax=1)
        ax[2].set_title("Model Confidence (Prob)")
        ax[2].axis('off')
        plt.colorbar(im, ax=ax[2], fraction=0.046, pad=0.04)

        # D. Prediction
        ax[3].imshow(preds_np[i, 0], cmap='gray', interpolation='nearest')
        ax[3].set_title(f"Prediction (>{THRESHOLD})")
        ax[3].axis('off')

        plt.tight_layout()
        plt.show()

# --- 3. EXECUTION BLOCK ---
# 1. Re-Init Model
model = SatMAESegmentation(img_size=224, patch_size=16, in_chans=10, num_frames=10).to(device)

# 2. Load Weights
try:
    if os.path.exists(MODEL_PATH):
        state_dict = torch.load(MODEL_PATH, map_location=device)
        model.load_state_dict(state_dict)
        print(f"Successfully loaded weights from {MODEL_PATH}")
    else:
        print(f"{MODEL_PATH} not found. Using current model weights (Evaluation might be poor if not trained).")
except Exception as e:
    print(f"Error loading weights: {e}")

# 3. Check if loaders exist
if 'loaders' in locals() and 'val' in loaders:
    # Run Evaluation
    evaluate_in_depth(model, loaders['val'], device, threshold=0.5)
    # Run Visualization
    visualize_sample_results(model, loaders['val'], device, num_samples=5)
else:
    print(" Loaders not found. Please run the Data Loading cell first.")

Output hidden; open in https://colab.research.google.com to view.