Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+Methane_benchmark_patches_summary_v3.xlsx filter=lfs diff=lfs merge=lfs -text

Methane_benchmark_patches_summary_v3.xlsx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:522e4342850c50e22decd218f5bd546493931f57cd10642cf6d42d2224ad4890
+size 349508

README.md

@@ -0,0 +1,27 @@
+# FAST-EO Use Case 2 - Methane Detection
+
+This directory contains all data and code necessary to recreate experiments conducted for fine-tuning Terramind-Base to detect methane in satellite images. It includes five distinct experiments along with their corresponding datasets. The attached Methane_benchmark_patches_summary_v3.xlsx file provides descriptions for every patch extracted from the Methane Benchmark Dataset (MBD) and defines the fold splits to ensure non-overlapping data. This Excel file is used by the runner scripts to partition the data, typically reserving one fold for testing.
+
+Each script includes usage instructions which can be accessed by applying the --help (or -h) flag.
+
+## Important: Ensure the Terramind package is installed before running any experiments.
+
+## Experiment 1: Fine tuning on Methane Benchmark Dataset
+
+The first experiment is fine tuning the model on Methane Benchmark Dataset. The dataset has been attatched in the directory `MBD_nan_S2_zscore`, and has been already normalized. The code for running the training is located in the `classification` directory, along with neccessary `dataset` and `dataloader` classes. 
+
+## Experiment 2: Fine tuning on MBD with text captions
+
+This experiment contains a modified verion of the Terramind-Based model, which concatinates the textual embeddings of the text captions for every image, with the visual embeddings of the base model. The text embeddings are calculated using the `all-MiniLM-L6-v2` model. All the code, along with embeddings calculation, and data, is available in the `classification_with_text` directory. The original captions are located in `classification_with_text/MBD_text`, and the embeddings are located inside the `combined_caption_embeddings.csv` file.
+
+## Experiment 3: Fine tuning and inference on Sentinel 2 with simulated atmospheric conditions
+
+This experiment checks how the Terramind-Base behaves on the Sentinel-2 data with simulated atmospheric conditions. The simulated data is both in the Top-of-Atmosphere and Bottom-of-Atmsphere variants. The model can be both trained on this data, or only run on it, to test how good it is at generalization when trained on different data.
+
+## Experiment 4: Fine tuning and inference on Intuition 1 with simulated atmosphric conditions
+
+This experiment checks how the Terramind-Base behaves on the Intuition-1 data with simulated atmospheric conditions. The simulated data is both in the Top-of-Atmosphere and Bottom-of-Atmsphere variants. The model can be both trained on this data, or only run on it, to test how good it is at generalization when trained on different data.
+
+## Experiment 5: Testing the detector on urban dataset without methane
+
+The urban dataset has been prepared to check whether the models really learned to detect methane from multispectral data or just look for urban signatures in the images. All of the images in this dataset do not contain methane, the goal is to run the models and see how many false positives are returned. Python script for loading and running the models were attatched.

classification/config/methane_classification_datamodule.py

@@ -0,0 +1,116 @@
+import pandas as pd
+import albumentations as A
+from typing import Optional, List
+from sklearn.model_selection import train_test_split
+from torch.utils.data import DataLoader
+from torchgeo.datamodules import NonGeoDataModule
+from methane_classification_dataset import MethaneClassificationDataset
+
+class MethaneClassificationDataModule(NonGeoDataModule):
+    def __init__(
+        self,
+        data_root: str,
+        excel_file: str,
+        batch_size: int = 8,
+        num_workers: int = 0,
+        val_split: float = 0.2,
+        seed: int = 42,
+        **kwargs
+    ):
+        # We pass "NonGeoDataset" just to satisfy the parent class, 
+        # but we instantiate specific datasets in setup()
+        super().__init__(MethaneClassificationDataset, batch_size, num_workers, **kwargs)
+
+        self.data_root = data_root
+        self.excel_file = excel_file
+        self.val_split = val_split
+        self.seed = seed
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        
+        # State variables for paths
+        self.train_paths = []
+        self.val_paths = []
+
+    def _get_training_transforms(self):
+        """Internal definition of training transforms"""
+        return A.Compose([
+            A.ElasticTransform(p=0.25),
+            A.RandomRotate90(p=0.5),
+            A.Flip(p=0.5),
+            A.ShiftScaleRotate(rotate_limit=90, shift_limit_x=0.05, shift_limit_y=0.05, p=0.5)
+        ])
+
+    def setup(self, stage: str = None):
+        # 1. Read the Excel File
+        try:
+            df = pd.read_csv(self.excel_file) if self.excel_file.endswith('.csv') else pd.read_excel(self.excel_file)
+        except Exception as e:
+            raise RuntimeError(f"Failed to load summary file: {e}")
+
+        # 2. Filter valid paths (checking if Fold column exists or just using all data)
+        # Assuming we just use all data in the file and split it 80/20 here.
+        # If you need specific Fold filtering, add that logic here.
+        all_paths = df['Filename'].tolist()
+
+        # 3. Perform the Split
+        self.train_paths, self.val_paths = train_test_split(
+            all_paths, 
+            test_size=self.val_split, 
+            random_state=self.seed
+        )
+
+        # 4. Instantiate Datasets
+        if stage in ("fit", "train"):
+            self.train_dataset = MethaneClassificationDataset(
+                root_dir=self.data_root,
+                excel_file=self.excel_file,
+                paths=self.train_paths,
+                transform=self._get_training_transforms(),
+            )
+        
+        if stage in ("fit", "validate", "val"):
+            self.val_dataset = MethaneClassificationDataset(
+                root_dir=self.data_root,
+                excel_file=self.excel_file,
+                paths=self.val_paths,
+                transform=None, # No transforms for validation
+            )
+
+        if stage in ("test", "predict"):
+            # For testing, you might want to use a specific hold-out set
+            # For now, reusing val_paths or you can add logic to load a test fold
+            self.test_dataset = MethaneClassificationDataset(
+                root_dir=self.data_root,
+                excel_file=self.excel_file,
+                paths=self.val_paths,
+                transform=None,
+            )
+
+
+    def train_dataloader(self):
+        return DataLoader(
+            self.train_dataset, 
+            batch_size=self.batch_size, 
+            shuffle=True, 
+            num_workers=self.num_workers, 
+            drop_last=True
+        )
+
+    def val_dataloader(self):
+        return DataLoader(
+            self.val_dataset, 
+            batch_size=self.batch_size, 
+            shuffle=False, 
+            num_workers=self.num_workers, 
+            drop_last=True
+        )
+
+    def test_dataloader(self):
+        return DataLoader(
+            self.test_dataset, 
+            batch_size=self.batch_size, 
+            shuffle=False, 
+            num_workers=self.num_workers, 
+            drop_last=True
+        )

classification/config/methane_classification_dataset.py

@@ -0,0 +1,79 @@
+import os
+import rasterio
+import torch
+from torchgeo.datasets import NonGeoDataset
+from torch.utils.data import DataLoader
+import torch.nn.functional as F
+import numpy as np
+import pandas as pd
+
+def min_max_normalize(data, new_min=0, new_max=1):
+    data = np.array(data, dtype=np.float32)  # Convert to NumPy array
+    
+    # Handle NaN, Inf values
+    data = np.nan_to_num(data, nan=np.nanmin(data), posinf=np.max(data), neginf=np.min(data))
+
+    old_min, old_max = np.min(data), np.max(data)
+    
+    if old_max == old_min:  # Prevent division by zero
+        return np.full_like(data, new_min, dtype=np.float32)  # Uniform array
+
+    return (data - old_min) / (old_max - old_min + 1e-10) * (new_max - new_min) + new_min
+
+class MethaneClassificationDataset(NonGeoDataset):
+    def __init__(self, root_dir, excel_file, paths, transform=None, mean=None, std=None):
+        super().__init__()
+        self.root_dir = root_dir
+        self.transform = transform
+        self.data_paths = []
+        self.mean = mean if mean else [0.485] * 12  # Default mean if not provided
+        self.std = std if std else [0.229] * 12  # Default std if not provided
+
+        # Collect paths for labelbinary.tif and sCube.tif in selected folders
+        for folder_name in paths:
+            subdir_path = os.path.join(root_dir, folder_name)
+            if os.path.isdir(subdir_path):
+                label_path = os.path.join(subdir_path, 'labelbinary.tif')
+                scube_path = os.path.join(subdir_path, 'sCube.tif')
+
+                if os.path.exists(label_path) and os.path.exists(scube_path):
+                    self.data_paths.append((label_path, scube_path))
+
+    def __len__(self):
+        return len(self.data_paths)
+
+    def __getitem__(self, idx):
+        label_path, scube_path = self.data_paths[idx]
+
+        # Load the label image (single band)
+        with rasterio.open(label_path) as label_src:
+            label_image = label_src.read(1)  # Shape: [512, 512]
+
+        # Load the sCube image (multi-band), drop the first band
+        with rasterio.open(scube_path) as scube_src:
+            scube_image = scube_src.read()  # Shape: [13, 512, 512]
+            
+            scube_image = scube_image[[0,1,2,3,4,5,6,7,8,9,11,12], :, :]  # Drop first band → Shape: [12, 512, 512]
+
+        # Convert to PyTorch tensors
+        scube_tensor = torch.from_numpy(scube_image).float()  # Shape: [12, 512, 512]
+        label_tensor = torch.from_numpy(label_image).float()  # Shape: [512, 512]
+
+        # Resize to [12, 224, 224] and [224, 224] respectively
+        scube_tensor = F.interpolate(scube_tensor.unsqueeze(0), size=(224, 224), mode='bilinear', align_corners=False).squeeze(0)
+        label_tensor = F.interpolate(label_tensor.unsqueeze(0).unsqueeze(0), size=(224, 224), mode='nearest').squeeze(0)
+
+        label_tensor = label_tensor.clip(0, 1)  # Clip values to [0, 1]
+        scube_tensor = torch.nan_to_num(scube_tensor, nan=0.0)  # Replace NaNs with 0
+
+        # Convert labels to binary
+        contains_methane = (label_tensor > 0).any().long()
+        
+        # Apply transformations (if any)
+        if self.transform:
+            transformed = self.transform(image=np.array(scube_tensor.permute(1, 2, 0)))
+            scube_tensor = transformed['image'].transpose(2, 0, 1)   # Convert back to [C, H, W]
+            
+
+        return {'image': scube_tensor, 'label': contains_methane, 'gt': label_image, 'sample': scube_path.split('/')[3]}
+

classification/config/train.yaml

@@ -0,0 +1,85 @@
+# Global Seed for reproducibility (matches script's set_seed)
+seed_everything: 42
+
+# ------------------------------------------------------------------
+# Trainer Configuration
+# ------------------------------------------------------------------
+trainer:
+  accelerator: auto       # Handles "cuda" if available, else "cpu"
+  strategy: auto
+  devices: 1
+  max_epochs: 100         # Matches args.epochs
+  default_root_dir: ./checkpoints
+  
+  # Callbacks to replicate the script's checkpointing and logging
+  callbacks:
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        monitor: val/loss
+        mode: min
+        save_top_k: 1
+        filename: "best_model"
+        save_last: true   # Saves 'last.ckpt' (similar to 'final_model.pth')
+        
+    - class_path: lightning.pytorch.callbacks.LearningRateMonitor
+      init_args:
+        logging_interval: epoch
+
+# ------------------------------------------------------------------
+# Model Configuration (TerraMind + UperNet)
+# ------------------------------------------------------------------
+model:
+  class_path: terratorch.tasks.ClassificationTask
+  init_args:
+    model_factory: EncoderDecoderFactory
+    loss: ce
+    ignore_index: -1
+    lr: 1.0e-5
+    # Optimizer settings matching _init_optimizer
+    optimizer: AdamW
+    optimizer_hparams:
+      weight_decay: 0.05
+      
+    # Scheduler settings matching ReduceLROnPlateau
+    scheduler: ReduceLROnPlateau
+    scheduler_hparams:
+      mode: min
+      patience: 5
+
+    # --------------------------------------------------------------
+    # Model Architecture (Exact match to script's model_config)
+    # --------------------------------------------------------------
+    model_args:
+      backbone: terramind_v1_base
+      backbone_pretrained: true
+      backbone_modalities:
+        - S2L2A
+      backbone_merge_method: mean
+      
+      decoder: UperNetDecoder
+      decoder_scale_modules: true
+      decoder_channels: 256
+      num_classes: 2
+      head_dropout: 0.3
+      
+      # Specific neck configuration for TerraMind
+      necks:
+        - name: ReshapeTokensToImage
+          remove_cls_token: false
+        - name: SelectIndices
+          indices: [2, 5, 8, 11]
+
+# ------------------------------------------------------------------
+# Data Configuration
+# ------------------------------------------------------------------
+data:
+  class_path: methane_classification_datamodule.MethaneClassificationDataModule
+  init_args:
+    data_root: ../../MBD_nan_S2_zscore/MBD_nan_S2_zscore
+    excel_file: ../../Methane_benchmark_patches_summary_v3.xlsx
+    batch_size: 8
+    val_split: 0.2
+    seed: 42
+    # Note: The procedural train_test_split logic from the script 
+    # (handling folds/splitting) should be encapsulated inside the 
+    # DataModule's setup() method for this config to work seamlessly.

classification/script/methane_classification_datamodule.py

@@ -0,0 +1,71 @@
+import os
+import rasterio
+import torch
+from torchgeo.datasets import NonGeoDataset
+from torch.utils.data import DataLoader
+from torchgeo.datamodules import NonGeoDataModule
+from methane_classification_dataset import MethaneClassificationDataset
+
+class MethaneClassificationDataModule(NonGeoDataModule):
+    """
+    A DataModule for handling MethaneClassificationDataset
+    """
+
+    def __init__(
+        self,
+        data_root: str,
+        excel_file: str,
+        paths: list,
+        batch_size: int = 8,
+        num_workers: int = 0,
+        train_transform: callable = None,
+        val_transform: callable = None,
+        test_transform: callable = None,
+        **kwargs
+    ):
+        super().__init__(MethaneClassificationDataset, batch_size, num_workers, **kwargs)
+
+        self.data_root = data_root
+        self.excel_file = excel_file
+        self.paths = paths
+        self.train_transform = train_transform
+        self.val_transform = val_transform
+        self.test_transform = test_transform
+
+    def setup(self, stage: str = None):
+        if stage in ("fit", "train"):
+            self.train_dataset = MethaneClassificationDataset(
+                root_dir=self.data_root,
+                excel_file=self.excel_file,
+                paths=self.paths,
+                transform=self.train_transform,
+            )
+        if stage in ("fit", "validate", "val"):
+            self.val_dataset = MethaneClassificationDataset(
+                root_dir=self.data_root,
+                excel_file=self.excel_file,
+                paths=self.paths,
+                transform=self.val_transform,
+            )
+        if stage in ("test", "predict"):
+            self.test_dataset = MethaneClassificationDataset(
+                root_dir=self.data_root,
+                excel_file=self.excel_file,
+                paths=self.paths,
+                transform=self.test_transform,
+            )
+
+    def train_dataloader(self):
+        return DataLoader(
+            self.train_dataset, batch_size=self.batch_size, shuffle=True, num_workers=self.num_workers, drop_last=True
+        )
+
+    def val_dataloader(self):
+        return DataLoader(
+            self.val_dataset, batch_size=self.batch_size, shuffle=False, num_workers=self.num_workers, drop_last=True
+        )
+
+    def test_dataloader(self):
+        return DataLoader(
+            self.test_dataset, batch_size=self.batch_size, shuffle=False, num_workers=self.num_workers, drop_last=True
+        )

classification/script/methane_classification_dataset.py

