#!/usr/bin/env python3 """ CSIRO Image2Biomass Prediction - Inference Pipeline ===================================================== Generates submission.csv from trained model ensemble. Supports: - Single model or ensemble of fold models - Test-Time Augmentation (TTA) - Multi-backbone ensembling - Post-processing (non-negative, consistency check) For Kaggle submission notebook: - All models must be saved in a Kaggle dataset - No internet access during inference """ import os import sys import json import glob import time 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 import timm try: import albumentations as A from albumentations.pytorch import ToTensorV2 HAS_ALBUMENTATIONS = True except ImportError: HAS_ALBUMENTATIONS = False from torchvision import transforms from PIL import Image # ============================================================ # Constants (must match training) # ============================================================ TARGET_COLS = ['Dry_Green_g', 'Dry_Dead_g', 'Dry_Clover_g', 'GDM_g', 'Dry_Total_g'] TARGET_WEIGHTS = [0.1, 0.1, 0.1, 0.2, 0.5] 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 (same as training) # ============================================================ class BiomassModel(nn.Module): def __init__( self, backbone_name: str = 'vit_base_patch14_dinov2.lvd142m', num_targets: int = 5, hidden_dim: int = 512, dropout: float = 0.3, pretrained: bool = False, # No pretrained for inference img_size: int = 224, use_ndvi: bool = False, separate_heads: bool = False, grad_checkpointing: bool = False, ): super().__init__() self.use_ndvi = use_ndvi self.separate_heads = separate_heads self.num_targets = num_targets 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: ndvi_feats = self.ndvi_embed(ndvi.unsqueeze(-1)) features = torch.cat([features, ndvi_feats], dim=-1) if self.separate_heads: outputs = [head(features) for head in self.heads] return torch.cat(outputs, dim=-1) else: return self.head(features) # ============================================================ # Dataset # ============================================================ class TestDataset(Dataset): def __init__(self, image_dir, df, transform=None, img_size=224, use_ndvi=False): self.image_dir = Path(image_dir) self.df = df.reset_index(drop=True) self.transform = transform self.img_size = img_size 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 = self.image_dir / f"{img_id}.jpg" if not img_path.exists(): img_path = self.image_dir / f"{img_id}.png" if not img_path.exists(): candidates = list(self.image_dir.glob(f"{img_id}.*")) if candidates: img_path = candidates[0] img = Image.open(img_path).convert('RGB') img = np.array(img) if self.transform is not None: if HAS_ALBUMENTATIONS: augmented = self.transform(image=img) img_tensor = augmented['image'] else: img = Image.fromarray(img) img_tensor = self.transform(img) else: img_tensor = transforms.ToTensor()(Image.fromarray(img)) result = {'image': img_tensor, 'image_id': str(img_id)} if self.use_ndvi and 'NDVI' in self.df.columns: result['ndvi'] = torch.tensor(row['NDVI'], dtype=torch.float32) return result # ============================================================ # Transforms # ============================================================ def get_val_transforms(img_size=224): if HAS_ALBUMENTATIONS: return 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(), ]) else: return transforms.Compose([ transforms.Resize(int(img_size * 1.14)), transforms.CenterCrop(img_size), transforms.ToTensor(), transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD), ]) def get_tta_transforms(img_size=224): """Returns list of TTA transforms.""" ttas = [get_val_transforms(img_size)] if HAS_ALBUMENTATIONS: # Horizontal flip ttas.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(), ])) # Vertical flip ttas.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(), ])) # Both flips ttas.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 ttas # ============================================================ # Inference # ============================================================ def load_model(checkpoint_path: str, device: torch.device) -> BiomassModel: """Load a trained model from checkpoint.""" checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False) args = checkpoint.get('args', {}) backbone_key = args.get('backbone', 'dinov2_base') backbone_cfg = BACKBONE_CONFIGS[backbone_key] img_size = args.get('img_size', None) or backbone_cfg.get('default_size', 224) model = BiomassModel( backbone_name=backbone_cfg['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(checkpoint['model_state_dict']) model = model.to(device) model.eval() print(f"Loaded model from {checkpoint_path}") print(f" Backbone: {backbone_key}, img_size: {img_size}") print(f" Best R²: {checkpoint.get('weighted_r2', 'N/A')}") return model, args @torch.no_grad() def predict_single(model, loader, device, log_transform=True): """Predict with a single model.""" model.eval() all_preds = [] all_ids = [] 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) all_preds.append(preds.cpu().numpy()) all_ids.extend(batch['image_id']) all_preds = np.concatenate(all_preds, axis=0) if log_transform: all_preds = np.expm1(all_preds) return all_preds, all_ids def predict_with_tta(model, test_df, image_dir, device, img_size=224, log_transform=True, batch_size=32, num_workers=4, use_ndvi=False, n_tta=4): """Predict with Test-Time Augmentation.""" tta_transforms = get_tta_transforms(img_size)[:n_tta] all_tta_preds = [] image_ids = None for tta_idx, tfm in enumerate(tta_transforms): dataset = TestDataset(image_dir, test_df, transform=tfm, img_size=img_size, use_ndvi=use_ndvi) loader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True) preds, ids = predict_single(model, loader, device, log_transform) all_tta_preds.append(preds) if image_ids is None: image_ids = ids print(f" TTA {tta_idx}: mean pred = {preds.mean():.2f}") # Average TTA predictions avg_preds = np.mean(all_tta_preds, axis=0) return avg_preds, image_ids def ensemble_predict( model_dirs: List[str], test_df: pd.DataFrame, image_dir: str, device: torch.device, batch_size: int = 32, num_workers: int = 4, n_tta: int = 4, model_weights: Optional[List[float]] = None, ): """ Ensemble predictions from multiple model checkpoints. Args: model_dirs: List of model checkpoint directories (each containing best_model.pth) test_df: Test dataframe image_dir: Test image directory device: torch device batch_size: Batch size for inference num_workers: DataLoader workers n_tta: Number of TTA augmentations model_weights: Optional weights for each model (default: equal weights) """ all_preds = [] image_ids = None for i, model_dir in enumerate(model_dirs): model_dir = Path(model_dir) # Find checkpoint ckpt_path = model_dir / 'best_model.pth' if not ckpt_path.exists(): # Search for any .pth file ckpt_files = list(model_dir.glob('*.pth')) if ckpt_files: ckpt_path = ckpt_files[0] else: print(f"No checkpoint found in {model_dir}, skipping") continue print(f"\nModel {i+1}/{len(model_dirs)}: {ckpt_path}") model, args = load_model(str(ckpt_path), device) backbone_key = args.get('backbone', 'dinov2_base') backbone_cfg = BACKBONE_CONFIGS[backbone_key] img_size = args.get('img_size', None) or backbone_cfg.get('default_size', 224) log_transform = args.get('log_transform', True) use_ndvi = args.get('use_ndvi', False) preds, ids = predict_with_tta( model, test_df, image_dir, device, img_size=img_size, log_transform=log_transform, batch_size=batch_size, num_workers=num_workers, use_ndvi=use_ndvi, n_tta=n_tta, ) all_preds.append(preds) if image_ids is None: image_ids = ids # Free memory del model torch.cuda.empty_cache() # Weighted ensemble if model_weights is None: model_weights = [1.0 / len(all_preds)] * len(all_preds) else: model_weights = np.array(model_weights) / sum(model_weights) ensemble_preds = np.zeros_like(all_preds[0]) for pred, weight in zip(all_preds, model_weights): ensemble_preds += weight * pred return ensemble_preds, image_ids # ============================================================ # Post-processing # ============================================================ def postprocess_predictions(preds: np.ndarray) -> np.ndarray: """ Post-process predictions: 1. Clip to non-negative values 2. Ensure Dry_Total >= sum of components (optional soft constraint) """ preds = np.clip(preds, 0, None) # Soft consistency: if Total < sum of components, adjust component_sum = preds[:, 0] + preds[:, 1] + preds[:, 2] total = preds[:, 4] # Where total is less than component sum, nudge total up mask = total < component_sum if mask.any(): preds[mask, 4] = component_sum[mask] * 1.0 # Set total = component_sum return preds def create_submission(preds: np.ndarray, image_ids: List[str], output_path: str = 'submission.csv'): """Create submission CSV in competition format.""" 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, preds[i, j])), }) sub_df = pd.DataFrame(rows) sub_df.to_csv(output_path, index=False) print(f"Submission saved: {output_path} ({len(sub_df)} rows)") # Validation n_images = len(image_ids) expected_rows = n_images * 5 assert len(sub_df) == expected_rows, f"Expected {expected_rows} rows, got {len(sub_df)}" # Print summary for j, name in enumerate(TARGET_COLS): col_preds = preds[:, j] print(f" {name}: mean={col_preds.mean():.2f}, std={col_preds.std():.2f}, " f"min={col_preds.min():.2f}, max={col_preds.max():.2f}") return sub_df # ============================================================ # Main inference entry point # ============================================================ def main(): import argparse parser = argparse.ArgumentParser() parser.add_argument('--data_dir', type=str, required=True) parser.add_argument('--model_dir', type=str, required=True, help='Directory with fold_0/best_model.pth, fold_1/best_model.pth, etc.') parser.add_argument('--output', type=str, default='submission.csv') parser.add_argument('--batch_size', type=int, default=32) parser.add_argument('--num_workers', type=int, default=4) parser.add_argument('--n_tta', type=int, default=4) parser.add_argument('--no_tta', action='store_true') args = parser.parse_args() device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print(f"Device: {device}") # Load test data data_dir = Path(args.data_dir) test_df = pd.read_csv(data_dir / 'test.csv') test_img_dir = data_dir / 'test_images' if not test_img_dir.exists(): test_img_dir = data_dir / 'test' print(f"Test samples: {len(test_df)}") # Find all fold checkpoints model_dir = Path(args.model_dir) fold_dirs = sorted(model_dir.glob('fold_*')) if not fold_dirs: # Single model fold_dirs = [model_dir] print(f"Found {len(fold_dirs)} fold model(s)") # Ensemble predict n_tta = 1 if args.no_tta else args.n_tta preds, image_ids = ensemble_predict( model_dirs=[str(d) for d in fold_dirs], test_df=test_df, image_dir=str(test_img_dir), device=device, batch_size=args.batch_size, num_workers=args.num_workers, n_tta=n_tta, ) # Post-process preds = postprocess_predictions(preds) # Create submission create_submission(preds, image_ids, args.output) print("\nDone!") if __name__ == '__main__': main()