#!/usr/bin/env python3 """ CSIRO Image2Biomass Prediction - Kaggle Inference Notebook ============================================================ This notebook loads trained models and generates submission.csv. Requirements: - Trained model weights saved as a Kaggle dataset - No internet access (all models pre-downloaded) Expected model dataset structure: /kaggle/input/biomass-models/ fold_0/best_model.pth fold_1/best_model.pth ... training_info.json """ import os import sys import json import time import warnings from pathlib import Path from typing import Dict, List, Optional, Tuple import numpy as np import pandas as pd import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import Dataset, DataLoader from torch.cuda.amp import autocast from PIL import Image warnings.filterwarnings('ignore') os.system('pip install -q timm albumentations') import timm import albumentations as A from albumentations.pytorch import ToTensorV2 # ============================================================ # Configuration # ============================================================ class CFG: COMPETITION = 'csiro-biomass' DATA_DIR = Path(f'/kaggle/input/{COMPETITION}') MODEL_DIR = Path('/kaggle/input/biomass-models') # Your uploaded model weights OUTPUT_DIR = Path('/kaggle/working') BATCH_SIZE = 32 NUM_WORKERS = 2 N_TTA = 4 # Number of TTA augmentations DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu' TARGET_COLS = ['Dry_Green_g', 'Dry_Dead_g', 'Dry_Clover_g', 'GDM_g', 'Dry_Total_g'] IMAGENET_MEAN = (0.485, 0.456, 0.406) IMAGENET_STD = (0.229, 0.224, 0.225) BACKBONE_CONFIGS = { 'dinov2_small': {'name': 'vit_small_patch14_dinov2.lvd142m', 'feat_dim': 384}, 'dinov2_base': {'name': 'vit_base_patch14_dinov2.lvd142m', 'feat_dim': 768}, 'dinov2_large': {'name': 'vit_large_patch14_dinov2.lvd142m', 'feat_dim': 1024}, 'dinov2_base_reg': {'name': 'vit_base_patch14_reg4_dinov2.lvd142m', 'feat_dim': 768}, 'convnext_large': {'name': 'convnext_large.fb_in22k_ft_in1k', 'feat_dim': 1536}, 'convnextv2_large': {'name': 'convnextv2_large.fcmae_ft_in22k_in1k', 'feat_dim': 1536}, 'efficientnet_b4': {'name': 'efficientnet_b4.ra2_in1k', 'feat_dim': 1792}, 'swin_large': {'name': 'swin_large_patch4_window7_224.ms_in22k_ft_in1k', 'feat_dim': 1536}, 'eva02_large': {'name': 'eva02_large_patch14_448.mim_m38m_ft_in22k_in1k', 'feat_dim': 1024}, } # ============================================================ # Model Definition (must match training) # ============================================================ class BiomassModel(nn.Module): def __init__(self, backbone_name, num_targets=5, hidden_dim=512, dropout=0.3, pretrained=False, img_size=224, use_ndvi=False, separate_heads=False): super().__init__() self.use_ndvi = use_ndvi self.separate_heads = separate_heads kwargs = {'pretrained': pretrained, 'num_classes': 0} if 'vit' in backbone_name or 'dinov2' in backbone_name: kwargs['img_size'] = img_size self.backbone = timm.create_model(backbone_name, **kwargs) feat_dim = self.backbone.num_features if use_ndvi: self.ndvi_embed = nn.Sequential(nn.Linear(1, 32), nn.GELU(), nn.Linear(32, 64)) feat_dim += 64 if separate_heads: self.heads = nn.ModuleList([ nn.Sequential( nn.LayerNorm(feat_dim), nn.Dropout(dropout), nn.Linear(feat_dim, hidden_dim), nn.GELU(), nn.Dropout(dropout * 0.5), nn.Linear(hidden_dim, 1), ) for _ in range(num_targets) ]) else: self.head = nn.Sequential( nn.LayerNorm(feat_dim), nn.Dropout(dropout), nn.Linear(feat_dim, hidden_dim), nn.GELU(), nn.Dropout(dropout * 0.5), nn.Linear(hidden_dim, hidden_dim // 2), nn.GELU(), nn.Dropout(dropout * 0.3), nn.Linear(hidden_dim // 2, num_targets), ) def forward(self, x, ndvi=None): features = self.backbone(x) if self.use_ndvi and ndvi is not None: features = torch.cat([features, self.ndvi_embed(ndvi.unsqueeze(-1))], dim=-1) if self.separate_heads: return torch.cat([h(features) for h in self.heads], dim=-1) return self.head(features) # ============================================================ # Dataset # ============================================================ class TestDataset(Dataset): def __init__(self, image_dir, df, transform, use_ndvi=False): self.image_dir = Path(image_dir) self.df = df.reset_index(drop=True) self.transform = transform self.use_ndvi = use_ndvi def __len__(self): return len(self.df) def __getitem__(self, idx): row = self.df.iloc[idx] img_id = row['image_id'] if 'image_id' in row.index else row.name img_path = None for ext in ['.jpg', '.jpeg', '.png', '.JPG']: p = self.image_dir / f"{img_id}{ext}" if p.exists(): img_path = p break if img_path is None: candidates = list(self.image_dir.glob(f"{img_id}*")) img_path = candidates[0] if candidates else self.image_dir / f"{img_id}.jpg" img = np.array(Image.open(img_path).convert('RGB')) img_tensor = self.transform(image=img)['image'] result = {'image': img_tensor, 'image_id': str(img_id)} if self.use_ndvi and 'NDVI' in self.df.columns: result['ndvi'] = torch.tensor(float(row['NDVI']), dtype=torch.float32) return result # ============================================================ # TTA Transforms # ============================================================ def get_tta_transforms(img_size=224, n_tta=4): tfms = [] # 0: Standard center crop tfms.append(A.Compose([ A.Resize(height=int(img_size * 1.14), width=int(img_size * 1.14)), A.CenterCrop(height=img_size, width=img_size), A.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD), ToTensorV2(), ])) # 1: HFlip tfms.append(A.Compose([ A.Resize(height=int(img_size * 1.14), width=int(img_size * 1.14)), A.CenterCrop(height=img_size, width=img_size), A.HorizontalFlip(p=1.0), A.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD), ToTensorV2(), ])) # 2: VFlip tfms.append(A.Compose([ A.Resize(height=int(img_size * 1.14), width=int(img_size * 1.14)), A.CenterCrop(height=img_size, width=img_size), A.VerticalFlip(p=1.0), A.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD), ToTensorV2(), ])) # 3: Both flips tfms.append(A.Compose([ A.Resize(height=int(img_size * 1.14), width=int(img_size * 1.14)), A.CenterCrop(height=img_size, width=img_size), A.HorizontalFlip(p=1.0), A.VerticalFlip(p=1.0), A.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD), ToTensorV2(), ])) return tfms[:n_tta] # ============================================================ # Inference Functions # ============================================================ def load_model(ckpt_path, device): ckpt = torch.load(ckpt_path, map_location=device, weights_only=False) args = ckpt.get('args', {}) # Resolve backbone name backbone_key = args.get('backbone', 'vit_base_patch14_dinov2.lvd142m') if backbone_key in BACKBONE_CONFIGS: backbone_name = BACKBONE_CONFIGS[backbone_key]['name'] else: backbone_name = backbone_key img_size = args.get('img_size', 224) model = BiomassModel( backbone_name=backbone_name, num_targets=5, hidden_dim=args.get('hidden_dim', 512), dropout=args.get('dropout', 0.3), pretrained=False, img_size=img_size, use_ndvi=args.get('use_ndvi', False), separate_heads=args.get('separate_heads', False), ) model.load_state_dict(ckpt['model_state_dict']) model = model.to(device).eval() return model, args @torch.no_grad() def predict(model, loader, device, log_transform=True): model.eval() preds_list, ids_list = [], [] for batch in loader: images = batch['image'].to(device) ndvi = batch.get('ndvi', None) if ndvi is not None: ndvi = ndvi.to(device) with autocast(dtype=torch.float16): preds = model(images, ndvi) preds_list.append(preds.cpu().numpy()) ids_list.extend(batch['image_id']) preds = np.concatenate(preds_list) if log_transform: preds = np.expm1(preds) return preds, ids_list def predict_tta(model, test_df, img_dir, device, img_size, log_transform, use_ndvi, batch_size, num_workers, n_tta): tta_tfms = get_tta_transforms(img_size, n_tta) all_preds = [] ids = None for i, tfm in enumerate(tta_tfms): ds = TestDataset(img_dir, test_df, tfm, use_ndvi) loader = DataLoader(ds, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True) p, image_ids = predict(model, loader, device, log_transform) all_preds.append(p) if ids is None: ids = image_ids return np.mean(all_preds, axis=0), ids # ============================================================ # Main Inference # ============================================================ device = torch.device(CFG.DEVICE) print(f"Device: {device}") # Find data for alt in ['/kaggle/input/csiro-biomass', '/kaggle/input/csiro-image2biomass-prediction', '/kaggle/input/csiro-image2biomass']: if Path(alt).exists(): CFG.DATA_DIR = Path(alt) break # Find model weights for alt in ['/kaggle/input/biomass-models', '/kaggle/input/biomass-weights', '/kaggle/working']: if Path(alt).exists() and list(Path(alt).glob('fold_*')): CFG.MODEL_DIR = Path(alt) break print(f"Data: {CFG.DATA_DIR}") print(f"Models: {CFG.MODEL_DIR}") # Load test data test_csv = None for fname in ['test.csv', 'Test.csv']: if (CFG.DATA_DIR / fname).exists(): test_csv = CFG.DATA_DIR / fname break test_df = pd.read_csv(test_csv) print(f"Test samples: {len(test_df)}") # Find test images test_img_dir = None for d in ['test_images', 'test', 'images/test']: if (CFG.DATA_DIR / d).exists(): test_img_dir = CFG.DATA_DIR / d break print(f"Test images: {test_img_dir}") # Find fold models fold_dirs = sorted(CFG.MODEL_DIR.glob('fold_*')) print(f"Found {len(fold_dirs)} fold models") # Ensemble prediction all_fold_preds = [] image_ids = None for fold_dir in fold_dirs: ckpt_path = fold_dir / 'best_model.pth' if not ckpt_path.exists(): continue print(f"\nLoading {ckpt_path}...") model, args = load_model(str(ckpt_path), device) img_size = args.get('img_size', 224) log_transform = args.get('log_transform', True) use_ndvi = args.get('use_ndvi', False) preds, ids = predict_tta( model, test_df, str(test_img_dir), device, img_size=img_size, log_transform=log_transform, use_ndvi=use_ndvi, batch_size=CFG.BATCH_SIZE, num_workers=CFG.NUM_WORKERS, n_tta=CFG.N_TTA, ) all_fold_preds.append(preds) if image_ids is None: image_ids = ids print(f" Mean predictions: {preds.mean(axis=0)}") del model torch.cuda.empty_cache() # Average across folds ensemble_preds = np.mean(all_fold_preds, axis=0) ensemble_preds = np.clip(ensemble_preds, 0, None) # Post-process: ensure total >= component sum comp_sum = ensemble_preds[:, 0] + ensemble_preds[:, 1] + ensemble_preds[:, 2] mask = ensemble_preds[:, 4] < comp_sum ensemble_preds[mask, 4] = comp_sum[mask] print(f"\nEnsemble predictions summary:") for i, name in enumerate(TARGET_COLS): col = ensemble_preds[:, i] print(f" {name}: mean={col.mean():.2f}, std={col.std():.2f}, " f"min={col.min():.2f}, max={col.max():.2f}") # Create submission rows = [] for i, img_id in enumerate(image_ids): for j, target_name in enumerate(TARGET_COLS): rows.append({ 'sample_id': f"{img_id}__{target_name}", 'target': float(max(0, ensemble_preds[i, j])), }) submission = pd.DataFrame(rows) submission.to_csv('submission.csv', index=False) print(f"\nSubmission saved: submission.csv ({len(submission)} rows)") print(submission.head(10)) # Verify format assert submission.columns.tolist() == ['sample_id', 'target'] assert len(submission) == len(test_df) * 5 print("\n✅ Submission format verified!")