@@ -0,0 +1,82 @@
+import os
+import rasterio
+import torch
+from torchgeo.datasets import NonGeoDataset
+from torch.utils.data import DataLoader
+import torch.nn.functional as F
+import numpy as np
+import pandas as pd
+
+def min_max_normalize(data, new_min=0, new_max=1):
+    data = np.array(data, dtype=np.float32)  # Convert to NumPy array
+    
+    # Handle NaN, Inf values
+    data = np.nan_to_num(data, nan=np.nanmin(data), posinf=np.max(data), neginf=np.min(data))
+
+    old_min, old_max = np.min(data), np.max(data)
+    
+    if old_max == old_min:  # Prevent division by zero
+        return np.full_like(data, new_min, dtype=np.float32)  # Uniform array
+
+    return (data - old_min) / (old_max - old_min + 1e-10) * (new_max - new_min) + new_min
+
+class MethaneClassificationDataset(NonGeoDataset):
+    def __init__(self, root_dir, excel_file, paths, transform=None, mean=None, std=None):
+        super().__init__()
+        self.root_dir = root_dir
+        self.transform = transform
+        self.data_paths = []
+        self.mean = mean if mean else [0.485] * 12  # Default mean if not provided
+        self.std = std if std else [0.229] * 12  # Default std if not provided
+
+        # Collect paths for labelbinary.tif and sCube.tif in selected folders
+        for folder_name in paths:
+            subdir_path = os.path.join(root_dir, folder_name)
+            if os.path.isdir(subdir_path):
+                label_path = os.path.join(subdir_path, 'labelbinary.tif')
+                scube_path = os.path.join(subdir_path, 'sCube.tif')
+
+                if os.path.exists(label_path) and os.path.exists(scube_path):
+                    self.data_paths.append((label_path, scube_path))
+
+    def __len__(self):
+        return len(self.data_paths)
+
+    def __getitem__(self, idx):
+        label_path, scube_path = self.data_paths[idx]
+
+        # Load the label image (single band)
+        with rasterio.open(label_path) as label_src:
+            label_image = label_src.read(1)  # Shape: [512, 512]
+
+        # Load the sCube image (multi-band), drop the first band
+        with rasterio.open(scube_path) as scube_src:
+            scube_image = scube_src.read()  # Shape: [13, 512, 512]
+            # Zrobić tak żeby nie było 10 bandu
+            
+            scube_image = scube_image[[0,1,2,3,4,5,6,7,8,9,11,12], :, :]  # Drop first band → Shape: [12, 512, 512]
+
+        # Convert to PyTorch tensors
+        scube_tensor = torch.from_numpy(scube_image).float()  # Shape: [12, 512, 512]
+        label_tensor = torch.from_numpy(label_image).float()  # Shape: [512, 512]
+
+        # Resize to [12, 224, 224] and [224, 224] respectively
+        scube_tensor = F.interpolate(scube_tensor.unsqueeze(0), size=(224, 224), mode='bilinear', align_corners=False).squeeze(0)
+        label_tensor = F.interpolate(label_tensor.unsqueeze(0).unsqueeze(0), size=(224, 224), mode='nearest').squeeze(0)
+
+        label_tensor = label_tensor.clip(0, 1)  # Clip values to [0, 1]
+        scube_tensor = torch.nan_to_num(scube_tensor, nan=0.0)  # Replace NaNs with 0
+
+        # Convert labels to binary
+        contains_methane = (label_tensor > 0).any().long()
+
+        # Convert to one-hot encoding
+        one_hot_label = F.one_hot(contains_methane, num_classes=2).float()
+        
+        # Apply transformations (if any)
+        if self.transform:
+            transformed = self.transform(image=np.array(scube_tensor.permute(1, 2, 0)))
+            scube_tensor = transformed['image'].transpose(2, 0, 1)   # Convert back to [C, H, W]
+            
+
+        return {'S2L2A': scube_tensor, 'label': one_hot_label, 'gt': label_image, 'sample': scube_path.split('/')[3]}

classification/script/train_classification_fine_tuning.py

@@ -0,0 +1,329 @@
+import argparse
+import logging
+import csv
+import random
+import warnings
+import time
+from pathlib import Path
+from typing import Dict, List, Tuple, Any, Optional
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import albumentations as A
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+from sklearn.model_selection import train_test_split
+from sklearn.metrics import (
+    accuracy_score, recall_score, f1_score, matthews_corrcoef, confusion_matrix
+)
+from rasterio.errors import NotGeoreferencedWarning
+import terramind
+
+# Local Imports
+from methane_classification_datamodule import MethaneClassificationDataModule
+
+# TerraTorch Imports
+from terratorch.tasks import ClassificationTask
+
+
+# --- Configuration & Setup ---
+
+# Configure Logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    datefmt='%Y-%m-%d %H:%M:%S'
+)
+logger = logging.getLogger(__name__)
+
+# Suppress Warnings
+logging.getLogger("rasterio._env").setLevel(logging.ERROR)
+warnings.simplefilter("ignore", NotGeoreferencedWarning)
+warnings.filterwarnings("ignore", category=FutureWarning)
+
+def set_seed(seed: int = 42):
+    """Sets the seed for reproducibility across random, numpy, and torch."""
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+def get_training_transforms() -> A.Compose:
+    """Returns the albumentations training pipeline."""
+    return A.Compose([
+        A.ElasticTransform(p=0.25),
+        A.RandomRotate90(p=0.5),
+        A.Flip(p=0.5),
+        A.ShiftScaleRotate(rotate_limit=90, shift_limit_x=0.05, shift_limit_y=0.05, p=0.5)
+    ])
+
+# --- Helper Classes ---
+
+class MetricTracker:
+    """Accumulates targets and predictions to calculate epoch-level metrics."""
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.all_targets = []
+        self.all_predictions = []
+        self.total_loss = 0.0
+        self.steps = 0
+
+    def update(self, loss: float, targets: torch.Tensor, probabilities: torch.Tensor):
+        self.total_loss += loss
+        self.steps += 1
+        # Store detached cpu numpy arrays to avoid VRAM leaks
+        self.all_targets.extend(torch.argmax(targets, dim=1).detach().cpu().numpy())
+        self.all_predictions.extend(torch.argmax(probabilities, dim=1).detach().cpu().numpy())
+
+    def compute(self) -> Dict[str, float]:
+        """Calculates aggregate metrics for the accumulated data."""
+        if not self.all_targets:
+            return {}
+        
+        # Calculate Confusion Matrix elements
+        tn, fp, fn, tp = confusion_matrix(self.all_targets, self.all_predictions, labels=[0, 1]).ravel()
+        
+        return {
+            "Loss": self.total_loss / max(self.steps, 1),
+            "Accuracy": accuracy_score(self.all_targets, self.all_predictions),
+            "Specificity": tn / (tn + fp) if (tn + fp) != 0 else 0.0,
+            "Sensitivity": recall_score(self.all_targets, self.all_predictions, average='binary', pos_label=1, zero_division=0),
+            "F1": f1_score(self.all_targets, self.all_predictions, average='binary', pos_label=1, zero_division=0),
+            "MCC": matthews_corrcoef(self.all_targets, self.all_predictions),
+        }
+
+class MethaneTrainer:
+    """
+    Handles the training lifecycle: Model setup, Training loop, Validation, and Checkpointing.
+    """
+    def __init__(self, args: argparse.Namespace):
+        self.args = args
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.save_dir = Path(args.save_dir) / f'fold{args.test_fold}'
+        self.save_dir.mkdir(parents=True, exist_ok=True)
+        
+        self.model = self._init_model()
+        self.optimizer, self.scheduler = self._init_optimizer()
+        self.criterion = self.task.criterion  # Retrieved from the TerraTorch task
+        
+        self.best_val_loss = float('inf')
+        
+        logger.info(f"Trainer initialized on device: {self.device}")
+
+    def _init_model(self) -> nn.Module:
+        """Initializes the TerraTorch Classification Task and Model."""
+        model_config = dict(
+            backbone="terramind_v1_base",
+            backbone_pretrained=True,
+            backbone_modalities=["S2L2A"], 
+            backbone_merge_method="mean",
+            decoder="UperNetDecoder",
+            decoder_scale_modules=True,
+            decoder_channels=256,
+            num_classes=2,
+            head_dropout=0.3,
+            necks=[
+                {"name": "ReshapeTokensToImage", "remove_cls_token": False},
+                {"name": "SelectIndices", "indices": [2, 5, 8, 11]},
+            ],
+        )
+
+        self.task = ClassificationTask(
+            model_args=model_config,
+            model_factory="EncoderDecoderFactory",
+            loss="ce",
+            lr=self.args.lr,
+            ignore_index=-1,
+            optimizer="AdamW",
+            optimizer_hparams={"weight_decay": self.args.weight_decay},
+        )
+        self.task.configure_models()
+        self.task.configure_losses()
+        return self.task.model.to(self.device)
+
+    def _init_optimizer(self):
+        optimizer = optim.AdamW(
+            self.model.parameters(), 
+            lr=self.args.lr, 
+            weight_decay=self.args.weight_decay
+        )
+        scheduler = optim.lr_scheduler.ReduceLROnPlateau(
+            optimizer, mode='min', patience=5, verbose=True
+        )
+        return optimizer, scheduler
+
+    def run_epoch(self, dataloader: DataLoader, stage: str = "train") -> Dict[str, float]:
+        """Runs a single epoch for either training or validation."""
+        is_train = stage == "train"
+        self.model.train() if is_train else self.model.eval()
+        
+        tracker = MetricTracker()
+        
+        # Context manager: enable grad only if training
+        with torch.set_grad_enabled(is_train):
+            pbar = tqdm(dataloader, desc=f"  {stage.capitalize()}", leave=False)
+            
+            for batch in pbar:
+                inputs = batch['S2L2A'].to(self.device)
+                targets = batch['label'].to(self.device)
+
+                # Forward Pass
+                outputs = self.model(x={"S2L2A": inputs})
+                probabilities = torch.softmax(outputs.output, dim=1)
+                loss = self.criterion(probabilities, targets)
+
+                if is_train:
+                    self.optimizer.zero_grad()
+                    loss.backward()
+                    self.optimizer.step()
+
+                # Update metrics
+                tracker.update(loss.item(), targets, probabilities)
+                
+                # Update progress bar description with live loss
+                pbar.set_postfix(loss=f"{loss.item():.4f}")
+
+        return tracker.compute()
+
+    def save_checkpoint(self, filename: str):
+        path = self.save_dir / filename
+        torch.save(self.model.state_dict(), path)
+        logger.info(f"Saved model to {path}")
+
+    def log_to_csv(self, epoch: int, train_metrics: Dict, val_metrics: Dict):
+        """Appends metrics to the CSV log file."""
+        csv_path = self.save_dir / 'train_val_metrics.csv'
+        file_exists = csv_path.exists()
+        
+        # Define headers based on metric keys
+        headers = ['Epoch'] + [f'Train_{k}' for k in train_metrics.keys()] + [f'Val_{k}' for k in val_metrics.keys()]
+        
+        with open(csv_path, mode='a', newline='') as f:
+            writer = csv.writer(f)
+            if not file_exists:
+                writer.writerow(headers)
+            
+            row = [epoch] + list(train_metrics.values()) + list(val_metrics.values())
+            writer.writerow(row)
+
+    def fit(self, train_loader: DataLoader, val_loader: DataLoader):
+        """Main training entry point."""
+        logger.info(f"Starting training for {self.args.epochs} epochs...")
+        start_time = time.time()
+
+        for epoch in range(1, self.args.epochs + 1):
+            logger.info(f"Epoch {epoch}/{self.args.epochs}")
+            
+            # Run Training & Validation
+            train_metrics = self.run_epoch(train_loader, stage="train")
+            val_metrics = self.run_epoch(val_loader, stage="validate")
+            
+            # Scheduler Step
+            self.scheduler.step(val_metrics['Loss'])
+            
+            # Logging
+            self.log_to_csv(epoch, train_metrics, val_metrics)
+            logger.info(
+                f"Train Loss: {train_metrics['Loss']:.4f} | "
+                f"Val Loss: {val_metrics['Loss']:.4f} | "
+                f"Val F1: {val_metrics['F1']:.4f}"
+            )
+
+            # Save Best Model
+            if val_metrics['Loss'] < self.best_val_loss:
+                self.best_val_loss = val_metrics['Loss']
+                self.save_checkpoint("best_model.pth")
+                logger.info(f"--> New best model (Val Loss: {self.best_val_loss:.4f})")
+
+        # End of training
+        self.save_checkpoint("final_model.pth")
+        logger.info(f"Training finished in {time.time() - start_time:.2f}s")
+
+
+# --- Data Utilities ---
+
+def get_data_loaders(args) -> Tuple[DataLoader, DataLoader]:
+    """Prepares DataModule and returns Train/Val loaders."""
+    
+    # Read Excel and Filter Folds
+    try:
+        df = pd.read_csv(args.excel_file) if args.excel_file.endswith('.csv') else pd.read_excel(args.excel_file)
+    except Exception as e:
+        logger.error(f"Failed to load summary file: {e}")
+        raise
+
+    # Determine training pool (all folds except test_fold)
+    all_folds = range(1, args.num_folds + 1)
+    train_pool_folds = [f for f in all_folds if f != args.test_fold]
+    
+    # Filter filenames
+    df_filtered = df[df['Fold'].isin(train_pool_folds)]
+    if df_filtered.empty:
+        raise ValueError(f"No data found for folds {train_pool_folds}. Check 'Fold' column in Excel.")
+    
+    paths = df_filtered['Filename'].tolist()
+    
+    # 80/20 Split
+    train_paths, val_paths = train_test_split(paths, test_size=0.2, random_state=args.seed)
+    
+    logger.info(f"Data Split - Train: {len(train_paths)}, Val: {len(val_paths)} (Test Fold: {args.test_fold})")
+
+    # Initialize DataModule
+    datamodule = MethaneClassificationDataModule(
+        data_root=args.root_dir,
+        excel_file=args.excel_file,
+        batch_size=args.batch_size,
+        paths=train_paths,
+        train_transform=get_training_transforms(),
+        val_transform=None,
+    )
+    
+    # Create Loaders
+    datamodule.paths = train_paths
+    datamodule.setup(stage="fit")
+    train_loader = datamodule.train_dataloader()
+    
+    datamodule.paths = val_paths
+    datamodule.setup(stage="validate")
+    val_loader = datamodule.val_dataloader()
+    
+    return train_loader, val_loader
+
+
+# --- Main Execution ---
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Methane Classification Training with TerraTorch")
+    
+    # Paths
+    parser.add_argument('--root_dir', type=str, required=True, help='Root directory for satellite images')
+    parser.add_argument('--excel_file', type=str, required=True, help='Path to summary Excel/CSV file')
+    parser.add_argument('--save_dir', type=str, default='./checkpoints', help='Directory to save outputs')
+    
+    # Training Config
+    parser.add_argument('--epochs', type=int, default=100)
+    parser.add_argument('--batch_size', type=int, default=8)
+    parser.add_argument('--lr', type=float, default=1e-5)
+    parser.add_argument('--weight_decay', type=float, default=0.05)
+    parser.add_argument('--num_folds', type=int, default=5)
+    parser.add_argument('--test_fold', type=int, default=2, help='Fold ID to hold out for testing')
+    parser.add_argument('--seed', type=int, default=42)
+    
+    return parser.parse_args()
+
+if __name__ == "__main__":
+    args = parse_args()
+    set_seed(args.seed)
+    
+    # Prepare Data
+    train_loader, val_loader = get_data_loaders(args)
+    
+    # Initialize Trainer and Start
+    trainer = MethaneTrainer(args)
+    trainer.fit(train_loader, val_loader)

classification_with_text/calculate_embeddings.py

@@ -0,0 +1,51 @@
+import json
+import pandas as pd
+from sentence_transformers import SentenceTransformer
+from pathlib import Path
+from tqdm import tqdm
+
+def extract_caption(text_block):
+    for line in text_block.splitlines():
+        if "CAPTION:" in line.upper():
+            return line.split("CAPTION:")[-1].strip()
+    return ""
+
+def load_captions_from_files(json_files):
+    all_paths = []
+    all_captions = []
+
+    for json_path in tqdm(json_files, desc="Reading files"):
+        with open(json_path, 'r', encoding='utf-8') as f:
+            data = json.load(f)
+
+        for img_path, outer_list in data.items():
+            if not outer_list or not outer_list[0]:
+                continue
+            text_block = outer_list[0][0]
+            caption = extract_caption(text_block)
+            if caption:
+                all_paths.append(img_path)
+                all_captions.append(caption)
+
+    return all_paths, all_captions
+
+def compute_and_save_embeddings(json_files, output_csv):
+    model = SentenceTransformer('all-MiniLM-L6-v2')
+    image_paths, captions = load_captions_from_files(json_files)
+
+    if not captions:
+        print("No valid captions found across input files.")
+        return
+
+    embeddings = model.encode(captions, show_progress_bar=True)
+    df = pd.DataFrame(embeddings)
+    df.insert(0, "image_path", image_paths)
+    df.to_csv(output_csv, index=False)
+    print(f"Saved {len(df)} embeddings from {len(json_files)} files to {output_csv}")
+
+# Example usage
+if __name__ == "__main__":
+    import glob
+    # Collect multiple JSON files
+    files = glob.glob("./MBD_text/*.json")  # or provide manually
+    compute_and_save_embeddings(files, "combined_caption_embeddings.csv")

