kaggle_inference_notebook.py

#!/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!")