{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.12.12","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":31259,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# CSIRO Biomass Inference with ViT\n\nThis notebook loads the trained ViT models (saved as `vit_fold_X_best.pth`) and generates a `submission.csv`.","metadata":{}},{"cell_type":"code","source":"import math\nimport os\nimport random\nfrom pathlib import Path\n\nimport numpy as np\nimport pandas as pd\nimport torch\nfrom PIL import Image\nfrom sklearn.model_selection import KFold\nfrom torch import nn\nfrom torch.utils.data import DataLoader, Dataset\nfrom torchvision import models, transforms\nfrom tqdm import tqdm  # Changed from tqdm.auto to tqdm to avoid widget errors\n\n# === Configuration ===\nDATA_ROOT = Path('/kaggle/input/csiro-biomass')\nif not DATA_ROOT.exists():\n    # Fallback to local path or current directory\n    local_path = Path('/Users/fukamikaito/kaggle/csiro-biomass')\n    if local_path.exists():\n        DATA_ROOT = local_path\n    else:\n        DATA_ROOT = Path('.')\n\nTRAIN_CSV = DATA_ROOT / 'train.csv'\nTEST_CSV = DATA_ROOT / 'test.csv'\nSAMPLE_SUBMISSION = DATA_ROOT / 'sample_submission.csv'\nOUTPUT_CSV = Path('/kaggle/working/submission.csv') if Path('/kaggle/working').exists() else Path('submission.csv')\n\n# Hyperparameters\nFOLDS = 5\nEPOCHS = 150\nBATCH_SIZE = 32\nIMG_SIZE = 224\nLR = 1e-4\nNUM_WORKERS = 0\nSEED = 42\n\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(f\"Using device: {DEVICE}\")\n\n# === Utils ===\ndef seed_everything(seed: int):\n    random.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\nseed_everything(SEED)\n\nTARGET_ORDER = [\n    \"Dry_Clover_g\",\n    \"Dry_Dead_g\",\n    \"Dry_Green_g\",\n    \"Dry_Total_g\",\n    \"GDM_g\",\n]\nTARGET_TO_IDX = {name: idx for idx, name in enumerate(TARGET_ORDER)}\n\n# === Dataset ===\nclass ImageTargetDataset(Dataset):\n    def __init__(self, image_paths, targets, data_root, transform=None):\n        self.image_paths = list(image_paths)\n        self.targets = None if targets is None else torch.tensor(targets, dtype=torch.float32)\n        self.data_root = Path(data_root)\n        self.transform = transform\n\n    def __len__(self):\n        return len(self.image_paths)\n\n    def __getitem__(self, idx):\n        path = self.data_root / self.image_paths[idx]\n        # Basic error handling for image loading\n        try:\n            image = Image.open(path).convert('RGB')\n            if self.transform:\n                image = self.transform(image)\n        except Exception as e:\n            image = torch.zeros((3, IMG_SIZE, IMG_SIZE))\n        \n        if self.targets is None:\n            return image, self.image_paths[idx]\n        else:\n            return image, self.targets[idx], self.image_paths[idx]\n\ndef get_transforms(image_size: int):\n    train_tfms = transforms.Compose([\n        transforms.RandomResizedCrop(image_size, scale=(0.8, 1.0)),\n        transforms.RandomHorizontalFlip(),\n        transforms.ToTensor(),\n        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n    ])\n    valid_tfms = transforms.Compose([\n        transforms.Resize((image_size, image_size)),\n        transforms.ToTensor(),\n        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n    ])\n    return train_tfms, valid_tfms\n\n# === Model ===\ndef build_model(num_outputs: int):\n    # ViT model (weights=None for no internet/pretraining)\n    model = models.vit_b_16(weights=None)\n    # Replace head\n    if hasattr(model.heads, 'head'):\n        in_features = model.heads.head.in_features\n        model.heads.head = nn.Linear(in_features, num_outputs)\n    else:\n        # Generic fallback for ViT structures\n        model.heads.head = nn.Linear(768, num_outputs)\n    return model\n\n# === Training Functions ===\ndef train_one_epoch(model, loader, optimizer, criterion, device):\n    model.train()\n    running_loss = 0.0\n    for images, targets, _ in tqdm(loader, desc='train', leave=False):\n        images = images.to(device, non_blocking=True)\n        targets = targets.to(device, non_blocking=True)\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    return running_loss / len(loader.dataset)\n\ndef predict(model, loader, device):\n    model.eval()\n    preds = []\n    paths = []\n    with torch.no_grad():\n        for batch in tqdm(loader, desc='predict', leave=False):\n            images = batch[0].to(device, non_blocking=True)\n            outputs = model(images)\n            preds.append(outputs.detach().cpu().numpy())\n            paths.extend(batch[-1])\n    preds = np.concatenate(preds, axis=0)\n    return preds, paths\n\ndef rmse(preds, targets):\n    return math.sqrt(np.mean((preds - targets) ** 2))\n\n# === Main Process ===\ndef main():\n    fold_models = []\n    \n    # 1. Training Phase\n    if TRAIN_CSV.exists():\n        print(\"Loading training data...\")\n        train_raw = pd.read_csv(TRAIN_CSV)\n        train_df = train_raw.pivot_table(index='image_path', columns='target_name', values='target')\n        train_df = train_df[TARGET_ORDER].reset_index()\n        \n        kf = KFold(n_splits=FOLDS, shuffle=True, random_state=SEED)\n        train_tfms, valid_tfms = get_transforms(IMG_SIZE)\n        \n        for fold_id, (train_idx, valid_idx) in enumerate(kf.split(train_df), start=1):\n            print(f\"\\n=== Fold {fold_id}/{FOLDS} ===\")\n            \n            train_sub = train_df.iloc[train_idx]\n            valid_sub = train_df.iloc[valid_idx]\n            \n            train_ds = ImageTargetDataset(\n                train_sub[\"image_path\"].values, \n                train_sub[TARGET_ORDER].values, \n                DATA_ROOT, \n                transform=train_tfms\n            )\n            valid_ds = ImageTargetDataset(\n                valid_sub[\"image_path\"].values, \n                valid_sub[TARGET_ORDER].values, \n                DATA_ROOT, \n                transform=valid_tfms\n            )\n            \n            train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=NUM_WORKERS, pin_memory=False)\n            valid_loader = DataLoader(valid_ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=NUM_WORKERS, pin_memory=False)\n            \n            model = build_model(len(TARGET_ORDER)).to(DEVICE)\n            criterion = nn.MSELoss()\n            optimizer = torch.optim.AdamW(model.parameters(), lr=LR)\n            \n            for epoch in range(1, EPOCHS + 1):\n                train_loss = train_one_epoch(model, train_loader, optimizer, criterion, DEVICE)\n                val_preds, _ = predict(model, valid_loader, DEVICE)\n                val_rmse = rmse(val_preds, valid_sub[TARGET_ORDER].values)\n                print(f\"Epoch {epoch} - Loss: {train_loss:.4f}, Val RMSE: {val_rmse:.4f}\")\n            \n            fold_models.append(model)\n    else:\n        print(\"Train CSV not found. Skipping training.\")\n\n    # 2. Inference Phase\n    if TEST_CSV.exists():\n        print(\"\\nStarting inference...\")\n        test_raw = pd.read_csv(TEST_CSV)\n        unique_test_images = test_raw[['image_path']].drop_duplicates().reset_index(drop=True)\n        \n        _, valid_tfms = get_transforms(IMG_SIZE)\n        test_ds = ImageTargetDataset(unique_test_images['image_path'].values, None, DATA_ROOT, transform=valid_tfms)\n        test_loader = DataLoader(test_ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=NUM_WORKERS, pin_memory=False)\n        \n        if not fold_models:\n            print(\"No models available. Initializing dummy model for submission check.\")\n            model = build_model(len(TARGET_ORDER)).to(DEVICE)\n            fold_models.append(model)\n            \n        test_preds_sum = None\n        for i, model in enumerate(fold_models):\n            print(f\"Predicting with model {i+1}...\")\n            preds, paths = predict(model, test_loader, DEVICE)\n            if test_preds_sum is None:\n                test_preds_sum = preds\n            else:\n                test_preds_sum += preds\n        \n        test_preds_avg = test_preds_sum / len(fold_models)\n        \n        # Map predictions back to sample_id\n        preds_map = {path: pred for path, pred in zip(unique_test_images['image_path'].values, test_preds_avg)}\n        \n        submission = test_raw[['sample_id', 'image_path', 'target_name']].copy()\n        \n        def get_pred_value(row):\n            path = row['image_path']\n            if path not in preds_map: return 0.0\n            \n            vec = preds_map[path]\n            target_name = row['target_name']\n            idx = TARGET_TO_IDX.get(target_name)\n            if idx is None: return 0.0\n            return float(vec[idx])\n\n        submission['target'] = submission.apply(get_pred_value, axis=1)\n        submission = submission[['sample_id', 'target']]\n        \n        submission.to_csv(OUTPUT_CSV, index=False)\n        print(f\"Saved submission to {OUTPUT_CSV}\")\n    else:\n        print(\"Test CSV not found.\")\n\nif __name__ == '__main__':\n    main()\n","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}