classification_with_text/script/methan_text_dataset.py

@@ -0,0 +1,81 @@
+import os
+import rasterio
+import torch
+from torchgeo.datasets import NonGeoDataset
+from torch.utils.data import DataLoader
+import torch.nn.functional as F
+import numpy as np
+import pandas as pd
+import json
+
+class MethaneTextDataset(NonGeoDataset):
+    def __init__(self, root_dir, paths, captions, transform=None):
+        super().__init__()
+        self.root_dir = root_dir
+        self.transform = transform
+        self.data_paths = []
+        self.captions_dict = captions
+
+        # Collect paths for labelbinary.tif and sCube.tif in selected folders
+        for folder_name in paths:
+            subdir_path = os.path.join(root_dir, folder_name)
+            if os.path.isdir(subdir_path):
+                filename_tokens = subdir_path.split("/")
+                label_path = os.path.join(subdir_path, 'labelbinary.tif')
+                scube_path = os.path.join(subdir_path, 'sCube.tif')
+                if os.path.exists(label_path) and os.path.exists(scube_path):
+                    self.data_paths.append((label_path, scube_path))
+                else:
+                    print(f"Warning: Missing files in {subdir_path}. Expected labelbinary.tif and sCube.tif.")
+
+
+    def __len__(self):
+        return len(self.data_paths)
+
+    def __getitem__(self, idx):
+        label_path, scube_path = self.data_paths[idx]
+        filename_tokens = label_path.split("/")
+        folder_name = filename_tokens[-2]
+
+        # Load the label image (single band)
+        with rasterio.open(label_path) as label_src:
+            label_image = label_src.read(1)  # Shape: [512, 512]
+
+        # Load the sCube image (multi-band), drop the first band
+        with rasterio.open(scube_path) as scube_src:
+            scube_image = scube_src.read()  # Shape: [13, 512, 512]
+            scube_image = scube_image[1:, :, :]  # Drop first band → Shape: [12, 512, 512]
+
+        # print(label_image.shape)
+        # print(scube_image.shape)
+        # Convert to PyTorch tensors
+        scube_tensor = torch.from_numpy(scube_image).float()  # Shape: [12, 512, 512]
+        label_tensor = torch.from_numpy(label_image).float()  # Shape: [512, 512]
+
+        # Resize to [12, 224, 224] and [224, 224] respectively
+        scube_tensor = F.interpolate(scube_tensor.unsqueeze(0), size=(224, 224), mode='bilinear', align_corners=False).squeeze(0)
+        label_tensor = F.interpolate(label_tensor.unsqueeze(0).unsqueeze(0), size=(224, 224), mode='nearest').squeeze(0)
+
+        label_tensor = label_tensor.clip(0, 1)  # Clip values to [0, 1]
+        scube_tensor = torch.nan_to_num(scube_tensor, nan=0.0)  # Replace NaNs with 0
+        # normalized_tensor = min_max_normalize(scube_tensor)
+        # Convert labels to binary
+        contains_methane = (label_tensor > 0).any().long()
+
+        # Convert to one-hot encoding
+        one_hot_label = F.one_hot(contains_methane, num_classes=2).float()
+        
+        # Apply transformations (if any)
+        if self.transform:
+            transformed = self.transform(image=np.array(scube_tensor.permute(1, 2, 0)))
+            scube_tensor = transformed['image'].transpose(2, 0, 1)   # Convert back to [C, H, W]
+
+        if folder_name in self.captions_dict:
+            caption = self.captions_dict[folder_name]
+        else:
+            # If the folder name is not in the captions_dict, set a default caption or None 
+            caption = 'No caption'
+        # caption = self.captions_dict[folder_name]
+            
+
+        return {'S2L2A': scube_tensor, 'label': one_hot_label, 'caption': caption}

classification_with_text/script/methane_text_datamodule.py

@@ -0,0 +1,72 @@
+import os
+import rasterio
+import torch
+from torchgeo.datasets import NonGeoDataset
+from torch.utils.data import DataLoader
+from torchgeo.datamodules import NonGeoDataModule
+
+from methan_text_dataset import MethaneTextDataset
+
+class MethaneTextDataModule(NonGeoDataModule):
+    """
+    A DataModule for handling MethaneClassificationDataset
+    """
+
+    def __init__(
+        self,
+        data_root: str,
+        paths: list,
+        captions: list,
+        batch_size: int = 8,
+        num_workers: int = 0,
+        train_transform: callable = None,
+        val_transform: callable = None,
+        test_transform: callable = None,
+        **kwargs
+    ):
+        super().__init__(MethaneTextDataset, batch_size, num_workers, **kwargs)
+
+        self.data_root = data_root
+        self.paths = paths
+        self.captions = captions
+        self.train_transform = train_transform
+        self.val_transform = val_transform
+        self.test_transform = test_transform
+
+    def setup(self, stage: str = None):
+        if stage in ("fit", "train"):
+            self.train_dataset = MethaneTextDataset(
+                root_dir=self.data_root,
+                paths=self.paths,
+                captions=self.captions,
+                transform=self.train_transform,
+            )
+        if stage in ("fit", "validate", "val"):
+            self.val_dataset = MethaneTextDataset(
+                root_dir=self.data_root,
+                paths=self.paths,
+                captions=self.captions,
+                transform=self.val_transform,
+            )
+        if stage in ("test", "predict"):
+            self.test_dataset = MethaneTextDataset(
+                root_dir=self.data_root,
+                paths=self.paths,
+                captions=self.captions,
+                transform=self.test_transform,
+            )
+
+    def train_dataloader(self):
+        return DataLoader(
+            self.train_dataset, batch_size=self.batch_size, shuffle=True, num_workers=self.num_workers, drop_last=True
+        )
+
+    def val_dataloader(self):
+        return DataLoader(
+            self.val_dataset, batch_size=self.batch_size, shuffle=False, num_workers=self.num_workers, drop_last=True
+        )
+
+    def test_dataloader(self):
+        return DataLoader(
+            self.test_dataset, batch_size=self.batch_size, shuffle=False, num_workers=self.num_workers, drop_last=True
+        )

classification_with_text/script/train_text.py

