{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":112509,"databundleVersionId":14254895,"sourceType":"competition"}],"dockerImageVersionId":31193,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.nn.parallel import DataParallel\nimport timm\nfrom PIL import Image\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nfrom sklearn.model_selection import KFold\nfrom tqdm import tqdm\nimport warnings\nwarnings.filterwarnings('ignore')\n\n# Configuration\nclass Config:\n    BASE_PATH = '/kaggle/input/csiro-biomass'\n    TRAIN_CSV = f'{BASE_PATH}/train.csv'\n    TRAIN_IMG_DIR = f'{BASE_PATH}/'\n    \n    MODEL_NAME = 'eva02_base_patch14_448.mim_in22k_ft_in22k_in1k'\n    IMAGE_SIZE = (448, 448)  # (height, width) - EVA-02 requires 448x448\n    BATCH_SIZE = 16\n    NUM_EPOCHS = 30\n    LEARNING_RATE = 1e-4\n    NUM_WORKERS = 4\n    K_FOLDS = 5\n    SEED = 42\n    DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'\n    NUM_TARGETS = 3  # Dry_Clover_g, Dry_Dead_g, Dry_Green_g\n    MODEL_SAVE_PATH = './biomass_models'\n\n# Set seed for reproducibility\ndef set_seed(seed):\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n\nset_seed(Config.SEED)\n\n# Create model save directory\nos.makedirs(Config.MODEL_SAVE_PATH, exist_ok=True)\n\n# Dataset\nclass BiomassDataset(Dataset):\n    def __init__(self, df, img_dir, transform=None):\n        # Group by image_path to get all targets for each image\n        self.image_groups = df.groupby('image_path')\n        self.image_paths = list(self.image_groups.groups.keys())\n        self.img_dir = img_dir\n        self.transform = transform\n        \n    def __len__(self):\n        return len(self.image_paths)\n    \n    def __getitem__(self, idx):\n        img_path = self.image_paths[idx]\n        full_path = os.path.join(self.img_dir, img_path)\n        \n        # Load image\n        image = Image.open(full_path).convert('RGB')\n        image = np.array(image)\n        \n        # Get targets for this image\n        img_data = self.image_groups.get_group(img_path)\n        \n        # Extract the 3 main targets we need to predict\n        targets = np.zeros(Config.NUM_TARGETS, dtype=np.float32)\n        for _, row in img_data.iterrows():\n            if row['target_name'] == 'Dry_Clover_g':\n                targets[0] = row['target']\n            elif row['target_name'] == 'Dry_Dead_g':\n                targets[1] = row['target']\n            elif row['target_name'] == 'Dry_Green_g':\n                targets[2] = row['target']\n        \n        # Apply transforms\n        if self.transform:\n            transformed = self.transform(image=image)\n            image = transformed['image']\n        \n        return image, targets\n\n# Augmentations\ndef get_train_transforms():\n    return A.Compose([\n        A.Resize(Config.IMAGE_SIZE[0], Config.IMAGE_SIZE[1]),\n        A.HorizontalFlip(p=0.5),\n        A.VerticalFlip(p=0.5),\n        A.Rotate(limit=15, p=0.5),\n        A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=0.5),\n        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n        ToTensorV2()\n    ])\n\ndef get_valid_transforms():\n    return A.Compose([\n        A.Resize(Config.IMAGE_SIZE[0], Config.IMAGE_SIZE[1]),\n        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n        ToTensorV2()\n    ])\n\n# Model\nclass BiomassModel(nn.Module):\n    def __init__(self, model_name, num_targets=3, pretrained=True):\n        super(BiomassModel, self).__init__()\n        self.backbone = timm.create_model(model_name, pretrained=pretrained, num_classes=0)\n        \n        # Get the number of features from the backbone\n        with torch.no_grad():\n            dummy_input = torch.randn(1, 3, Config.IMAGE_SIZE[0], Config.IMAGE_SIZE[1])\n            features = self.backbone(dummy_input)\n            num_features = features.shape[1]\n        \n        # Regression head\n        self.head = nn.Sequential(\n            nn.Linear(num_features, 512),\n            nn.ReLU(),\n            nn.Dropout(0.3),\n            nn.Linear(512, 256),\n            nn.ReLU(),\n            nn.Dropout(0.2),\n            nn.Linear(256, num_targets)\n        )\n        \n    def forward(self, x):\n        features = self.backbone(x)\n        output = self.head(features)\n        return output\n\n# Training function\ndef train_epoch(model, dataloader, criterion, optimizer, device):\n    model.train()\n    running_loss = 0.0\n    \n    pbar = tqdm(dataloader, desc='Training')\n    for images, targets in pbar:\n        images = images.to(device)\n        targets = targets.to(device)\n        \n        optimizer.zero_grad()\n        outputs = model(images)\n        loss = criterion(outputs, targets)\n        loss.backward()\n        optimizer.step()\n        \n        running_loss += loss.item() * images.size(0)\n        pbar.set_postfix({'loss': loss.item()})\n    \n    epoch_loss = running_loss / len(dataloader.dataset)\n    return epoch_loss\n\n# Validation function\ndef validate_epoch(model, dataloader, criterion, device):\n    model.eval()\n    running_loss = 0.0\n    \n    with torch.no_grad():\n        pbar = tqdm(dataloader, desc='Validation')\n        for images, targets in pbar:\n            images = images.to(device)\n            targets = targets.to(device)\n            \n            outputs = model(images)\n            loss = criterion(outputs, targets)\n            \n            running_loss += loss.item() * images.size(0)\n            pbar.set_postfix({'loss': loss.item()})\n    \n    epoch_loss = running_loss / len(dataloader.dataset)\n    return epoch_loss\n\n# Main training loop with K-Fold\ndef train_kfold():\n    # Load data\n    df = pd.read_csv(Config.TRAIN_CSV)\n    \n    # Get unique images for splitting\n    unique_images = df['image_path'].unique()\n    \n    # K-Fold split\n    kfold = KFold(n_splits=Config.K_FOLDS, shuffle=True, random_state=Config.SEED)\n    \n    fold_scores = []\n    \n    for fold, (train_idx, val_idx) in enumerate(kfold.split(unique_images)):\n        print(f'\\n{\"=\"*50}')\n        print(f'Fold {fold + 1}/{Config.K_FOLDS}')\n        print(f'{\"=\"*50}\\n')\n        \n        # Split data\n        train_images = unique_images[train_idx]\n        val_images = unique_images[val_idx]\n        \n        train_df = df[df['image_path'].isin(train_images)].reset_index(drop=True)\n        val_df = df[df['image_path'].isin(val_images)].reset_index(drop=True)\n        \n        # Create datasets\n        train_dataset = BiomassDataset(train_df, Config.TRAIN_IMG_DIR, get_train_transforms())\n        val_dataset = BiomassDataset(val_df, Config.TRAIN_IMG_DIR, get_valid_transforms())\n        \n        # Create dataloaders\n        train_loader = DataLoader(\n            train_dataset, \n            batch_size=Config.BATCH_SIZE, \n            shuffle=True, \n            num_workers=Config.NUM_WORKERS,\n            pin_memory=True\n        )\n        val_loader = DataLoader(\n            val_dataset, \n            batch_size=Config.BATCH_SIZE, \n            shuffle=False, \n            num_workers=Config.NUM_WORKERS,\n            pin_memory=True\n        )\n        \n        # Create model\n        model = BiomassModel(Config.MODEL_NAME, num_targets=Config.NUM_TARGETS)\n        \n        # Use DataParallel for multi-GPU training\n        if torch.cuda.device_count() > 1:\n            print(f\"Using {torch.cuda.device_count()} GPUs!\")\n            model = DataParallel(model)\n        \n        model = model.to(Config.DEVICE)\n        \n        # Loss and optimizer\n        criterion = nn.MSELoss()\n        optimizer = torch.optim.AdamW(model.parameters(), lr=Config.LEARNING_RATE)\n        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=Config.NUM_EPOCHS)\n        \n        # Training loop\n        best_val_loss = float('inf')\n        \n        for epoch in range(Config.NUM_EPOCHS):\n            print(f'\\nEpoch {epoch + 1}/{Config.NUM_EPOCHS}')\n            \n            train_loss = train_epoch(model, train_loader, criterion, optimizer, Config.DEVICE)\n            val_loss = validate_epoch(model, val_loader, criterion, Config.DEVICE)\n            \n            scheduler.step()\n            \n            print(f'Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f}')\n            \n            # Save best model\n            if val_loss < best_val_loss:\n                best_val_loss = val_loss\n                model_path = os.path.join(Config.MODEL_SAVE_PATH, f'best_model_fold{fold+1}.pth')\n                \n                # Save the model (handle DataParallel)\n                if isinstance(model, DataParallel):\n                    torch.save(model.module.state_dict(), model_path)\n                else:\n                    torch.save(model.state_dict(), model_path)\n                \n                print(f'Saved best model with Val Loss: {best_val_loss:.4f}')\n        \n        fold_scores.append(best_val_loss)\n        print(f'\\nFold {fold + 1} Best Val Loss: {best_val_loss:.4f}')\n    \n    print(f'\\n{\"=\"*50}')\n    print(f'Average Val Loss across all folds: {np.mean(fold_scores):.4f} (+/- {np.std(fold_scores):.4f})')\n    print(f'{\"=\"*50}')\n    \n    return fold_scores\n\n# Run training\nif __name__ == '__main__':\n    print(f'Using device: {Config.DEVICE}')\n    print(f'Number of GPUs: {torch.cuda.device_count()}')\n    \n    fold_scores = train_kfold()\n    \n    print('\\nTraining completed!')\n    print(f'Models saved in: {Config.MODEL_SAVE_PATH}')","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-11-12T00:11:05.424253Z","iopub.execute_input":"2025-11-12T00:11:05.424537Z","iopub.status.idle":"2025-11-12T01:58:59.328426Z","shell.execute_reply.started":"2025-11-12T00:11:05.424516Z","shell.execute_reply":"2025-11-12T01:58:59.327524Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}