@@ -0,0 +1,448 @@
+import argparse
+import logging
+import csv
+import random
+import warnings
+import time
+import json
+from pathlib import Path
+from functools import partial
+from typing import Dict, List, Tuple, Any, Optional
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import albumentations as A
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+from sklearn.model_selection import train_test_split
+from sklearn.metrics import (
+    accuracy_score, recall_score, f1_score, matthews_corrcoef, confusion_matrix
+)
+from rasterio.errors import NotGeoreferencedWarning
+from sentence_transformers import SentenceTransformer
+
+# --- CRITICAL IMPORTS ---
+import terramind 
+from terratorch.tasks import ClassificationTask
+from terratorch.registry import TERRATORCH_BACKBONE_REGISTRY, TERRATORCH_DECODER_REGISTRY
+from terramind.models.terramind_register import build_terrammind_vit
+
+# Local Imports
+from methane_text_datamodule import MethaneTextDataModule
+
+# --- Configuration & Setup ---
+
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    datefmt='%Y-%m-%d %H:%M:%S'
+)
+logger = logging.getLogger(__name__)
+
+logging.getLogger("rasterio._env").setLevel(logging.ERROR)
+warnings.simplefilter("ignore", NotGeoreferencedWarning)
+warnings.filterwarnings("ignore", category=FutureWarning)
+
+# --- Global Constants ---
+PRETRAINED_BANDS = {
+    'untok_sen2l2a@224': [
+        "COASTAL_AEROSOL", "BLUE", "GREEN", "RED", "RED_EDGE_1", "RED_EDGE_2",
+        "RED_EDGE_3", "NIR_BROAD", "NIR_NARROW", "WATER_VAPOR", "SWIR_1", "SWIR_2",
+    ]
+}
+
+def set_seed(seed: int = 42):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+def get_training_transforms() -> A.Compose:
+    return A.Compose([
+        A.ElasticTransform(p=0.25),
+        A.RandomRotate90(p=0.5),
+        A.Flip(p=0.5),
+        A.ShiftScaleRotate(rotate_limit=90, shift_limit_x=0.05, shift_limit_y=0.05, p=0.5)
+    ])
+
+# --- Custom Model Components (From Notebook) ---
+
+# Initialize Sentence Transformer globally to avoid reloading
+try:
+    EMBB_MODEL = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
+    # Move to GPU if available for faster encoding during training if needed, 
+    # though usage in forward() implies dynamic encoding.
+    if torch.cuda.is_available():
+        EMBB_MODEL = EMBB_MODEL.to("cuda")
+except Exception as e:
+    logger.warning(f"Could not load SentenceTransformer: {e}")
+    EMBB_MODEL = None
+
+class TerraMindWithText(nn.Module):
+    def __init__(self, terramind_kwargs: dict):
+        super().__init__()
+        self.terramind = build_terrammind_vit(
+            variant='terramind_v1_base',
+            encoder_depth=12,
+            dim=768,
+            num_heads=12,
+            mlp_ratio=4,
+            qkv_bias=False,
+            proj_bias=False,
+            mlp_bias=False,
+            norm_layer=partial(nn.LayerNorm, eps=1e-6),
+            act_layer=nn.SiLU,
+            gated_mlp=True,
+            pretrained_bands=PRETRAINED_BANDS,
+            **terramind_kwargs
+        )
+        self.out_channels = [768] * 12
+        # self.project = nn.Linear(768 + 512, 768*192) # Referenced in notebook but seemingly unused in forward
+
+    def forward(self, x, captions):
+        vision_features = self.terramind(x)  # shape: (batch_size, 768)
+        
+        # Encode captions
+        # Note: embb_model.encode returns numpy or tensor. Ensure it is on correct device.
+        with torch.no_grad():
+            captions_embed = EMBB_MODEL.encode(captions, convert_to_tensor=True, show_progress_bar=False)
+        
+        # Ensure dimensionality matches what decoder expects (Squeeze if necessary, though encode usually returns [B, D])
+        if len(captions_embed.shape) == 3:
+            captions_embed = captions_embed.squeeze()
+            
+        return vision_features + [captions_embed]
+
+@TERRATORCH_BACKBONE_REGISTRY.register
+def terramind_v1_base_with_text(**kwargs):
+    return TerraMindWithText(terramind_kwargs=kwargs)
+
+@TERRATORCH_DECODER_REGISTRY.register
+class SimpleDecoder(nn.Module):
+    includes_head = True
+
+    def __init__(self, input_dim=768, num_classes=2, caption_dim=384):
+        super().__init__()
+        # Handle input_dim if passed as list (common in TerraTorch)
+        dim = input_dim[0] if isinstance(input_dim, (list, tuple)) else input_dim
+        
+        self.image_conv = nn.Sequential(
+            nn.Conv2d(dim, 512, kernel_size=3, padding=1),
+            nn.BatchNorm2d(512),
+            nn.ReLU(inplace=True),
+            nn.Dropout2d(0.3),
+            nn.Conv2d(512, 256, kernel_size=3, padding=1),
+            nn.BatchNorm2d(256),
+            nn.ReLU(inplace=True),
+            nn.Dropout2d(0.3)
+        )
+
+        self.caption_mlp = nn.Sequential(
+            nn.Linear(caption_dim, 512),
+            nn.ReLU(inplace=True),
+            nn.Dropout(0.3),
+            nn.Linear(512, 256),
+            nn.ReLU(inplace=True),
+            nn.Dropout(0.3)
+        )
+
+        self.cross_attention = nn.MultiheadAttention(
+            embed_dim=256, num_heads=8, dropout=0.1, batch_first=True
+        )
+
+        self.fusion_conv = nn.Sequential(
+            nn.Conv2d(512, 256, kernel_size=3, padding=1),
+            nn.BatchNorm2d(256),
+            nn.ReLU(inplace=True),
+            nn.Dropout2d(0.3),
+            nn.Conv2d(256, 128, kernel_size=3, padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU(inplace=True),
+            nn.Dropout2d(0.3)
+        )
+
+        self.conv_head = nn.Sequential(
+            nn.Conv2d(128, 64, kernel_size=3, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True),
+            nn.Dropout2d(0.3),
+            nn.Conv2d(64, 1, kernel_size=1)
+        )
+
+        self.out_channels = 1
+
+    def forward(self, features: list[torch.Tensor]) -> torch.Tensor:
+        # features list contains: [vision_feat_0, ..., vision_feat_11, caption_embed]
+        caption_embed = features[-1] # [B, 384]
+        image_features = features[:12]
+        
+        # Average vision tokens
+        x = torch.stack(image_features, dim=1).mean(dim=1)  # [B, 196, 768]
+            
+        B, N, C = x.shape
+        H = W = int(N ** 0.5) 
+        
+        x = x.permute(0, 2, 1).view(B, C, H, W)  # [B, 768, 14, 14]
+        img_features = self.image_conv(x)        # [B, 256, 14, 14]
+        
+        # Ensure caption embed has batch dim
+        if caption_embed.dim() == 1:
+            caption_embed = caption_embed.unsqueeze(0)
+            
+        caption_features = self.caption_mlp(caption_embed)  # [B, 256]
+        
+        # Expand caption to spatial dims
+        caption_spatial = caption_features.unsqueeze(-1).unsqueeze(-1)
+        caption_spatial = caption_spatial.expand(B, -1, H, W)  # [B, 256, 14, 14]
+        
+        # Fuse
+        fused_features = torch.cat([img_features, caption_spatial], dim=1)  # [B, 512, 14, 14]
+        fused = self.fusion_conv(fused_features)  # [B, 128, 14, 14]
+        
+        output = self.conv_head(fused)  # [B, 1, 14, 14]
+        return output
+
+# --- Helper Classes ---
+
+class MetricTracker:
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.all_targets = []
+        self.all_predictions = []
+        self.total_loss = 0.0
+        self.steps = 0
+
+    def update(self, loss: float, targets: torch.Tensor, probabilities: torch.Tensor):
+        self.total_loss += loss
+        self.steps += 1
+        self.all_targets.extend(torch.argmax(targets, dim=1).detach().cpu().numpy())
+        self.all_predictions.extend(torch.argmax(probabilities, dim=1).detach().cpu().numpy())
+
+    def compute(self) -> Dict[str, float]:
+        if not self.all_targets:
+            return {}
+        
+        tn, fp, fn, tp = confusion_matrix(self.all_targets, self.all_predictions, labels=[0, 1]).ravel()
+        
+        return {
+            "Loss": self.total_loss / max(self.steps, 1),
+            "Accuracy": accuracy_score(self.all_targets, self.all_predictions),
+            "Specificity": tn / (tn + fp) if (tn + fp) != 0 else 0.0,
+            "Sensitivity": recall_score(self.all_targets, self.all_predictions, average='binary', pos_label=1, zero_division=0),
+            "F1": f1_score(self.all_targets, self.all_predictions, average='binary', pos_label=1, zero_division=0),
+            "MCC": matthews_corrcoef(self.all_targets, self.all_predictions),
+        }
+
+class MethaneTextTrainer:
+    def __init__(self, args: argparse.Namespace):
+        self.args = args
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.save_dir = Path(args.save_dir) / f'fold{args.test_fold}'
+        self.save_dir.mkdir(parents=True, exist_ok=True)
+        
+        self.model = self._init_model()
+        self.optimizer, self.scheduler = self._init_optimizer()
+        self.criterion = self.task.criterion
+        self.best_val_loss = float('inf')
+        
+        logger.info(f"Trainer initialized on device: {self.device}")
+
+    def _init_model(self) -> nn.Module:
+        model_args = dict(
+            backbone="terramind_v1_base_with_text",
+            backbone_pretrained=True,
+            backbone_modalities=["S2L2A"], 
+            backbone_merge_method="mean",
+            num_classes=2,
+            head_dropout=0.3,
+            decoder="SimpleDecoder",
+        )
+
+        self.task = ClassificationTask(
+            model_args=model_args,
+            model_factory="EncoderDecoderFactory",
+            loss="ce",
+            lr=self.args.lr,
+            ignore_index=-1,
+            optimizer="AdamW",
+            optimizer_hparams={"weight_decay": self.args.weight_decay},
+        )
+        self.task.configure_models()
+        self.task.configure_losses()
+        return self.task.model.to(self.device)
+
+    def _init_optimizer(self):
+        optimizer = optim.AdamW(self.model.parameters(), lr=self.args.lr, weight_decay=self.args.weight_decay)
+        scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', patience=5, verbose=True)
+        return optimizer, scheduler
+
+    def run_epoch(self, dataloader: DataLoader, stage: str = "train") -> Dict[str, float]:
+        is_train = stage == "train"
+        self.model.train() if is_train else self.model.eval()
+        tracker = MetricTracker()
+        
+        with torch.set_grad_enabled(is_train):
+            pbar = tqdm(dataloader, desc=f"  {stage.capitalize()}", leave=False)
+            for batch in pbar:
+                # Prepare Inputs
+                inputs = batch['S2L2A'].to(self.device)
+                captions = batch['caption'] # List of strings
+                targets = batch['label'].to(self.device)
+
+                # Forward Pass (Note: passing captions explicitly)
+                # The Task wrapper might expect x dict, but our custom backbone forward handles 'captions'
+                outputs = self.model(x={"S2L2A": inputs}, captions=captions)
+                probabilities = torch.softmax(outputs.output, dim=1)
+                loss = self.criterion(probabilities, targets)
+
+                if is_train:
+                    self.optimizer.zero_grad()
+                    loss.backward()
+                    self.optimizer.step()
+
+                tracker.update(loss.item(), targets, probabilities)
+                pbar.set_postfix(loss=f"{loss.item():.4f}")
+
+        return tracker.compute()
+
+    def log_to_csv(self, epoch: int, train_metrics: Dict, val_metrics: Dict):
+        csv_path = self.save_dir / 'train_val_metrics.csv'
+        headers = ['Epoch'] + [f'Train_{k}' for k in train_metrics.keys()] + [f'Val_{k}' for k in val_metrics.keys()]
+        
+        with open(csv_path, mode='a', newline='') as f:
+            writer = csv.writer(f)
+            if not csv_path.exists():
+                writer.writerow(headers)
+            writer.writerow([epoch] + list(train_metrics.values()) + list(val_metrics.values()))
+
+    def fit(self, train_loader: DataLoader, val_loader: DataLoader):
+        logger.info(f"Starting training for {self.args.epochs} epochs...")
+        start_time = time.time()
+
+        for epoch in range(1, self.args.epochs + 1):
+            logger.info(f"Epoch {epoch}/{self.args.epochs}")
+            
+            train_metrics = self.run_epoch(train_loader, stage="train")
+            val_metrics = self.run_epoch(val_loader, stage="validate")
+            
+            self.scheduler.step(val_metrics['Loss'])
+            self.log_to_csv(epoch, train_metrics, val_metrics)
+            
+            logger.info(f"Train Loss: {train_metrics['Loss']:.4f} | Val Loss: {val_metrics['Loss']:.4f} | Val F1: {val_metrics['F1']:.4f}")
+
+            if val_metrics['Loss'] < self.best_val_loss:
+                self.best_val_loss = val_metrics['Loss']
+                torch.save(self.model.state_dict(), self.save_dir / "best_model.pth")
+                logger.info(f"--> New best model saved")
+
+        torch.save(self.model.state_dict(), self.save_dir / "final_model.pth")
+        logger.info(f"Training finished in {time.time() - start_time:.2f}s")
+
+# --- Data Utilities ---
+
+def read_captions(json_path: Path, captions_dict: Dict) -> Dict:
+    """Reads captions from JSON and populates dictionary."""
+    if not json_path.exists():
+        logger.warning(f"Caption file not found: {json_path}")
+        return captions_dict
+        
+    try:
+        with open(json_path, "r", encoding="utf-8") as file:
+            data = json.load(file)
+
+        for file_path_str, text_list in data.items():
+            if text_list and isinstance(text_list, list) and text_list[0]:
+                text_content = text_list[0][0]
+                caption_start = text_content.find("CAPTION:")
+                if caption_start != -1:
+                    caption = text_content[caption_start + len("CAPTION:"):].strip()
+                    # Extract folder name (assumes specific directory structure from notebook)
+                    # "path\\to\\folder\\image.ext" -> "folder"
+                    path_parts = file_path_str.replace("\\", "/").split("/") 
+                    if len(path_parts) >= 2:
+                        last_directory = path_parts[-2]
+                        captions_dict[last_directory] = caption
+    except Exception as e:
+        logger.error(f"Error reading captions {json_path}: {e}")
+        
+    return captions_dict
+
+def get_paths_for_fold(excel_file: str, folds: List[int]) -> List[str]:
+    df = pd.read_excel(excel_file)
+    df_filtered = df[df['Fold'].isin(folds)]
+    return df_filtered['Filename'].tolist()
+
+def get_data_loaders(args) -> Tuple[DataLoader, DataLoader]:
+    # 1. Load Captions
+    captions_dict = {}
+    captions_dict = read_captions(Path(args.methane_captions), captions_dict)
+    captions_dict = read_captions(Path(args.no_methane_captions), captions_dict)
+    logger.info(f"Loaded {len(captions_dict)} captions.")
+
+    # 2. Get File Paths
+    all_folds = range(1, args.num_folds + 1)
+    train_pool_folds = [f for f in all_folds if f != args.test_fold]
+    paths = get_paths_for_fold(args.excel_file, train_pool_folds)
+    
+    # 3. Split
+    train_paths, val_paths = train_test_split(paths, test_size=0.2, random_state=args.seed)
+    logger.info(f"Train: {len(train_paths)}, Val: {len(val_paths)}")
+
+    # 4. DataModule
+    datamodule = MethaneTextDataModule(
+        data_root=args.root_dir,
+        paths=paths, # Initial dummy
+        captions=captions_dict,
+        train_transform=get_training_transforms(),
+        batch_size=args.batch_size,
+    )
+    
+    # Train Loader
+    datamodule.paths = train_paths
+    datamodule.setup(stage="train")
+    train_loader = datamodule.train_dataloader()
+    
+    # Val Loader
+    datamodule.paths = val_paths
+    datamodule.setup(stage="validate")
+    val_loader = datamodule.val_dataloader()
+    
+    return train_loader, val_loader
+
+# --- Main Execution ---
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Methane Text-Multimodal Training")
+    
+    # Data Paths
+    parser.add_argument('--root_dir', type=str, required=True, help='Root directory for images')
+    parser.add_argument('--excel_file', type=str, required=True, help='Path to Summary Excel')
+    parser.add_argument('--methane_captions', type=str, required=True, help='Path to Methane JSON captions')
+    parser.add_argument('--no_methane_captions', type=str, required=True, help='Path to No-Methane JSON captions')
+    parser.add_argument('--save_dir', type=str, default='./checkpoints', help='Output directory')
+    
+    # Hyperparameters
+    parser.add_argument('--epochs', type=int, default=100)
+    parser.add_argument('--batch_size', type=int, default=4)
+    parser.add_argument('--lr', type=float, default=5e-5)
+    parser.add_argument('--weight_decay', type=float, default=0.05)
+    parser.add_argument('--num_folds', type=int, default=5)
+    parser.add_argument('--test_fold', type=int, default=2)
+    parser.add_argument('--seed', type=int, default=42)
+    
+    return parser.parse_args()
+
+if __name__ == "__main__":
+    args = parse_args()
+    set_seed(args.seed)
+    
+    train_loader, val_loader = get_data_loaders(args)
+    
+    trainer = MethaneTextTrainer(args)
+    trainer.fit(train_loader, val_loader)

intuition1_classification_finetuning/config/methane_simulated_datamodule.py

@@ -0,0 +1,116 @@
+import pandas as pd
+import albumentations as A
+from sklearn.model_selection import train_test_split
+from torch.utils.data import DataLoader
+from torchgeo.datamodules import NonGeoDataModule
+from methane_simulated_dataset import MethaneSimulatedDataset
+
+class MethaneSimulatedDataModule(NonGeoDataModule):
+    def __init__(
+        self,
+        data_root: str,
+        excel_file: str,
+        batch_size: int = 8,
+        num_workers: int = 0,
+        val_split: float = 0.2,
+        seed: int = 42,
+        test_fold: int = 4,   # Default test fold from your script
+        num_folds: int = 5,
+        **kwargs
+    ):
+        super().__init__(MethaneSimulatedDataset, batch_size, num_workers, **kwargs)
+
+        self.data_root = data_root
+        self.excel_file = excel_file
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.val_split = val_split
+        self.seed = seed
+        self.test_fold = test_fold
+        self.num_folds = num_folds
+        
+        self.train_paths = []
+        self.val_paths = []
+
+    def _get_training_transforms(self):
+        return A.Compose([
+            A.ElasticTransform(p=0.25),
+            A.RandomRotate90(p=0.5),
+            A.Flip(p=0.5),
+            A.ShiftScaleRotate(rotate_limit=90, shift_limit_x=0.05, shift_limit_y=0.05, p=0.5)
+        ])
+
+    def _get_simulated_paths(self, paths):
+        """Logic to rename files to I1/TOA format"""
+        simulated_paths = []
+        for path in paths:
+            try:
+                tokens = path.split('_')
+                if len(tokens) >= 5:
+                    simulated_path = f"{tokens[0]}_toarefl_{tokens[3]}_{tokens[4]}"
+                    simulated_paths.append(simulated_path)
+                else:
+                    simulated_paths.append(path)
+            except Exception:
+                simulated_paths.append(path)
+        return simulated_paths
+
+    def setup(self, stage: str = None):
+        # 1. Read Excel
+        try:
+            df = pd.read_excel(self.excel_file)
+        except Exception as e:
+            raise RuntimeError(f"Failed to load excel: {e}")
+
+        # 2. Filter Folds (Exclude test_fold)
+        all_folds = list(range(1, self.num_folds + 1))
+        train_pool_folds = [f for f in all_folds if f != self.test_fold]
+        
+        df_filtered = df[df['Fold'].isin(train_pool_folds)]
+        raw_paths = df_filtered['Filename'].tolist()
+
+        # 3. Apply Path Renaming Logic
+        paths = self._get_simulated_paths(raw_paths)
+
+        # 4. Train/Val Split
+        self.train_paths, self.val_paths = train_test_split(
+            paths, 
+            test_size=self.val_split, 
+            random_state=self.seed
+        )
+
+        # 5. Instantiate Datasets
+        if stage in ("fit", "train"):
+            self.train_dataset = MethaneSimulatedDataset(
+                root_dir=self.data_root,
+                excel_file=self.excel_file,
+                paths=self.train_paths,
+                transform=self._get_training_transforms(),
+            )
+        
+        if stage in ("fit", "validate", "val"):
+            self.val_dataset = MethaneSimulatedDataset(
+                root_dir=self.data_root,
+                excel_file=self.excel_file,
+                paths=self.val_paths,
+                transform=None,
+            )
+
+    def train_dataloader(self):
+        return DataLoader(self.train_dataset, batch_size=self.batch_size, shuffle=True, 
+                          num_workers=self.num_workers, drop_last=True)
+
+    def val_dataloader(self):
+        return DataLoader(self.val_dataset, batch_size=self.batch_size, shuffle=False, 
+                          num_workers=self.num_workers, drop_last=True)
+
+    # def on_after_batch_transfer(self, batch, dataloader_idx):
+    #     # 1. Run TorchGeo default (expects 'image')
+    #     batch = super().on_after_batch_transfer(batch, dataloader_idx)
+        
+    #     # 2. Wrap into TerraMind format {'S2L2A': ...}
+    #     if 'image' in batch:
+    #         s2_data = batch['image']
+    #         batch['image'] = {'S2L2A': s2_data}
+            
+    #     return batch

intuition1_classification_finetuning/config/methane_simulated_dataset.py

@@ -0,0 +1,66 @@
+import os
+import rasterio
+import torch
+from torchgeo.datasets import NonGeoDataset
+import torch.nn.functional as F
+import numpy as np
+
+class MethaneSimulatedDataset(NonGeoDataset):
+    def __init__(self, root_dir, excel_file, paths, transform=None):
+        super().__init__()
+        self.root_dir = root_dir
+        self.transform = transform
+        self.data_paths = []
+
+        # Collect paths
+        for folder_name in paths:
+            subdir_path = os.path.join(root_dir, folder_name)
+            if os.path.isdir(subdir_path):
+                # Note: Filenames here seem to match the folder name based on your script
+                label_path = os.path.join(subdir_path, folder_name + '_mask.tif')
+                scube_path = os.path.join(subdir_path, folder_name + '_hsi.dat')
+
+                if os.path.exists(label_path) and os.path.exists(scube_path):
+                    self.data_paths.append((label_path, scube_path))
+
+    def __len__(self):
+        return len(self.data_paths)
+
+    def __getitem__(self, idx):
+        label_path, scube_path = self.data_paths[idx]
+
+        # Load label
+        with rasterio.open(label_path) as label_src:
+            label_image = label_src.read()
+
+        # Load sCube (I1/TOA data)
+        with rasterio.open(scube_path) as scube_src:
+            scube_image = scube_src.read()
+            # Read only first 12 bands
+            scube_image = scube_image[:12, :, :]  
+
+        # Convert to Tensors
+        scube_tensor = torch.from_numpy(scube_image).float()
+        label_tensor = torch.from_numpy(label_image).float()
+
+        # Resize
+        scube_tensor = F.interpolate(scube_tensor.unsqueeze(0), size=(224, 224), mode='bilinear', align_corners=False).squeeze(0)
+        label_tensor = F.interpolate(label_tensor.unsqueeze(0), size=(224, 224), mode='nearest').squeeze(0)
+
+        label_tensor = label_tensor.clip(0, 1)
+        scube_tensor = torch.nan_to_num(scube_tensor, nan=0.0)
+
+        # Convert labels to binary index (0 or 1)
+        contains_methane = (label_tensor > 0).any().long()
+
+        # Apply transformations
+        if self.transform:
+            # Albumentations expects [H, W, C]
+            transformed = self.transform(image=np.array(scube_tensor.permute(1, 2, 0)))
+            scube_tensor = transformed['image'].transpose(2, 0, 1)
+
+        return {
+            'image': scube_tensor,       # <--- Named 'image' for TorchGeo
+            'label': contains_methane,   # <--- Index for CrossEntropy
+            'sample': scube_path
+        }

intuition1_classification_finetuning/config/train.yaml

@@ -0,0 +1,66 @@
+seed_everything: 42
+
+trainer:
+  accelerator: auto
+  strategy: auto
+  devices: 1
+  max_epochs: 100
+  default_root_dir: ./checkpoints_i1
+  callbacks:
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        monitor: val/loss
+        mode: min
+        save_top_k: 1
+        filename: "best_model"
+        save_last: true
+    - class_path: lightning.pytorch.callbacks.LearningRateMonitor
+      init_args:
+        logging_interval: epoch
+
+model:
+  class_path: terratorch.tasks.ClassificationTask
+  init_args:
+    model_factory: EncoderDecoderFactory
+    loss: ce
+    ignore_index: -1
+    lr: 1.0e-5  # Top-level LR
+
+    optimizer: AdamW
+    optimizer_hparams:
+      weight_decay: 0.05
+      
+    scheduler: ReduceLROnPlateau
+    scheduler_hparams:
+      mode: min
+      patience: 5
+
+    model_args:
+      backbone: terramind_v1_base
+      backbone_pretrained: true
+      backbone_modalities:
+        - S2L2A
+      backbone_merge_method: mean
+      
+      decoder: UperNetDecoder
+      decoder_scale_modules: true
+      decoder_channels: 256
+      num_classes: 2
+      head_dropout: 0.3
+      
+      necks:
+        - name: ReshapeTokensToImage
+          remove_cls_token: false
+        - name: SelectIndices
+          indices: [2, 5, 8, 11]
+
+data:
+  class_path: methane_simulated_datamodule.MethaneSimulatedDataModule
+  init_args:
+    data_root: # Place the data root here
+    excel_file: ../../Methane_benchmark_patches_summary_v3.xlsx   # Update this!
+    batch_size: 8
+    val_split: 0.2
+    seed: 42
+    test_fold: 4
+    num_folds: 5

intuition1_classification_finetuning/script/methane_simulated_datamodule.py

@@ -0,0 +1,72 @@
+import os
+import rasterio
+import torch
+from torchgeo.datasets import NonGeoDataset
+from torch.utils.data import DataLoader
+from torchgeo.datamodules import NonGeoDataModule
+from methane_simulated_dataset import MethaneSimulatedDataset
+# from methane_classification_dataset import MethaneClassificationDataset
+
+class MethaneSimulatedDataModule(NonGeoDataModule):
+    """
+    A DataModule for handling MethaneClassificationDataset
+    """
+
+    def __init__(
+        self,
+        data_root: str,
+        excel_file: str,
+        paths: list,
+        batch_size: int = 8,
+        num_workers: int = 0,
+        train_transform: callable = None,
+        val_transform: callable = None,
+        test_transform: callable = None,
+        **kwargs
+    ):
+        super().__init__(MethaneSimulatedDataset, batch_size, num_workers, **kwargs)
+
+        self.data_root = data_root
+        self.excel_file = excel_file
+        self.paths = paths
+        self.train_transform = train_transform
+        self.val_transform = val_transform
+        self.test_transform = test_transform
+
+    def setup(self, stage: str = None):
+        if stage in ("fit", "train"):
+            self.train_dataset = MethaneSimulatedDataset(
+                root_dir=self.data_root,
+                excel_file=self.excel_file,
+                paths=self.paths,
+                transform=self.train_transform,
+            )
+        if stage in ("fit", "validate", "val"):
+            self.val_dataset = MethaneSimulatedDataset(
+                root_dir=self.data_root,
+                excel_file=self.excel_file,
+                paths=self.paths,
+                transform=self.val_transform,
+            )
+        if stage in ("test", "predict"):
+            self.test_dataset = MethaneSimulatedDataset(
+                root_dir=self.data_root,
+                excel_file=self.excel_file,
+                paths=self.paths,
+                transform=self.test_transform,
+            )
+
+    def train_dataloader(self):
+        return DataLoader(
+            self.train_dataset, batch_size=self.batch_size, shuffle=True, num_workers=self.num_workers, drop_last=True
+        )
+
+    def val_dataloader(self):
+        return DataLoader(
+            self.val_dataset, batch_size=self.batch_size, shuffle=False, num_workers=self.num_workers, drop_last=True
+        )
+
+    def test_dataloader(self):
+        return DataLoader(
+            self.test_dataset, batch_size=self.batch_size, shuffle=False, num_workers=self.num_workers, drop_last=True
+        )

intuition1_classification_finetuning/script/methane_simulated_dataset.py

@@ -0,0 +1,81 @@
+import os
+import rasterio
+import torch
+from torchgeo.datasets import NonGeoDataset
+from torch.utils.data import DataLoader
+import torch.nn.functional as F
+import numpy as np
+import pandas as pd
+
+def min_max_normalize(data, new_min=0, new_max=1):
+    data = np.array(data, dtype=np.float32)  # Convert to NumPy array
+    
+    # Handle NaN, Inf values
+    data = np.nan_to_num(data, nan=np.nanmin(data), posinf=np.max(data), neginf=np.min(data))
+
+    old_min, old_max = np.min(data), np.max(data)
+    
+    if old_max == old_min:  # Prevent division by zero
+        return np.full_like(data, new_min, dtype=np.float32)  # Uniform array
+
+    return (data - old_min) / (old_max - old_min + 1e-10) * (new_max - new_min) + new_min
+
+class MethaneSimulatedDataset(NonGeoDataset):
+    def __init__(self, root_dir, excel_file, paths, transform=None):
+        super().__init__()
+        self.root_dir = root_dir
+        self.transform = transform
+        self.data_paths = []
+
+        # Collect paths for labelbinary.tif and sCube.tif in selected folders
+        for folder_name in paths:
+            subdir_path = os.path.join(root_dir, folder_name)
+            if os.path.isdir(subdir_path):
+                label_path = os.path.join(subdir_path, folder_name + '_mask.tif')
+                scube_path = os.path.join(subdir_path, folder_name + '_hsi.dat')
+
+                if os.path.exists(label_path) and os.path.exists(scube_path):
+                    self.data_paths.append((label_path, scube_path))
+
+    def __len__(self):
+        return len(self.data_paths)
+
+    def __getitem__(self, idx):
+        label_path, scube_path = self.data_paths[idx]
+
+        # Load the label image (single band)
+        with rasterio.open(label_path) as label_src:
+            label_image = label_src.read()  # Shape: [512, 512]
+
+
+        # Load the sCube image (multi-band), drop the first band
+        with rasterio.open(scube_path) as scube_src:
+            scube_image = scube_src.read()  # Shape: [13, 512, 512]
+
+            # Read only the first 12 bands for testing purposes
+            # Map the bands later on 
+            scube_image = scube_image[:12, :, :]  
+
+        # Convert to PyTorch tensors
+        scube_tensor = torch.from_numpy(scube_image).float()  # Shape: [12, 512, 512]
+        label_tensor = torch.from_numpy(label_image).float()  # Shape: [512, 512]
+
+        # Resize to [12, 224, 224] and [224, 224] respectively
+        scube_tensor = F.interpolate(scube_tensor.unsqueeze(0), size=(224, 224), mode='bilinear', align_corners=False).squeeze(0)
+        label_tensor = F.interpolate(label_tensor.unsqueeze(0), size=(224, 224), mode='nearest').squeeze(0)
+
+        label_tensor = label_tensor.clip(0, 1)  # Clip values to [0, 1]
+        scube_tensor = torch.nan_to_num(scube_tensor, nan=0.0)  # Replace NaNs with 0
+        # normalized_tensor = min_max_normalize(scube_tensor)
+        # Convert labels to binary
+        contains_methane = (label_tensor > 0).any().long()
+
+        # Convert to one-hot encoding
+        one_hot_label = F.one_hot(contains_methane, num_classes=2).float()
+        
+        # Apply transformations (if any)
+        if self.transform:
+            transformed = self.transform(image=np.array(scube_tensor.permute(1, 2, 0)))
+            scube_tensor = transformed['image'].transpose(2, 0, 1)   # Convert back to [C, H, W]
+
+        return {'S2L2A': scube_tensor, 'label': one_hot_label, 'sample': scube_path}

intuition1_classification_finetuning/script/train_simulated_I1.py

@@ -0,0 +1,309 @@
+import argparse
+import logging
+import csv
+import random
+import warnings
+import time
+from pathlib import Path
+from typing import Dict, List, Tuple, Any, Optional
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import albumentations as A
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+from sklearn.model_selection import train_test_split
+from sklearn.metrics import (
+    accuracy_score, recall_score, f1_score, matthews_corrcoef, confusion_matrix
+)
+from rasterio.errors import NotGeoreferencedWarning
+
+# --- CRITICAL IMPORTS ---
+import terramind
+from terratorch.tasks import ClassificationTask
+
+# Local Imports
+from methane_simulated_datamodule import MethaneSimulatedDataModule
+
+# --- Configuration & Setup ---
+
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    datefmt='%Y-%m-%d %H:%M:%S'
+)
+logger = logging.getLogger(__name__)
+
+logging.getLogger("rasterio._env").setLevel(logging.ERROR)
+warnings.simplefilter("ignore", NotGeoreferencedWarning)
+warnings.filterwarnings("ignore", category=FutureWarning)
+
+def set_seed(seed: int = 42):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+def get_training_transforms() -> A.Compose:
+    return A.Compose([
+        A.ElasticTransform(p=0.25),
+        A.RandomRotate90(p=0.5),
+        A.Flip(p=0.5),
+        A.ShiftScaleRotate(rotate_limit=90, shift_limit_x=0.05, shift_limit_y=0.05, p=0.5)
+    ])
+
+# --- Path Utilities ---
+
+def get_simulated_paths(paths: List[str]) -> List[str]:
+    """
+    Modifies filenames to match the I1/TOA naming convention.
+    Converts 'ang2015..._S2_...' -> 'ang2015..._toarefl_...'
+    """
+    simulated_paths = []
+    for path in paths:
+        try:
+            tokens = path.split('_')
+            # Logic: {ID}_toarefl_{Coord1}_{Coord2}
+            # Adjusts original filename tokens to target format
+            if len(tokens) >= 5:
+                simulated_path = f"{tokens[0]}_toarefl_{tokens[3]}_{tokens[4]}"
+                simulated_paths.append(simulated_path)
+            else:
+                simulated_paths.append(path)
+        except Exception as e:
+            logger.warning(f"Could not parse path {path}: {e}")
+            simulated_paths.append(path)
+    return simulated_paths
+
+def get_paths_for_fold(excel_file: str, folds: List[int]) -> List[str]:
+    try:
+        df = pd.read_excel(excel_file)
+        df_filtered = df[df['Fold'].isin(folds)]
+        return df_filtered['Filename'].tolist()
+    except Exception as e:
+        logger.error(f"Error reading Excel file: {e}")
+        raise
+
+# --- Helper Classes ---
+
+class MetricTracker:
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.all_targets = []
+        self.all_predictions = []
+        self.total_loss = 0.0
+        self.steps = 0
+
+    def update(self, loss: float, targets: torch.Tensor, probabilities: torch.Tensor):
+        self.total_loss += loss
+        self.steps += 1
+        self.all_targets.extend(torch.argmax(targets, dim=1).detach().cpu().numpy())
+        self.all_predictions.extend(torch.argmax(probabilities, dim=1).detach().cpu().numpy())
+
+    def compute(self) -> Dict[str, float]:
+        if not self.all_targets:
+            return {}
+        
+        tn, fp, fn, tp = confusion_matrix(self.all_targets, self.all_predictions, labels=[0, 1]).ravel()
+        
+        return {
+            "Loss": self.total_loss / max(self.steps, 1),
+            "Accuracy": accuracy_score(self.all_targets, self.all_predictions),
+            "Specificity": tn / (tn + fp) if (tn + fp) != 0 else 0.0,
+            "Sensitivity": recall_score(self.all_targets, self.all_predictions, average='binary', pos_label=1, zero_division=0),
+            "F1": f1_score(self.all_targets, self.all_predictions, average='binary', pos_label=1, zero_division=0),
+            "MCC": matthews_corrcoef(self.all_targets, self.all_predictions),
+            "TP": int(tp), "TN": int(tn), "FP": int(fp), "FN": int(fn)
+        }
+
+class TrainerI1:
+    def __init__(self, args: argparse.Namespace):
+        self.args = args
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.save_dir = Path(args.save_dir) / f'fold{args.test_fold}'
+        self.save_dir.mkdir(parents=True, exist_ok=True)
+        
+        self.model = self._init_model()
+        self.optimizer, self.scheduler = self._init_optimizer()
+        self.criterion = self.task.criterion
+        self.best_val_loss = float('inf')
+        
+        logger.info(f"Trainer initialized on device: {self.device}")
+
+    def _init_model(self) -> nn.Module:
+        model_args = dict(
+            backbone="terramind_v1_base",
+            backbone_pretrained=True,
+            backbone_modalities=["S2L2A"], 
+            backbone_merge_method="mean",
+            decoder="UperNetDecoder",
+            decoder_scale_modules=True,
+            decoder_channels=256,
+            num_classes=2,
+            head_dropout=0.3,
+            necks=[
+                {"name": "ReshapeTokensToImage", "remove_cls_token": False},
+                {"name": "SelectIndices", "indices": [2, 5, 8, 11]},
+                {"name": "LearnedInterpolateToPyramidal"},
+            ],
+        )
+
+        self.task = ClassificationTask(
+            model_args=model_args,
+            model_factory="EncoderDecoderFactory",
+            loss="ce",
+            lr=self.args.lr,
+            ignore_index=-1,
+            optimizer="AdamW",
+            optimizer_hparams={"weight_decay": self.args.weight_decay},
+        )
+        self.task.configure_models()
+        self.task.configure_losses()
+        return self.task.model.to(self.device)
+
+    def _init_optimizer(self):
+        optimizer = optim.AdamW(self.model.parameters(), lr=self.args.lr, weight_decay=self.args.weight_decay)
+        scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', patience=5, verbose=True)
+        return optimizer, scheduler
+
+    def run_epoch(self, dataloader: DataLoader, stage: str = "train") -> Dict[str, float]:
+        is_train = stage == "train"
+        self.model.train() if is_train else self.model.eval()
+        tracker = MetricTracker()
+        
+        with torch.set_grad_enabled(is_train):
+            pbar = tqdm(dataloader, desc=f"  {stage.capitalize()}", leave=False)
+            for batch in pbar:
+                inputs = batch['S2L2A'].to(self.device)
+                targets = batch['label'].to(self.device)
+
+                outputs = self.model(x={"S2L2A": inputs})
+                probabilities = torch.softmax(outputs.output, dim=1)
+                loss = self.criterion(probabilities, targets)
+
+                if is_train:
+                    self.optimizer.zero_grad()
+                    loss.backward()
+                    self.optimizer.step()
+
+                tracker.update(loss.item(), targets, probabilities)
+                pbar.set_postfix(loss=f"{loss.item():.4f}")
+
+        return tracker.compute()
+
+    def fit(self, train_loader: DataLoader, val_loader: DataLoader):
+        logger.info(f"Starting training for {self.args.epochs} epochs...")
+        start_time = time.time()
+        
+        # Initialize CSV logging
+        csv_path = self.save_dir / 'train_val_metrics.csv'
+        with open(csv_path, 'w', newline='') as f:
+            writer = csv.writer(f)
+            writer.writerow([
+                'Epoch', 'Train_Loss', 'Train_F1', 'Train_Acc', 
+                'Val_Loss', 'Val_F1', 'Val_Acc', 'Val_Spec', 'Val_Sens'
+            ])
+
+        for epoch in range(1, self.args.epochs + 1):
+            logger.info(f"Epoch {epoch}/{self.args.epochs}")
+            
+            train_metrics = self.run_epoch(train_loader, stage="train")
+            val_metrics = self.run_epoch(val_loader, stage="validate")
+            
+            self.scheduler.step(val_metrics['Loss'])
+            
+            # Log to CSV
+            with open(csv_path, 'a', newline='') as f:
+                writer = csv.writer(f)
+                writer.writerow([
+                    epoch, 
+                    train_metrics.get('Loss'), train_metrics.get('F1'), train_metrics.get('Accuracy'),
+                    val_metrics.get('Loss'), val_metrics.get('F1'), val_metrics.get('Accuracy'),
+                    val_metrics.get('Specificity'), val_metrics.get('Sensitivity')
+                ])
+
+            logger.info(f"Train Loss: {train_metrics['Loss']:.4f} | Val Loss: {val_metrics['Loss']:.4f} | Val F1: {val_metrics['F1']:.4f}")
+
+            # Save Best Model
+            if val_metrics['Loss'] < self.best_val_loss:
+                self.best_val_loss = val_metrics['Loss']
+                torch.save(self.model.state_dict(), self.save_dir / "best_model.pth")
+                logger.info(f"--> New best model saved")
+
+        # Save Final Model
+        torch.save(self.model.state_dict(), self.save_dir / "final_model.pth")
+        logger.info(f"Training finished in {time.time() - start_time:.2f}s")
+
+# --- Data Utilities ---
+
+def get_data_loaders(args) -> Tuple[DataLoader, DataLoader]:
+    # 1. Determine Folds
+    all_folds = list(range(1, args.num_folds + 1))
+    train_pool_folds = [f for f in all_folds if f != args.test_fold]
+    
+    # 2. Get Paths & Convert to TOA/I1 format
+    # Note: Using get_simulated_paths to transform names as done in the notebook
+    paths = get_paths_for_fold(args.excel_file, train_pool_folds)
+    paths = get_simulated_paths(paths)
+    
+    # 3. Train/Val Split (80/20)
+    train_paths, val_paths = train_test_split(paths, test_size=0.2, random_state=args.seed)
+    
+    logger.info(f"Data Split - Train: {len(train_paths)}, Val: {len(val_paths)} (Test Fold: {args.test_fold})")
+
+    # 4. Initialize DataModule
+    datamodule = MethaneSimulatedDataModule(
+        data_root=args.root_dir,
+        excel_file=args.excel_file,
+        batch_size=args.batch_size,
+        paths=paths, # Initial dummy
+        train_transform=get_training_transforms(),
+        val_transform=None,
+    )
+    
+    # 5. Create Loaders
+    datamodule.paths = train_paths
+    datamodule.setup(stage="fit")
+    train_loader = datamodule.train_dataloader()
+    
+    datamodule.paths = val_paths
+    datamodule.setup(stage="validate")
+    val_loader = datamodule.val_dataloader()
+    
+    return train_loader, val_loader
+
+# --- Main Execution ---
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Methane I1 (TOA Refl) Training")
+    
+    # Paths
+    parser.add_argument('--root_dir', type=str, required=True, help='Root directory for I1/TOA data')
+    parser.add_argument('--excel_file', type=str, required=True, help='Path to Summary Excel')
+    parser.add_argument('--save_dir', type=str, default='./checkpoints_i1', help='Output directory')
+    
+    # Hyperparameters
+    parser.add_argument('--epochs', type=int, default=100)
+    parser.add_argument('--batch_size', type=int, default=1)
+    parser.add_argument('--lr', type=float, default=1e-5)
+    parser.add_argument('--weight_decay', type=float, default=0.05)
+    parser.add_argument('--num_folds', type=int, default=5)
+    parser.add_argument('--test_fold', type=int, default=4, help='Fold ID to hold out')
+    parser.add_argument('--seed', type=int, default=42)
+    
+    return parser.parse_args()
+
+if __name__ == "__main__":
+    args = parse_args()
+    set_seed(args.seed)
+    
+    train_loader, val_loader = get_data_loaders(args)
+    
+    trainer = TrainerI1(args)
+    trainer.fit(train_loader, val_loader)

sentinel2_classification_finetuning/config/methane_simulated_datamodule.py

@@ -0,0 +1,119 @@
+import pandas as pd
+import albumentations as A
+from sklearn.model_selection import train_test_split
+from torch.utils.data import DataLoader
+from torchgeo.datamodules import NonGeoDataModule
+from methane_simulated_dataset import MethaneSimulatedDataset
+
+class MethaneSimulatedDataModule(NonGeoDataModule):
+    def __init__(
+        self,
+        data_root: str,
+        excel_file: str,
+        batch_size: int = 8,
+        num_workers: int = 0,
+        val_split: float = 0.2,
+        seed: int = 42,
+        test_fold: int = 4,
+        num_folds: int = 5,
+        sim_tag: str = "toarefl",  # <--- New arg for 'toarefl'/'boarefl'
+        **kwargs
+    ):
+        super().__init__(MethaneSimulatedDataset, batch_size, num_workers, **kwargs)
+
+        self.data_root = data_root
+        self.excel_file = excel_file
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.val_split = val_split
+        self.seed = seed
+        self.test_fold = test_fold
+        self.num_folds = num_folds
+        self.sim_tag = sim_tag
+        
+        self.train_paths = []
+        self.val_paths = []
+
+    def _get_training_transforms(self):
+        return A.Compose([
+            A.ElasticTransform(p=0.25),
+            A.RandomRotate90(p=0.5),
+            A.Flip(p=0.5),
+            A.ShiftScaleRotate(rotate_limit=90, shift_limit_x=0.05, shift_limit_y=0.05, p=0.5)
+        ])
+
+    def _get_simulated_paths(self, paths):
+        """Internal logic to rename files based on sim_tag"""
+        simulated_paths = []
+        for path in paths:
+            try:
+                tokens = path.split('_')
+                if len(tokens) >= 5:
+                    # Logic: {ID}_{tag}_{Coord1}_{Coord2}
+                    simulated_path = f"{tokens[0]}_{self.sim_tag}_{tokens[3]}_{tokens[4]}"
+                    simulated_paths.append(simulated_path)
+                else:
+                    simulated_paths.append(path)
+            except Exception:
+                simulated_paths.append(path)
+        return simulated_paths
+
+    def setup(self, stage: str = None):
+        # 1. Read Excel
+        try:
+            df = pd.read_excel(self.excel_file)
+        except Exception as e:
+            raise RuntimeError(f"Failed to load excel: {e}")
+
+        # 2. Filter Folds (Exclude test_fold)
+        all_folds = list(range(1, self.num_folds + 1))
+        train_pool_folds = [f for f in all_folds if f != self.test_fold]
+        
+        df_filtered = df[df['Fold'].isin(train_pool_folds)]
+        raw_paths = df_filtered['Filename'].tolist()
+
+        # 3. Apply Path Renaming Logic
+        paths = self._get_simulated_paths(raw_paths)
+
+        # 4. Train/Val Split
+        self.train_paths, self.val_paths = train_test_split(
+            paths, 
+            test_size=self.val_split, 
+            random_state=self.seed
+        )
+
+        # 5. Instantiate Datasets
+        if stage in ("fit", "train"):
+            self.train_dataset = MethaneSimulatedDataset(
+                root_dir=self.data_root,
+                excel_file=self.excel_file,
+                paths=self.train_paths,
+                transform=self._get_training_transforms(),
+            )
+        
+        if stage in ("fit", "validate", "val"):
+            self.val_dataset = MethaneSimulatedDataset(
+                root_dir=self.data_root,
+                excel_file=self.excel_file,
+                paths=self.val_paths,
+                transform=None,
+            )
+
+    def train_dataloader(self):
+        return DataLoader(self.train_dataset, batch_size=self.batch_size, shuffle=True, 
+                          num_workers=self.num_workers, drop_last=True)
+
+    def val_dataloader(self):
+        return DataLoader(self.val_dataset, batch_size=self.batch_size, shuffle=False, 
+                          num_workers=self.num_workers, drop_last=True)
+
+    def on_after_batch_transfer(self, batch, dataloader_idx):
+        # 1. Run TorchGeo default (expects 'image')
+        batch = super().on_after_batch_transfer(batch, dataloader_idx)
+        
+        # 2. Wrap into TerraMind format {'S2L2A': ...}
+        if 'image' in batch:
+            s2_data = batch['image']
+            batch['image'] = {'S2L2A': s2_data}
+            
+        return batch

sentinel2_classification_finetuning/config/methane_simulated_dataset.py

@@ -0,0 +1,70 @@
+import os
+import rasterio
+import torch
+from torchgeo.datasets import NonGeoDataset
+import torch.nn.functional as F
+import numpy as np
+
+class MethaneSimulatedDataset(NonGeoDataset):
+    def __init__(self, root_dir, excel_file, paths, transform=None):
+        super().__init__()
+        self.root_dir = root_dir
+        self.transform = transform
+        self.data_paths = []
+
+        # Collect paths
+        for folder_name in paths:
+            subdir_path = os.path.join(root_dir, folder_name)
+            if os.path.isdir(subdir_path):
+                # Construct paths based on folder name
+                label_path = os.path.join(subdir_path, folder_name + '_mask.tif')
+                scube_path = os.path.join(subdir_path, folder_name + '_hsi.dat')
+
+                if os.path.exists(label_path) and os.path.exists(scube_path):
+                    self.data_paths.append((label_path, scube_path))
+
+    def __len__(self):
+        return len(self.data_paths)
+
+    def __getitem__(self, idx):
+        label_path, scube_path = self.data_paths[idx]
+
+        # Load label
+        with rasterio.open(label_path) as label_src:
+            label_image = label_src.read()
+
+        # Load sCube (Try explicit ENVI driver first for .dat files)
+        try:
+            with rasterio.open(scube_path, driver='ENVI') as scube_src:
+                scube_image = scube_src.read()
+                scube_image = scube_image[:12, :, :]  # Read first 12 bands
+        except Exception:
+            # Fallback if driver auto-detection is needed
+            with rasterio.open(scube_path) as scube_src:
+                scube_image = scube_src.read()
+                scube_image = scube_image[:12, :, :]
+
+        # Convert to Tensors
+        scube_tensor = torch.from_numpy(scube_image).float()
+        label_tensor = torch.from_numpy(label_image).float()
+
+        # Resize
+        scube_tensor = F.interpolate(scube_tensor.unsqueeze(0), size=(224, 224), mode='bilinear', align_corners=False).squeeze(0)
+        label_tensor = F.interpolate(label_tensor.unsqueeze(0), size=(224, 224), mode='nearest').squeeze(0)
+
+        label_tensor = label_tensor.clip(0, 1)
+        scube_tensor = torch.nan_to_num(scube_tensor, nan=0.0)
+
+        # Convert labels to binary index (0 or 1)
+        contains_methane = (label_tensor > 0).any().long()
+
+        # Apply transformations
+        if self.transform:
+            transformed = self.transform(image=np.array(scube_tensor.permute(1, 2, 0)))
+            scube_tensor = transformed['image'].transpose(2, 0, 1)
+
+        return {
+            'image': scube_tensor,       # <--- 'image' for TorchGeo
+            'label': contains_methane,   # <--- Index for CE Loss
+            'sample': scube_path
+        }

sentinel2_classification_finetuning/config/train.yaml

@@ -0,0 +1,67 @@
+seed_everything: 42
+
+trainer:
+  accelerator: auto
+  strategy: auto
+  devices: 1
+  max_epochs: 100
+  default_root_dir: ./checkpoints_s2_simulated
+  callbacks:
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        monitor: val/loss
+        mode: min
+        save_top_k: 1
+        filename: "best_model"
+        save_last: true
+    - class_path: lightning.pytorch.callbacks.LearningRateMonitor
+      init_args:
+        logging_interval: epoch
+
+model:
+  class_path: terratorch.tasks.ClassificationTask
+  init_args:
+    model_factory: EncoderDecoderFactory
+    loss: ce
+    ignore_index: -1
+    lr: 1.0e-5  # Top-level LR
+
+    optimizer: AdamW
+    optimizer_hparams:
+      weight_decay: 0.05
+      
+    scheduler: ReduceLROnPlateau
+    scheduler_hparams:
+      mode: min
+      patience: 5
+
+    model_args:
+      backbone: terramind_v1_base
+      backbone_pretrained: true
+      backbone_modalities:
+        - S2L2A
+      backbone_merge_method: mean
+      
+      decoder: UperNetDecoder
+      decoder_scale_modules: true
+      decoder_channels: 256
+      num_classes: 2
+      head_dropout: 0.3
+      
+      necks:
+        - name: ReshapeTokensToImage
+          remove_cls_token: false
+        - name: SelectIndices
+          indices: [2, 5, 8, 11]
+
+data:
+  class_path: methane_simulated_datamodule.MethaneSimulatedDataModule
+  init_args:
+    data_root: /path/to/data_root  # UPDATE THIS
+    excel_file: ../../Methane_benchmark_patches_summary_v3.xlsx   
+    batch_size: 8
+    val_split: 0.2
+    seed: 42
+    test_fold: 4
+    num_folds: 5
+    sim_tag: toarefl  # Change to 'boarefl' if needed

sentinel2_classification_finetuning/script/inference_s2_simulated.py

@@ -0,0 +1,241 @@
+import argparse
+import logging
+import csv
+import random
+import warnings
+import time
+import os
+from pathlib import Path
+from typing import Dict, List, Tuple, Any, Optional
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+import albumentations as A
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+from rasterio.errors import NotGeoreferencedWarning
+
+# --- CRITICAL IMPORTS ---
+import terramind
+from terratorch.tasks import ClassificationTask
+
+# Local Imports
+from methane_simulated_datamodule import MethaneSimulatedDataModule
+
+# --- Configuration & Setup ---
+
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    datefmt='%Y-%m-%d %H:%M:%S'
+)
+logger = logging.getLogger(__name__)
+
+logging.getLogger("rasterio._env").setLevel(logging.ERROR)
+warnings.simplefilter("ignore", NotGeoreferencedWarning)
+warnings.filterwarnings("ignore", category=FutureWarning)
+
+def set_seed(seed: int = 42):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+def get_inference_transforms() -> A.Compose:
+    return None
+
+# --- Path Utilities (Crucial for Simulated Data) ---
+
+def get_simulated_paths(paths: List[str]) -> List[str]:
+    """
+    Modifies filenames to match the simulated dataset naming convention.
+    Original: 'MBD_0001_S2_...' -> Simulated: 'MBD_toarefl_S2_...'
+    """
+    simulated_paths = []
+    for path in paths:
+        try:
+            tokens = path.split('_')
+            # Reconstruct filename based on notebook logic
+            if len(tokens) >= 5:
+                # e.g., MBD_toarefl_S2_123_456
+                simulated_path = f"{tokens[0]}_toarefl_{tokens[3]}_{tokens[4]}"
+                simulated_paths.append(simulated_path)
+            else:
+                simulated_paths.append(path) 
+        except Exception as e:
+            logger.warning(f"Could not parse path {path}: {e}")
+            simulated_paths.append(path)
+    return simulated_paths
+
+# --- Inference Class ---
+
+class SimulatedInference:
+    def __init__(self, args: argparse.Namespace):
+        self.args = args
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.output_dir = Path(args.output_dir)
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        
+        logger.info(f"Initializing Inference on device: {self.device}")
+        
+        self.model = self._init_model()
+        self._load_checkpoint(args.checkpoint)
+
+    def _init_model(self) -> nn.Module:
+        model_args = dict(
+            backbone="terramind_v1_base",
+            backbone_pretrained=False,
+            backbone_modalities=["S2L2A"], 
+            backbone_merge_method="mean",
+            decoder="UperNetDecoder",
+            decoder_scale_modules=True,
+            decoder_channels=256,
+            num_classes=2,
+            head_dropout=0.3,
+            necks=[
+                {"name": "ReshapeTokensToImage", "remove_cls_token": False},
+                {"name": "SelectIndices", "indices": [2, 5, 8, 11]},
+                {"name": "LearnedInterpolateToPyramidal"},
+            ],
+        )
+
+        task = ClassificationTask(
+            model_args=model_args,
+            model_factory="EncoderDecoderFactory",
+            loss="ce",
+            ignore_index=-1
+        )
+        task.configure_models()
+        return task.model.to(self.device)
+
+    def _load_checkpoint(self, checkpoint_path: str):
+        path = Path(checkpoint_path)
+        if not path.exists():
+            raise FileNotFoundError(f"Checkpoint not found at {path}")
+            
+        logger.info(f"Loading weights from {path}...")
+        checkpoint = torch.load(path, map_location=self.device)
+        
+        if 'state_dict' in checkpoint:
+            state_dict = checkpoint['state_dict']
+        else:
+            state_dict = checkpoint
+            
+        self.model.load_state_dict(state_dict, strict=False)
+        self.model.eval()
+
+    def run_inference(self, dataloader: DataLoader, sample_names: List[str]):
+        """
+        Generates predictions and matches them with provided sample identifiers.
+        """
+        results = {}
+        
+        logger.info(f"Starting inference on {len(sample_names)} samples...")
+        
+        # Iterator to match predictions with original filenames
+        name_iter = iter(sample_names)
+        
+        with torch.no_grad():
+            for batch in tqdm(dataloader, desc="Predicting"):
+                inputs = batch['S2L2A'].to(self.device)
+                
+                # Forward Pass
+                outputs = self.model(x={"S2L2A": inputs})
+                probabilities = torch.softmax(outputs.output, dim=1)
+                
+                # Get binary prediction (0 or 1)
+                predictions = torch.argmax(probabilities, dim=1)
+                batch_preds = predictions.cpu().numpy()
+                
+                # Assign to Sample Names
+                for pred in batch_preds:
+                    try:
+                        sample_name = next(name_iter)
+                        results[sample_name] = int(pred)
+                    except StopIteration:
+                        logger.error("More predictions than sample names! Check sync.")
+                        break
+
+        if len(results) != len(sample_names):
+            logger.warning(f"Mismatch: Expected {len(sample_names)} results, got {len(results)}.")
+
+        # Save CSV
+        self._save_results(results)
+
+    def _save_results(self, results: Dict[str, int]):
+        csv_path = self.output_dir / "simulated_predictions.csv"
+        with open(csv_path, mode='w', newline='') as f:
+            writer = csv.writer(f)
+            writer.writerow(['Sample_ID', 'Prediction'])
+            for sample, pred in results.items():
+                writer.writerow([sample, pred])
+        logger.info(f"Predictions saved to {csv_path}")
+
+# --- Data Loading ---
+
+def get_dataloader_and_names(args) -> Tuple[DataLoader, List[str]]:
+    # 1. Read Excel to get base filenames
+    try:
+        df = pd.read_excel(args.excel_file)
+        # If specific folds are requested, filter them
+        if args.folds:
+            folds_to_use = [int(f) for f in args.folds.split(',')]
+            df = df[df['Fold'].isin(folds_to_use)]
+            logger.info(f"Filtered to folds: {folds_to_use}")
+        
+        raw_paths = df['Filename'].tolist()
+        logger.info(f"Loaded {len(raw_paths)} paths from Excel.")
+    except Exception as e:
+        logger.error(f"Error reading Excel: {e}")
+        raise
+
+    # 2. Transform paths to Simulated format
+    simulated_paths = get_simulated_paths(raw_paths)
+    
+    # 3. Initialize DataModule
+    datamodule = MethaneSimulatedDataModule(
+        data_root=args.root_dir,
+        excel_file=args.excel_file,
+        batch_size=args.batch_size,
+        paths=simulated_paths,
+        train_transform=None,
+        val_transform=get_inference_transforms(),
+    )
+    
+    # 4. Setup
+    datamodule.paths = simulated_paths
+    datamodule.setup(stage="test") 
+    
+    # Try getting test_dataloader, else train/val
+    loader = datamodule.test_dataloader() if hasattr(datamodule, 'test_dataloader') else datamodule.train_dataloader()
+    
+    return loader, simulated_paths
+
+# --- Main Execution ---
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Methane Simulated S2 Inference")
+    
+    parser.add_argument('--root_dir', type=str, required=True, help='Root directory for simulated data')
+    parser.add_argument('--excel_file', type=str, required=True, help='Path to Summary Excel file')
+    parser.add_argument('--checkpoint', type=str, required=True, help='Path to model checkpoint (.pth)')
+    parser.add_argument('--output_dir', type=str, default='./inference_results', help='Directory to save results')
+    parser.add_argument('--folds', type=str, default=None, help='Comma-separated list of folds to infer on (e.g., "4" or "1,2"). If None, uses all.')
+    parser.add_argument('--batch_size', type=int, default=1, help='Inference batch size')
+    parser.add_argument('--seed', type=int, default=42)
+    
+    return parser.parse_args()
+
+if __name__ == "__main__":
+    args = parse_args()
+    set_seed(args.seed)
+    
+    # 1. Prepare Data & Names
+    dataloader, sample_names = get_dataloader_and_names(args)
+    
+    # 2. Run Inference
+    engine = SimulatedInference(args)
+    engine.run_inference(dataloader, sample_names)

sentinel2_classification_finetuning/script/methane_simulated_datamodule.py

@@ -0,0 +1,72 @@
+import os
+import rasterio
+import torch
+from torchgeo.datasets import NonGeoDataset
+from torch.utils.data import DataLoader
+from torchgeo.datamodules import NonGeoDataModule
+from methane_simulated_dataset import MethaneSimulatedDataset
+# from methane_classification_dataset import MethaneClassificationDataset
+
+class MethaneSimulatedDataModule(NonGeoDataModule):
+    """
+    A DataModule for handling MethaneClassificationDataset
+    """
+
+    def __init__(
+        self,
+        data_root: str,
+        excel_file: str,
+        paths: list,
+        batch_size: int = 8,
+        num_workers: int = 0,
+        train_transform: callable = None,
+        val_transform: callable = None,
+        test_transform: callable = None,
+        **kwargs
+    ):
+        super().__init__(MethaneSimulatedDataset, batch_size, num_workers, **kwargs)
+
+        self.data_root = data_root
+        self.excel_file = excel_file
+        self.paths = paths
+        self.train_transform = train_transform
+        self.val_transform = val_transform
+        self.test_transform = test_transform
+
+    def setup(self, stage: str = None):
+        if stage in ("fit", "train"):
+            self.train_dataset = MethaneSimulatedDataset(
+                root_dir=self.data_root,
+                excel_file=self.excel_file,
+                paths=self.paths,
+                transform=self.train_transform,
+            )
+        if stage in ("fit", "validate", "val"):
+            self.val_dataset = MethaneSimulatedDataset(
+                root_dir=self.data_root,
+                excel_file=self.excel_file,
+                paths=self.paths,
+                transform=self.val_transform,
+            )
+        if stage in ("test", "predict"):
+            self.test_dataset = MethaneSimulatedDataset(
+                root_dir=self.data_root,
+                excel_file=self.excel_file,
+                paths=self.paths,
+                transform=self.test_transform,
+            )
+
+    def train_dataloader(self):
+        return DataLoader(
+            self.train_dataset, batch_size=self.batch_size, shuffle=True, num_workers=self.num_workers, drop_last=True
+        )
+
+    def val_dataloader(self):
+        return DataLoader(
+            self.val_dataset, batch_size=self.batch_size, shuffle=False, num_workers=self.num_workers, drop_last=True
+        )
+
+    def test_dataloader(self):
+        return DataLoader(
+            self.test_dataset, batch_size=self.batch_size, shuffle=False, num_workers=self.num_workers, drop_last=True
+        )

sentinel2_classification_finetuning/script/methane_simulated_dataset.py

@@ -0,0 +1,81 @@
+import os
+import rasterio
+import torch
+from torchgeo.datasets import NonGeoDataset
+from torch.utils.data import DataLoader
+import torch.nn.functional as F
+import numpy as np
+import pandas as pd
+
+def min_max_normalize(data, new_min=0, new_max=1):
+    data = np.array(data, dtype=np.float32)  # Convert to NumPy array
+    
+    # Handle NaN, Inf values
+    data = np.nan_to_num(data, nan=np.nanmin(data), posinf=np.max(data), neginf=np.min(data))
+
+    old_min, old_max = np.min(data), np.max(data)
+    
+    if old_max == old_min:  # Prevent division by zero
+        return np.full_like(data, new_min, dtype=np.float32)  # Uniform array
+
+    return (data - old_min) / (old_max - old_min + 1e-10) * (new_max - new_min) + new_min
+
+class MethaneSimulatedDataset(NonGeoDataset):
+    def __init__(self, root_dir, excel_file, paths, transform=None):
+        super().__init__()
+        self.root_dir = root_dir
+        self.transform = transform
+        self.data_paths = []
+
+        # Collect paths for labelbinary.tif and sCube.tif in selected folders
+        for folder_name in paths:
+            subdir_path = os.path.join(root_dir, folder_name)
+            if os.path.isdir(subdir_path):
+                label_path = os.path.join(subdir_path, folder_name + '_mask.tif')
+                scube_path = os.path.join(subdir_path, folder_name + '_hsi.dat')
+
+                if os.path.exists(label_path) and os.path.exists(scube_path):
+                    self.data_paths.append((label_path, scube_path))
+
+    def __len__(self):
+        return len(self.data_paths)
+
+    def __getitem__(self, idx):
+        label_path, scube_path = self.data_paths[idx]
+
+        # Load the label image (single band)
+        with rasterio.open(label_path) as label_src:
+            label_image = label_src.read()  # Shape: [512, 512]
+
+
+        # Load the sCube image (multi-band), drop the first band
+        with rasterio.open(scube_path) as scube_src:
+            scube_image = scube_src.read()  # Shape: [13, 512, 512]
+
+            # Read only the first 12 bands for testing purposes
+            # Map the bands later on 
+            scube_image = scube_image[:12, :, :]  
+
+        # Convert to PyTorch tensors
+        scube_tensor = torch.from_numpy(scube_image).float()  # Shape: [12, 512, 512]
+        label_tensor = torch.from_numpy(label_image).float()  # Shape: [512, 512]
+
+        # Resize to [12, 224, 224] and [224, 224] respectively
+        scube_tensor = F.interpolate(scube_tensor.unsqueeze(0), size=(224, 224), mode='bilinear', align_corners=False).squeeze(0)
+        label_tensor = F.interpolate(label_tensor.unsqueeze(0), size=(224, 224), mode='nearest').squeeze(0)
+
+        label_tensor = label_tensor.clip(0, 1)  # Clip values to [0, 1]
+        scube_tensor = torch.nan_to_num(scube_tensor, nan=0.0)  # Replace NaNs with 0
+        # normalized_tensor = min_max_normalize(scube_tensor)
+        # Convert labels to binary
+        contains_methane = (label_tensor > 0).any().long()
+
+        # Convert to one-hot encoding
+        one_hot_label = F.one_hot(contains_methane, num_classes=2).float()
+        
+        # Apply transformations (if any)
+        if self.transform:
+            transformed = self.transform(image=np.array(scube_tensor.permute(1, 2, 0)))
+            scube_tensor = transformed['image'].transpose(2, 0, 1)   # Convert back to [C, H, W]
+
+        return {'S2L2A': scube_tensor, 'label': one_hot_label, 'sample': scube_path}

sentinel2_classification_finetuning/script/train_simulated_s2.py

@@ -0,0 +1,314 @@
+import argparse
+import logging
+import csv
+import random
+import warnings
+import time
+from pathlib import Path
+from typing import Dict, List, Tuple, Any, Optional
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import albumentations as A
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+from sklearn.model_selection import train_test_split
+from sklearn.metrics import (
+    accuracy_score, recall_score, f1_score, matthews_corrcoef, confusion_matrix
+)
+from rasterio.errors import NotGeoreferencedWarning
+
+# --- CRITICAL IMPORTS ---
+import terramind
+from terratorch.tasks import ClassificationTask
+
+# Local Imports
+from methane_simulated_datamodule import MethaneSimulatedDataModule
+
+# --- Configuration & Setup ---
+
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    datefmt='%Y-%m-%d %H:%M:%S'
+)
+logger = logging.getLogger(__name__)
+
+logging.getLogger("rasterio._env").setLevel(logging.ERROR)
+warnings.simplefilter("ignore", NotGeoreferencedWarning)
+warnings.filterwarnings("ignore", category=FutureWarning)
+
+def set_seed(seed: int = 42):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+def get_training_transforms() -> A.Compose:
+    return A.Compose([
+        A.ElasticTransform(p=0.25),
+        A.RandomRotate90(p=0.5),
+        A.Flip(p=0.5),
+        A.ShiftScaleRotate(rotate_limit=90, shift_limit_x=0.05, shift_limit_y=0.05, p=0.5)
+    ])
+
+# --- Path Utilities ---
+
+def get_simulated_paths(paths: List[str], tag: str = "toarefl") -> List[str]:
+    """
+    Modifies filenames to match the I1/TOA naming convention.
+    Converts 'ang2015..._S2_...' -> 'ang2015..._{tag}_...'
+    """
+    simulated_paths = []
+    for path in paths:
+        try:
+            tokens = path.split('_')
+            # Logic: {ID}_{tag}_{Coord1}_{Coord2}
+            # Adjusts original filename tokens to target format
+            if len(tokens) >= 5:
+                simulated_path = f"{tokens[0]}_{tag}_{tokens[3]}_{tokens[4]}"
+                simulated_paths.append(simulated_path)
+            else:
+                simulated_paths.append(path)
+        except Exception as e:
+            logger.warning(f"Could not parse path {path}: {e}")
+            simulated_paths.append(path)
+    return simulated_paths
+
+def get_paths_for_fold(excel_file: str, folds: List[int]) -> List[str]:
+    try:
+        df = pd.read_excel(excel_file)
+        df_filtered = df[df['Fold'].isin(folds)]
+        return df_filtered['Filename'].tolist()
+    except Exception as e:
+        logger.error(f"Error reading Excel file: {e}")
+        raise
+
+# --- Helper Classes ---
+
+class MetricTracker:
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.all_targets = []
+        self.all_predictions = []
+        self.total_loss = 0.0
+        self.steps = 0
+
+    def update(self, loss: float, targets: torch.Tensor, probabilities: torch.Tensor):
+        self.total_loss += loss
+        self.steps += 1
+        self.all_targets.extend(torch.argmax(targets, dim=1).detach().cpu().numpy())
+        self.all_predictions.extend(torch.argmax(probabilities, dim=1).detach().cpu().numpy())
+
+    def compute(self) -> Dict[str, float]:
+        if not self.all_targets:
+            return {}
+        
+        tn, fp, fn, tp = confusion_matrix(self.all_targets, self.all_predictions, labels=[0, 1]).ravel()
+        
+        return {
+            "Loss": self.total_loss / max(self.steps, 1),
+            "Accuracy": accuracy_score(self.all_targets, self.all_predictions),
+            "Specificity": tn / (tn + fp) if (tn + fp) != 0 else 0.0,
+            "Sensitivity": recall_score(self.all_targets, self.all_predictions, average='binary', pos_label=1, zero_division=0),
+            "F1": f1_score(self.all_targets, self.all_predictions, average='binary', pos_label=1, zero_division=0),
+            "MCC": matthews_corrcoef(self.all_targets, self.all_predictions),
+            "TP": int(tp), "TN": int(tn), "FP": int(fp), "FN": int(fn)
+        }
+
+class TrainerI1:
+    def __init__(self, args: argparse.Namespace):
+        self.args = args
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.save_dir = Path(args.save_dir) / f'fold{args.test_fold}'
+        self.save_dir.mkdir(parents=True, exist_ok=True)
+        
+        self.model = self._init_model()
+        self.optimizer, self.scheduler = self._init_optimizer()
+        self.criterion = self.task.criterion
+        self.best_val_loss = float('inf')
+        
+        logger.info(f"Trainer initialized on device: {self.device}")
+
+    def _init_model(self) -> nn.Module:
+        model_args = dict(
+            backbone="terramind_v1_base",
+            backbone_pretrained=True,
+            backbone_modalities=["S2L2A"], 
+            backbone_merge_method="mean",
+            decoder="UperNetDecoder",
+            decoder_scale_modules=True,
+            decoder_channels=256,
+            num_classes=2,
+            head_dropout=0.3,
+            necks=[
+                {"name": "ReshapeTokensToImage", "remove_cls_token": False},
+                {"name": "SelectIndices", "indices": [2, 5, 8, 11]},
+                {"name": "LearnedInterpolateToPyramidal"},
+            ],
+        )
+
+        self.task = ClassificationTask(
+            model_args=model_args,
+            model_factory="EncoderDecoderFactory",
+            loss="ce",
+            lr=self.args.lr,
+            ignore_index=-1,
+            optimizer="AdamW",
+            optimizer_hparams={"weight_decay": self.args.weight_decay},
+        )
+        self.task.configure_models()
+        self.task.configure_losses()
+        return self.task.model.to(self.device)
+
+    def _init_optimizer(self):
+        optimizer = optim.AdamW(self.model.parameters(), lr=self.args.lr, weight_decay=self.args.weight_decay)
+        scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', patience=5, verbose=True)
+        return optimizer, scheduler
+
+    def run_epoch(self, dataloader: DataLoader, stage: str = "train") -> Dict[str, float]:
+        is_train = stage == "train"
+        self.model.train() if is_train else self.model.eval()
+        tracker = MetricTracker()
+        
+        with torch.set_grad_enabled(is_train):
+            pbar = tqdm(dataloader, desc=f"  {stage.capitalize()}", leave=False)
+            for batch in pbar:
+                inputs = batch['S2L2A'].to(self.device)
+                targets = batch['label'].to(self.device)
+
+                outputs = self.model(x={"S2L2A": inputs})
+                probabilities = torch.softmax(outputs.output, dim=1)
+                loss = self.criterion(probabilities, targets)
+
+                if is_train:
+                    self.optimizer.zero_grad()
+                    loss.backward()
+                    self.optimizer.step()
+
+                tracker.update(loss.item(), targets, probabilities)
+                pbar.set_postfix(loss=f"{loss.item():.4f}")
+
+        return tracker.compute()
+
+    def fit(self, train_loader: DataLoader, val_loader: DataLoader):
+        logger.info(f"Starting training for {self.args.epochs} epochs...")
+        start_time = time.time()
+        
+        # Initialize CSV logging
+        csv_path = self.save_dir / 'train_val_metrics.csv'
+        with open(csv_path, 'w', newline='') as f:
+            writer = csv.writer(f)
+            writer.writerow([
+                'Epoch', 'Train_Loss', 'Train_F1', 'Train_Acc', 
+                'Val_Loss', 'Val_F1', 'Val_Acc', 'Val_Spec', 'Val_Sens'
+            ])
+
+        for epoch in range(1, self.args.epochs + 1):
+            logger.info(f"Epoch {epoch}/{self.args.epochs}")
+            
+            train_metrics = self.run_epoch(train_loader, stage="train")
+            val_metrics = self.run_epoch(val_loader, stage="validate")
+            
+            self.scheduler.step(val_metrics['Loss'])
+            
+            # Log to CSV
+            with open(csv_path, 'a', newline='') as f:
+                writer = csv.writer(f)
+                writer.writerow([
+                    epoch, 
+                    train_metrics.get('Loss'), train_metrics.get('F1'), train_metrics.get('Accuracy'),
+                    val_metrics.get('Loss'), val_metrics.get('F1'), val_metrics.get('Accuracy'),
+                    val_metrics.get('Specificity'), val_metrics.get('Sensitivity')
+                ])
+
+            logger.info(f"Train Loss: {train_metrics['Loss']:.4f} | Val Loss: {val_metrics['Loss']:.4f} | Val F1: {val_metrics['F1']:.4f}")
+
+            # Save Best Model
+            if val_metrics['Loss'] < self.best_val_loss:
+                self.best_val_loss = val_metrics['Loss']
+                torch.save(self.model.state_dict(), self.save_dir / "best_model.pth")
+                logger.info(f"--> New best model saved")
+
+        # Save Final Model
+        torch.save(self.model.state_dict(), self.save_dir / "final_model.pth")
+        logger.info(f"Training finished in {time.time() - start_time:.2f}s")
+
+# --- Data Utilities ---
+
+def get_data_loaders(args) -> Tuple[DataLoader, DataLoader]:
+    # 1. Determine Folds
+    all_folds = list(range(1, args.num_folds + 1))
+    train_pool_folds = [f for f in all_folds if f != args.test_fold]
+    
+    # 2. Get Paths
+    paths = get_paths_for_fold(args.excel_file, train_pool_folds)
+    
+    # 3. Apply Tag (Dynamic Tagging)
+    paths = get_simulated_paths(paths, tag=args.sim_tag)
+    
+    # 4. Train/Val Split (80/20)
+    train_paths, val_paths = train_test_split(paths, test_size=0.2, random_state=args.seed)
+    
+    logger.info(f"Data Split - Train: {len(train_paths)}, Val: {len(val_paths)} (Test Fold: {args.test_fold})")
+
+    # 5. Initialize DataModule
+    datamodule = MethaneSimulatedDataModule(
+        data_root=args.root_dir,
+        excel_file=args.excel_file,
+        batch_size=args.batch_size,
+        paths=paths, # Initial dummy
+        train_transform=get_training_transforms(),
+        val_transform=None,
+    )
+    
+    # 6. Create Loaders
+    datamodule.paths = train_paths
+    datamodule.setup(stage="fit")
+    train_loader = datamodule.train_dataloader()
+    
+    datamodule.paths = val_paths
+    datamodule.setup(stage="validate")
+    val_loader = datamodule.val_dataloader()
+    
+    return train_loader, val_loader
+
+# --- Main Execution ---
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Methane I1 (Simulated) Training")
+    
+    # Paths
+    parser.add_argument('--root_dir', type=str, required=True, help='Root directory for I1/TOA/BOA data')
+    parser.add_argument('--excel_file', type=str, required=True, help='Path to Summary Excel')
+    parser.add_argument('--save_dir', type=str, default='./checkpoints_i1', help='Output directory')
+    
+    # Simulation Tag Config
+    parser.add_argument('--sim_tag', type=str, default='toarefl', 
+                        help='String identifier in filename (e.g. "toarefl" or "boarefl")')
+
+    # Hyperparameters
+    parser.add_argument('--epochs', type=int, default=100)
+    parser.add_argument('--batch_size', type=int, default=2, help='Batch size (must be >1 for BatchNorm)')
+    parser.add_argument('--lr', type=float, default=1e-5)
+    parser.add_argument('--weight_decay', type=float, default=0.05)
+    parser.add_argument('--num_folds', type=int, default=5)
+    parser.add_argument('--test_fold', type=int, default=4, help='Fold ID to hold out')
+    parser.add_argument('--seed', type=int, default=42)
+    
+    return parser.parse_args()
+
+if __name__ == "__main__":
+    args = parse_args()
+    set_seed(args.seed)
+    
+    train_loader, val_loader = get_data_loaders(args)
+    
+    trainer = TrainerI1(args)
+    trainer.fit(train_loader, val_loader)

urban_inference/methane_urban_datamodule.py

@@ -0,0 +1,71 @@
+import os
+import rasterio
+import torch
+from torchgeo.datasets import NonGeoDataset
+from torch.utils.data import DataLoader
+from torchgeo.datamodules import NonGeoDataModule
+from methane_urban_dataset import MethaneUrbanDataset
+
+class MethaneUrbanDataModule(NonGeoDataModule):
+    """
+    A DataModule for handling MethaneClassificationDataset
+    """
+
+    def __init__(
+        self,
+        data_root: str,
+        excel_file: str,
+        paths: list,
+        batch_size: int = 8,
+        num_workers: int = 0,
+        train_transform: callable = None,
+        val_transform: callable = None,
+        test_transform: callable = None,
+        **kwargs
+    ):
+        super().__init__(MethaneUrbanDataset, batch_size, num_workers, **kwargs)
+
+        self.data_root = data_root
+        self.excel_file = excel_file
+        self.paths = paths
+        self.train_transform = train_transform
+        self.val_transform = val_transform
+        self.test_transform = test_transform
+
+    def setup(self, stage: str = None):
+        if stage in ("fit", "train"):
+            self.train_dataset = MethaneUrbanDataset(
+                root_dir=self.data_root,
+                excel_file=self.excel_file,
+                paths=self.paths,
+                transform=self.train_transform,
+            )
+        if stage in ("fit", "validate", "val"):
+            self.val_dataset = MethaneUrbanDataset(
+                root_dir=self.data_root,
+                excel_file=self.excel_file,
+                paths=self.paths,
+                transform=self.val_transform,
+            )
+        if stage in ("test", "predict"):
+            self.test_dataset = MethaneUrbanDataset(
+                root_dir=self.data_root,
+                excel_file=self.excel_file,
+                paths=self.paths,
+                transform=self.test_transform,
+            )
+
+    def train_dataloader(self):
+        return DataLoader(
+            self.train_dataset, batch_size=self.batch_size, shuffle=True, num_workers=self.num_workers, drop_last=True
+        )
+
+    def val_dataloader(self):
+        return DataLoader(
+            self.val_dataset, batch_size=self.batch_size, shuffle=False, num_workers=self.num_workers, drop_last=True
+        )
+
+    def test_dataloader(self):
+        return DataLoader(
+            self.test_dataset, batch_size=self.batch_size, shuffle=False, num_workers=self.num_workers, drop_last=True
+        )

urban_inference/methane_urban_dataset.py

@@ -0,0 +1,82 @@
+import os
+import rasterio
+import torch
+from torchgeo.datasets import NonGeoDataset
+from torch.utils.data import DataLoader
+import torch.nn.functional as F
+import numpy as np
+import pandas as pd
+
+def min_max_normalize(data, new_min=0, new_max=1):
+    data = np.array(data, dtype=np.float32)  # Convert to NumPy array
+    
+    # Handle NaN, Inf values
+    data = np.nan_to_num(data, nan=np.nanmin(data), posinf=np.max(data), neginf=np.min(data))
+
+    old_min, old_max = np.min(data), np.max(data)
+    
+    if old_max == old_min:  # Prevent division by zero
+        return np.full_like(data, new_min, dtype=np.float32)  # Uniform array
+
+    return (data - old_min) / (old_max - old_min + 1e-10) * (new_max - new_min) + new_min
+
+class MethaneUrbanDataset(NonGeoDataset):
+    def __init__(self, root_dir, excel_file, paths, transform=None, mean=None, std=None):
+        super().__init__()
+        self.root_dir = root_dir
+        self.transform = transform
+        self.data_paths = []
+        self.mean = mean if mean else [0.485] * 12  # Default mean if not provided
+        self.std = std if std else [0.229] * 12  # Default std if not provided
+
+        # Collect paths for labelbinary.tif and sCube.tif in selected folders
+        for folder_name in paths:
+            subdir_path = next((os.path.join(root_dir, d) for d in os.listdir(root_dir) if d.startswith(folder_name) and os.path.isdir(os.path.join(root_dir, d))), None)
+            if subdir_path is not None and os.path.isdir(subdir_path):
+                label_path = os.path.join(subdir_path, 'hsi')
+                scube_path = os.path.join(subdir_path, 'hsi')
+                # print(scube_path)
+                if os.path.exists(scube_path):
+                    self.data_paths.append((label_path, scube_path))
+
+    def __len__(self):
+        return len(self.data_paths)
+
+    def __getitem__(self, idx):
+        label_path, scube_path = self.data_paths[idx]
+
+
+        # Load the sCube image (multi-band), drop the first band
+        with rasterio.open(scube_path) as scube_src:
+            scube_image = scube_src.read()  # Shape: [13, 512, 512]
+            scube_image = scube_image[[0,1,2,3,4,5,6,7,8,9,11,12], :, :]  # Drop first band → Shape: [12, 512, 512]
+            # scube_image = np.zeros((12, scube_src.height, scube_src.width), dtype=np.float32)  # Initialize with zeros
+            # selected_bands = [2, 3, 4, 5, 6, 7, 8]   # Bands to read
+            # for i, band in enumerate(selected_bands):
+                # scube_image[i, :, :] = scube_src.read(band + 1) 
+            # print(scube_path, scube_image.shape)
+
+        # Convert to PyTorch tensors
+        scube_tensor = torch.from_numpy(scube_image).float()  # Shape: [12, 512, 512]
+
+        # Resize to [12, 224, 224] and [224, 224] respectively
+        scube_tensor = F.interpolate(scube_tensor.unsqueeze(0), size=(224, 224), mode='bilinear', align_corners=False).squeeze(0)
+
+        scube_tensor = torch.nan_to_num(scube_tensor, nan=0.0)  # Replace NaNs with 0
+
+        # scube_tensor = (scube_tensor - scube_tensor.mean(dim=(1, 2), keepdim=True)) / (scube_tensor.std(dim=(1, 2), keepdim=True) + 1e-10)
+        
+        # Convert labels to binary
+        
+        contains_methane = torch.zeros(1).long()
+
+        # Convert to one-hot encoding
+        # one_hot_label = F.one_hot(contains_methane, num_classes=2).float()
+        
+        # Apply transformations (if any)
+        if self.transform:
+            transformed = self.transform(image=np.array(scube_tensor.permute(1, 2, 0)))
+            scube_tensor = transformed['image'].transpose(2, 0, 1)   # Convert back to [C, H, W]
+            
+
+        return {'S2L2A': scube_tensor, 'label': contains_methane, 'sample': scube_path.split('/')[2]}

urban_inference/urban_inference.py

@@ -0,0 +1,221 @@
+import argparse
+import logging
+import csv
+import random
+import warnings
+import time
+import os
+from pathlib import Path
+from typing import Dict, List, Tuple, Any, Optional
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+import albumentations as A
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+from rasterio.errors import NotGeoreferencedWarning
+
+# --- CRITICAL IMPORTS ---
+import terramind
+from terratorch.tasks import ClassificationTask
+
+# Local Imports
+from methane_urban_datamodule import MethaneUrbanDataModule
+
+# --- Configuration & Setup ---
+
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    datefmt='%Y-%m-%d %H:%M:%S'
+)
+logger = logging.getLogger(__name__)
+
+logging.getLogger("rasterio._env").setLevel(logging.ERROR)
+warnings.simplefilter("ignore", NotGeoreferencedWarning)
+warnings.filterwarnings("ignore", category=FutureWarning)
+
+def set_seed(seed: int = 42):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+def get_inference_transforms() -> A.Compose:
+    return None
+
+# --- Inference Class ---
+
+class UrbanInference:
+    def __init__(self, args: argparse.Namespace):
+        self.args = args
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.output_dir = Path(args.output_dir)
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        
+        logger.info(f"Initializing Inference on device: {self.device}")
+        
+        self.model = self._init_model()
+        self._load_checkpoint(args.checkpoint)
+
+    def _init_model(self) -> nn.Module:
+        model_args = dict(
+            backbone="terramind_v1_base",
+            backbone_pretrained=False, 
+            backbone_modalities=["S2L2A"], 
+            backbone_merge_method="mean",
+            decoder="UperNetDecoder",
+            decoder_scale_modules=True,
+            decoder_channels=256,
+            num_classes=2,
+            head_dropout=0.3,
+            necks=[
+                {"name": "ReshapeTokensToImage", "remove_cls_token": False},
+                {"name": "SelectIndices", "indices": [2, 5, 8, 11]},
+            ],
+        )
+
+        task = ClassificationTask(
+            model_args=model_args,
+            model_factory="EncoderDecoderFactory",
+            loss="ce",
+            ignore_index=-1
+        )
+        task.configure_models()
+        return task.model.to(self.device)
+
+    def _load_checkpoint(self, checkpoint_path: str):
+        path = Path(checkpoint_path)
+        if not path.exists():
+            raise FileNotFoundError(f"Checkpoint not found at {path}")
+            
+        logger.info(f"Loading weights from {path}...")
+        checkpoint = torch.load(path, map_location=self.device)
+        
+        if 'state_dict' in checkpoint:
+            state_dict = checkpoint['state_dict']
+        else:
+            state_dict = checkpoint
+            
+        self.model.load_state_dict(state_dict, strict=False)
+        self.model.eval()
+
+    def run_inference(self, dataloader: DataLoader, sample_names: List[str]):
+        """
+        Generates binary predictions and matches them with provided sample_names (folder names).
+        """
+        sample_results = {}
+        
+        logger.info(f"Starting inference on {len(sample_names)} samples...")
+        
+        # Iterator for sample names to match sequential predictions
+        name_iter = iter(sample_names)
+        
+        with torch.no_grad():
+            for batch in tqdm(dataloader, desc="Predicting"):
+                inputs = batch['S2L2A'].to(self.device)
+                
+                # Forward Pass
+                outputs = self.model(x={"S2L2A": inputs})
+                probabilities = torch.softmax(outputs.output, dim=1)
+                
+                # Get binary prediction (0 or 1)
+                predictions = torch.argmax(probabilities, dim=1)
+                batch_preds = predictions.cpu().numpy()
+                
+                # Assign Directory Names to Predictions
+                for pred in batch_preds:
+                    try:
+                        dir_name = next(name_iter)
+                        sample_results[dir_name] = int(pred)
+                    except StopIteration:
+                        logger.error("More predictions generated than sample names provided! Check dataloader sync.")
+                        break
+
+        # Check if we missed any samples
+        if len(sample_results) != len(sample_names):
+            logger.warning(f"Mismatch: Expected {len(sample_names)} results, got {len(sample_results)}.")
+
+        # Save CSV
+        self._save_results(sample_results)
+
+    def _save_results(self, results: Dict[str, int]):
+        csv_path = self.output_dir / "inference_predictions.csv"
+        with open(csv_path, mode='w', newline='') as f:
+            writer = csv.writer(f)
+            writer.writerow(['Sample_Directory', 'Prediction'])
+            for sample, pred in results.items():
+                writer.writerow([sample, pred])
+        logger.info(f"Predictions saved to {csv_path}")
+
+# --- Data Loading ---
+
+def get_dataloader_and_names(args) -> Tuple[DataLoader, List[str]]:
+    root_path = Path(args.root_dir)
+    if not root_path.exists():
+        raise FileNotFoundError(f"Data directory {args.root_dir} not found.")
+    
+    paths = None
+    if args.excel_file:
+        try:
+            df = pd.read_excel(args.excel_file)
+            # Ensure we get the directory name prefix if using 'Filename' column
+            paths = df['Filename'].apply(lambda x: str(x).split('_')[0]).tolist()
+            logger.info(f"Filtered {len(paths)} samples from Excel.")
+        except Exception as e:
+            logger.error(f"Error reading Excel: {e}")
+            raise
+
+    if paths is None:
+        # Fallback to all subdirectories
+        # SORTING is crucial here to match the sequential dataloader
+        paths = sorted([d.name for d in root_path.iterdir() if d.is_dir()])
+        logger.info(f"Found {len(paths)} samples in directory (Sorted).")
+
+    # Initialize DataModule
+    datamodule = MethaneUrbanDataModule(
+        data_root=args.root_dir,
+        excel_file=None, 
+        batch_size=args.batch_size,
+        paths=paths,
+        train_transform=None,
+        val_transform=get_inference_transforms(), 
+        test_transform=get_inference_transforms()
+    )
+    
+    # Setup for test stage
+    datamodule.paths = paths
+    datamodule.setup(stage="test")
+    
+    # Get loader (prefer test_dataloader)
+    loader = datamodule.test_dataloader() if hasattr(datamodule, 'test_dataloader') else datamodule.train_dataloader()
+    
+    return loader, paths
+
+# --- Main Execution ---
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Methane Urban Inference (Directory Names)")
+    
+    parser.add_argument('--root_dir', type=str, required=True, help='Root directory containing sample folders')
+    parser.add_argument('--checkpoint', type=str, required=True, help='Path to model checkpoint (.pth)')
+    parser.add_argument('--excel_file', type=str, help='Optional Excel file to filter specific samples')
+    parser.add_argument('--output_dir', type=str, default='./inference_results', help='Directory to save results')
+    parser.add_argument('--batch_size', type=int, default=1, help='Inference batch size')
+    parser.add_argument('--seed', type=int, default=42)
+    
+    return parser.parse_args()
+
+if __name__ == "__main__":
+    args = parse_args()
+    set_seed(args.seed)
+    
+    # 1. Prepare Data & Capture Directory Names
+    dataloader, sample_names = get_dataloader_and_names(args)
+    
+    # 2. Run Inference
+    engine = UrbanInference(args)
+    engine.run_inference(dataloader, sample_names)