{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":13836,"databundleVersionId":1718836,"sourceType":"competition"}],"dockerImageVersionId":30919,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        # Check if the file extension is not '.jpg'\n        if not filename.lower().endswith('.jpg'):\n            print(os.path.join(dirname, filename))","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-04-04T01:08:08.0211Z","iopub.execute_input":"2025-04-04T01:08:08.021383Z","iopub.status.idle":"2025-04-04T01:08:30.0175Z","shell.execute_reply.started":"2025-04-04T01:08:08.021353Z","shell.execute_reply":"2025-04-04T01:08:30.01683Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import sys\nimport os\nimport random\nimport json\nimport numpy as np\nimport pandas as pd\nfrom PIL import Image\nfrom tqdm import tqdm\nfrom pathlib import Path\nfrom typing import Optional, Tuple, Dict, List, Union\n\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.optim.lr_scheduler import _LRScheduler\nfrom torch.cuda.amp import autocast, GradScaler\n\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\nfrom sklearn.model_selection import StratifiedKFold\nfrom sklearn.metrics import accuracy_score\n\nimport timm\nimport math\nimport matplotlib.pyplot as plt\n\n# Define paths for data and outputs\nPATHS = {\n    'TRAIN_CSV': '/kaggle/input/cassava-leaf-disease-classification/train.csv',\n    'TEST_CSV': '/kaggle/input/cassava-leaf-disease-classification/sample_submission.csv',\n    'DISEASE_MAP': '/kaggle/input/cassava-leaf-disease-classification/label_num_to_disease_map.json',\n    'TRAIN_IMAGES': '/kaggle/input/cassava-leaf-disease-classification/train_images',\n    'TEST_IMAGES': '/kaggle/input/cassava-leaf-disease-classification/test_images',\n    'OUTPUT': '/kaggle/working/submission.csv',\n    'WEIGHTS': '/kaggle/working/weights',\n    'PLOTS': '/kaggle/working/plots'  # New path for saving plots\n}\n\n# Create weights and plots directories\nos.makedirs(PATHS['WEIGHTS'], exist_ok=True)\nos.makedirs(PATHS['PLOTS'], exist_ok=True)\n\n\n# Custom Scheduler with warmup and cosine decay\nclass WarmupCosineScheduler(_LRScheduler):\n    \"\"\"\n    Scheduler with warmup and cosine decay as specified in requirements\n    \"\"\"\n    def __init__(self, optimizer, warmup_epochs, max_epochs, min_lr=1e-6, max_lr=2e-4, final_lr=3.17e-6, last_epoch=-1):\n        self.warmup_epochs = warmup_epochs\n        self.max_epochs = max_epochs\n        self.min_lr = min_lr\n        self.max_lr = max_lr\n        self.final_lr = final_lr\n        super(WarmupCosineScheduler, self).__init__(optimizer, last_epoch)\n\n    def get_lr(self):\n        if self.last_epoch < self.warmup_epochs:\n            # Linear warmup\n            alpha = self.last_epoch / self.warmup_epochs\n            factor = alpha\n            # Ensure factor is positive\n            factor = max(0, factor)\n            return [self.min_lr + factor * (self.max_lr - self.min_lr) for _ in self.base_lrs]\n        else:\n            # Cosine decay from max_lr to final_lr\n            progress = (self.last_epoch - self.warmup_epochs) / (self.max_epochs - self.warmup_epochs)\n            progress = min(1.0, progress)\n            # Add epsilon to avoid numerical issues\n            cosine_factor = 0.5 * (1 + math.cos(math.pi * progress))\n            return [self.final_lr + cosine_factor * (self.max_lr - self.final_lr) for _ in self.base_lrs]\n\n\n# Sigmoid Focal Loss with Label Smoothing\nclass SigmoidFocalLossWithLabelSmoothing(nn.Module):\n    \"\"\"\n    Sigmoid Focal Loss with Label Smoothing as specified in requirements\n    \"\"\"\n    def __init__(self, gamma=2.0, alpha=0.25, smoothing=0.1, reduction='mean'):\n        super(SigmoidFocalLossWithLabelSmoothing, self).__init__()\n        self.gamma = gamma\n        self.alpha = alpha\n        self.smoothing = smoothing\n        self.reduction = reduction\n\n    def forward(self, logits, targets):\n        num_classes = logits.size(1)\n        \n        # Apply label smoothing\n        smoothed_targets = torch.zeros_like(logits)\n        smoothed_targets.fill_(self.smoothing / (num_classes - 1))\n        smoothed_targets.scatter_(1, targets.unsqueeze(1), 1.0 - self.smoothing)\n        \n        # Get probabilities\n        probs = torch.sigmoid(logits)\n        \n        # Calculate focal loss with label smoothing\n        targets_one_hot = torch.zeros_like(logits)\n        targets_one_hot.scatter_(1, targets.unsqueeze(1), 1)\n        \n        pt = (1 - probs) * targets_one_hot + probs * (1 - targets_one_hot)\n        focal_weight = (self.alpha * targets_one_hot + (1 - self.alpha) * (1 - targets_one_hot)) * pt.pow(self.gamma)\n        \n        loss = -focal_weight * (\n            smoothed_targets * torch.log(probs + 1e-8) + \n            (1 - smoothed_targets) * torch.log(1 - probs + 1e-8)\n        )\n        \n        if self.reduction == 'mean':\n            return loss.mean()\n        elif self.reduction == 'sum':\n            return loss.sum()\n        return loss\n\n\nclass Config:\n    \"\"\"Configuration class for model training and inference.\"\"\"\n    \n    def __init__(self):\n        # Model parameters\n        self.seed: int = 719\n        self.model_name: str = 'tf_efficientnet_b4.ns_jft_in1k'  # Updated model name to avoid deprecation warning\n        self.image_size: int = 512\n        self.drop_connect_rate: float = 0.4  # Custom drop connect rate\n        self.dropout_rate: float = 0.5      # Dropout rate for custom head\n        \n        # Training parameters - MODIFIED FOR DEMO RUN\n        self.num_epochs: int = 15           # Reduced to 15 epochs for demo run\n        self.train_batch_size: int = 16\n        self.valid_batch_size: int = 32\n        \n        # Learning rate parameters\n        self.min_lr: float = 1e-6\n        self.max_lr: float = 2e-4\n        self.final_lr: float = 3.17e-6\n        self.warmup_epochs: int = 3\n        \n        # Loss function parameters\n        self.focal_gamma: float = 2.0\n        self.focal_alpha: float = 0.25\n        self.label_smoothing: float = 0.1\n        \n        # Early stopping parameters\n        self.early_stopping_patience: int = 5\n        self.early_stopping_min_delta: float = 0.001\n        \n        # Other parameters\n        self.num_workers: int = 4\n        self.grad_accum_steps: int = 1\n        self.device: str = 'cuda' if torch.cuda.is_available() else 'cpu'\n        self.num_folds: int = 5\n        \n        # MODIFIED FOR DEMO RUN - Only use a single fold\n        self.used_epochs: list = list(range(15))  # Consider all 15 epochs\n        self.used_folds: list = [0]               # Only use fold 0 for testing\n        \n        # Normalization parameters (Global mean and std for 2020 Cassava dataset)\n        self.mean: list = [0.4253, 0.4559, 0.3395]  # Updated for Cassava dataset\n        self.std: list = [0.2236, 0.2261, 0.2339]   # Updated for Cassava dataset\n        self.weight_decay: float = 1e-4\n        \n        # TTA parameters\n        self.tta_patches: int = 6  # 4 overlapping + 2 center crops\n        self.tta_augmentations: int = 2  # 2 augmentations per patch\n        \n        # Checkpoint parameters\n        self.save_checkpoint_freq: int = 5  # Save checkpoint every N epochs\n        self.save_best_model: bool = True   # Save best model during training\n\n    @staticmethod\n    def load_disease_map(path: str) -> Dict[str, str]:\n        \"\"\"Load disease mapping from JSON file.\"\"\"\n        with open(path, 'r') as f:\n            return json.load(f)\n\n\ndef seed_everything(seed: int) -> None:\n    \"\"\"Set random seeds for reproducibility.\"\"\"\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = True\n\n\nclass CassavaDataset(Dataset):\n    \"\"\"Dataset class for Cassava Leaf Disease Classification.\"\"\"\n    \n    def __init__(\n        self,\n        df: pd.DataFrame,\n        data_root: str,\n        transforms: Optional[A.Compose] = None,\n        output_label: bool = True\n    ):\n        super().__init__()\n        self.df = df.reset_index(drop=True)\n        self.transforms = transforms\n        self.data_root = Path(data_root)\n        self.output_label = output_label\n    \n    def __len__(self) -> int:\n        return len(self.df)\n    \n    def __getitem__(self, index: int) -> Tuple[torch.Tensor, Optional[int]]:\n        image_id = self.df.iloc[index]['image_id']\n        if not str(image_id).lower().endswith(('.jpg', '.jpeg', '.png')):\n            image_id = f\"{image_id}.jpg\"\n        \n        image_path = self.data_root / image_id\n        try:\n            image = self._load_image(str(image_path))\n        except Exception as e:\n            print(f\"Error loading image {image_path}: {str(e)}\")\n            print(f\"Trying alternative extensions...\")\n            \n            for ext in ['.jpg', '.jpeg', '.png']:\n                try:\n                    alt_path = self.data_root / f\"{self.df.iloc[index]['image_id']}{ext}\"\n                    if alt_path.exists():\n                        image = self._load_image(str(alt_path))\n                        break\n                except:\n                    continue\n            else:\n                print(f\"Could not load image with any extension. Using blank image.\")\n                image = np.zeros((512, 512, 3), dtype=np.uint8)\n        \n        if self.transforms:\n            image = self.transforms(image=image)['image']\n        \n        if self.output_label:\n            target = self.df.iloc[index]['label']\n            return image, target\n        return image\n    \n    @staticmethod\n    def _load_image(path: str) -> np.ndarray:\n        \"\"\"Load and convert image to RGB.\"\"\"\n        image = np.array(Image.open(path).convert('RGB'))\n        if image is None:\n            raise ValueError(f\"Failed to load image at {path}\")\n        return image\n\n\nclass CustomEfficientNet(nn.Module):\n    \"\"\"\n    Custom EfficientNet with modified drop path rate and custom head\n    with global average pooling and dropout\n    \"\"\"\n    def __init__(\n        self,\n        model_name: str,\n        num_classes: int,\n        pretrained: bool = True,\n        drop_connect_rate: float = 0.4,  # This will be mapped to drop_path_rate\n        dropout_rate: float = 0.5\n    ):\n        super().__init__()\n        # Create base model - using drop_path_rate instead of drop_connect_rate\n        try:\n            self.model = timm.create_model(\n                model_name,\n                pretrained=pretrained,\n                drop_path_rate=drop_connect_rate  # This is the key change\n            )\n            \n            # Get the number of features from the last layer\n            if hasattr(self.model, 'classifier'):\n                n_features = self.model.classifier.in_features\n                self.model.classifier = nn.Identity()\n            elif hasattr(self.model, 'fc'):\n                n_features = self.model.fc.in_features\n                self.model.fc = nn.Identity()\n            else:\n                raise AttributeError(\"Model doesn't have a standard classifier or fc layer\")\n            \n            # Create a custom head with global average pooling and dropout\n            self.custom_head = nn.Sequential(\n                nn.Dropout(dropout_rate),\n                nn.Linear(n_features, num_classes)\n            )\n            \n        except Exception as e:\n            print(f\"Error creating model: {e}\")\n            print(f\"Trying fallback to a different model architecture...\")\n            # Fallback to a different EfficientNet version\n            self.model = timm.create_model(\n                'tf_efficientnet_b3.ns_jft_in1k',  # Fallback model\n                pretrained=True,\n                drop_path_rate=drop_connect_rate\n            )\n            n_features = self.model.classifier.in_features\n            self.model.classifier = nn.Identity()\n            \n            self.custom_head = nn.Sequential(\n                nn.Dropout(dropout_rate),\n                nn.Linear(n_features, num_classes)\n            )\n    \n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        # Forward pass through the base model\n        features = self.model(x)\n        # Forward pass through the custom head\n        return self.custom_head(features)\n\n\nclass DataTransforms:\n    \"\"\"Data augmentation and preprocessing transforms.\"\"\"\n    \n    @staticmethod\n    def get_train_transforms(config: Config) -> A.Compose:\n        \"\"\"Return training data augmentation transforms.\"\"\"\n        return A.Compose([\n            A.RandomResizedCrop(config.image_size, config.image_size),\n            A.Transpose(p=0.5),\n            A.HorizontalFlip(p=0.5),\n            A.VerticalFlip(p=0.5),\n            A.ShiftScaleRotate(p=0.5),\n            A.HueSaturationValue(\n                hue_shift_limit=0.2,\n                sat_shift_limit=0.2,\n                val_shift_limit=0.2,\n                p=0.5\n            ),\n            A.RandomBrightnessContrast(\n                brightness_limit=(-0.1, 0.1),\n                contrast_limit=(-0.1, 0.1),\n                p=0.5\n            ),\n            A.Normalize(\n                mean=config.mean,\n                std=config.std,\n                max_pixel_value=255.0,\n                p=1.0\n            ),\n            A.CoarseDropout(\n                max_holes=8,\n                max_height=config.image_size // 8,\n                max_width=config.image_size // 8,\n                min_holes=5,\n                min_height=config.image_size // 16,\n                min_width=config.image_size // 16,\n                fill_value=0,\n                p=0.5\n            ),\n            ToTensorV2(p=1.0),\n        ], p=1.)\n    \n    @staticmethod\n    def get_valid_transforms(config: Config) -> A.Compose:\n        \"\"\"Return validation data preprocessing transforms.\"\"\"\n        return A.Compose([\n            A.CenterCrop(config.image_size, config.image_size, p=1.),\n            A.Resize(config.image_size, config.image_size),\n            A.Normalize(\n                mean=config.mean,\n                std=config.std,\n                max_pixel_value=255.0,\n                p=1.0\n            ),\n            ToTensorV2(p=1.0),\n        ], p=1.)\n    \n    @staticmethod\n    def get_specific_tta_transforms(config: Config, aug_type: int) -> A.Compose:\n        \"\"\"Return specific TTA transforms based on augmentation type.\"\"\"\n        if aug_type == 0:  # No augmentation (center crop)\n            return A.Compose([\n                A.CenterCrop(config.image_size, config.image_size),\n                A.Normalize(mean=config.mean, std=config.std, max_pixel_value=255.0),\n                ToTensorV2(),\n            ])\n        elif aug_type == 1:  # Horizontal flip\n            return A.Compose([\n                A.CenterCrop(config.image_size, config.image_size),\n                A.HorizontalFlip(p=1.0),\n                A.Normalize(mean=config.mean, std=config.std, max_pixel_value=255.0),\n                ToTensorV2(),\n            ])\n        elif aug_type == 2:  # Vertical flip\n            return A.Compose([\n                A.CenterCrop(config.image_size, config.image_size),\n                A.VerticalFlip(p=1.0),\n                A.Normalize(mean=config.mean, std=config.std, max_pixel_value=255.0),\n                ToTensorV2(),\n            ])\n        elif aug_type == 3:  # Transpose\n            return A.Compose([\n                A.CenterCrop(config.image_size, config.image_size),\n                A.Transpose(p=1.0),\n                A.Normalize(mean=config.mean, std=config.std, max_pixel_value=255.0),\n                ToTensorV2(),\n            ])\n        elif aug_type == 4:  # Rotate 90\n            return A.Compose([\n                A.CenterCrop(config.image_size, config.image_size),\n                A.Rotate(limit=(90, 90), p=1.0),\n                A.Normalize(mean=config.mean, std=config.std, max_pixel_value=255.0),\n                ToTensorV2(),\n            ])\n        else:  # Rotate 270\n            return A.Compose([\n                A.CenterCrop(config.image_size, config.image_size),\n                A.Rotate(limit=(270, 270), p=1.0),\n                A.Normalize(mean=config.mean, std=config.std, max_pixel_value=255.0),\n                ToTensorV2(),\n            ])\n\n\nclass EarlyStopping:\n    \"\"\"Early stopping implementation with model weights restoration.\"\"\"\n    \n    def __init__(\n        self,\n        patience: int = 5,\n        min_delta: float = 0.001,\n        restore_best_weights: bool = True\n    ):\n        self.patience = patience\n        self.min_delta = min_delta\n        self.restore_best_weights = restore_best_weights\n        self.best_score = None\n        self.counter = 0\n        self.early_stop = False\n        self.best_weights = None\n        self.best_epoch = -1\n    \n    def __call__(self, model: nn.Module, val_score: float, epoch: int) -> bool:\n        \"\"\"Update early stopping state based on validation score.\"\"\"\n        if self.best_score is None:\n            self.best_score = val_score\n            self.best_weights = model.state_dict().copy()\n            self.best_epoch = epoch\n            return False\n        \n        if val_score > self.best_score + self.min_delta:\n            self.best_score = val_score\n            self.counter = 0\n            self.best_weights = model.state_dict().copy()\n            self.best_epoch = epoch\n        else:\n            self.counter += 1\n            if self.counter >= self.patience:\n                self.early_stop = True\n                return True\n        \n        return False\n    \n    def restore_model(self, model: nn.Module) -> None:\n        \"\"\"Restore model to best weights.\"\"\"\n        if self.restore_best_weights and self.best_weights is not None:\n            model.load_state_dict(self.best_weights)\n            print(f\"Restored model to best weights from epoch {self.best_epoch}.\")\n\n\nclass Trainer:\n    \"\"\"Class for model training with warmup + cosine decay and early stopping.\"\"\"\n    \n    def __init__(self, model: nn.Module, config: Config):\n        self.model = model\n        self.config = config\n        self.criterion = SigmoidFocalLossWithLabelSmoothing(\n            gamma=config.focal_gamma,\n            alpha=config.focal_alpha,\n            smoothing=config.label_smoothing\n        )\n        self.scaler = GradScaler()\n        self.device = torch.device(config.device)\n        \n        # Initialize optimizer\n        self.optimizer = optim.AdamW(\n            self.model.parameters(),\n            lr=config.min_lr,\n            weight_decay=config.weight_decay\n        )\n        \n        # Initialize scheduler with warmup and cosine decay\n        self.scheduler = WarmupCosineScheduler(\n            self.optimizer,\n            warmup_epochs=config.warmup_epochs,\n            max_epochs=config.num_epochs,\n            min_lr=config.min_lr,\n            max_lr=config.max_lr,\n            final_lr=config.final_lr\n        )\n        \n        # Initialize early stopping\n        self.early_stopping = EarlyStopping(\n            patience=config.early_stopping_patience,\n            min_delta=config.early_stopping_min_delta\n        )\n    \n    def train_epoch(self, train_loader: DataLoader) -> Tuple[float, float]:\n        \"\"\"Train the model for one epoch.\"\"\"\n        self.model.train()\n        total_loss = 0.0\n        all_predictions = []\n        all_targets = []\n        batch_count = 0\n        \n        # Initialize optimizer gradients at the beginning\n        self.optimizer.zero_grad()\n        \n        pbar = tqdm(train_loader, desc='Training')\n        \n        for batch_idx, (images, batch_targets) in enumerate(pbar):\n            batch_count += 1\n            images = images.to(self.device)\n            batch_targets = batch_targets.to(self.device)\n            \n            # Forward pass with automatic mixed precision\n            with autocast():\n                outputs = self.model(images)\n                loss = self.criterion(outputs, batch_targets)\n            \n            # Scale the loss and perform backward pass\n            self.scaler.scale(loss).backward()\n            \n            # Gradient accumulation\n            if (batch_idx + 1) % self.config.grad_accum_steps == 0:\n                self.scaler.step(self.optimizer)\n                self.scaler.update()\n                self.optimizer.zero_grad()\n            \n            # Record loss and predictions\n            total_loss += loss.item()\n            \n            # Get predicted class\n            _, predictions = torch.max(outputs, dim=1)\n            all_predictions.extend(predictions.cpu().numpy())\n            all_targets.extend(batch_targets.cpu().numpy())\n            \n            # Update progress bar\n            pbar.set_postfix({\n                'loss': f'{loss.item():.4f}',\n                'lr': f'{self.optimizer.param_groups[0][\"lr\"]:.6f}'\n            })\n        \n        # Make sure to update optimizer for any remaining gradients\n        if batch_count % self.config.grad_accum_steps != 0:\n            self.scaler.step(self.optimizer)\n            self.scaler.update()\n            self.optimizer.zero_grad()\n        \n        # Compute metrics\n        accuracy = accuracy_score(all_targets, all_predictions)\n        average_loss = total_loss / batch_count\n        \n        return average_loss, accuracy\n    \n    @torch.no_grad()\n    def validate(self, valid_loader: DataLoader) -> Tuple[float, float]:\n        \"\"\"Evaluate the model on the validation set.\"\"\"\n        self.model.eval()\n        total_loss = 0.0\n        all_predictions = []\n        all_targets = []\n        batch_count = 0\n        \n        pbar = tqdm(valid_loader, desc='Validating')\n        \n        for images, batch_targets in pbar:\n            batch_count += 1\n            images = images.to(self.device)\n            batch_targets = batch_targets.to(self.device)\n            \n            # Forward pass\n            outputs = self.model(images)\n            loss = self.criterion(outputs, batch_targets)\n            \n            # Record loss and predictions\n            total_loss += loss.item()\n            \n            # Get predicted class\n            _, predictions = torch.max(outputs, dim=1)\n            all_predictions.extend(predictions.cpu().numpy())\n            all_targets.extend(batch_targets.cpu().numpy())\n            \n            # Update progress bar\n            pbar.set_postfix({'loss': f'{loss.item():.4f}'})\n        \n        # Compute metrics\n        accuracy = accuracy_score(all_targets, all_predictions)\n        average_loss = total_loss / batch_count\n        \n        return average_loss, accuracy\n\n\ndef plot_training_curves(history: Dict, fold: int) -> None:\n    \"\"\"Plot training and validation loss and accuracy curves.\"\"\"\n    # Create figure with 2 subplots\n    plt.figure(figsize=(12, 10))\n    \n    # Plot training and validation loss\n    plt.subplot(2, 1, 1)\n    plt.plot(history['train_loss'], label='Training Loss')\n    plt.plot(history['val_loss'], label='Validation Loss')\n    plt.title(f'Loss Curves - Fold {fold}')\n    plt.xlabel('Epoch')\n    plt.ylabel('Loss')\n    plt.legend()\n    plt.grid(True)\n    \n    # Plot training and validation accuracy\n    plt.subplot(2, 1, 2)\n    plt.plot(history['train_acc'], label='Training Accuracy')\n    plt.plot(history['val_acc'], label='Validation Accuracy')\n    plt.title(f'Accuracy Curves - Fold {fold}')\n    plt.xlabel('Epoch')\n    plt.ylabel('Accuracy')\n    plt.legend()\n    plt.grid(True)\n    \n    # Adjust layout and save the figure\n    plt.tight_layout()\n    plot_path = os.path.join(PATHS['PLOTS'], f'training_curves_fold_{fold}.png')\n    plt.savefig(plot_path)\n    plt.close()\n    print(f\"Saved training curves plot to {plot_path}\")\n    \n    # Plot learning rate curve\n    plt.figure(figsize=(10, 5))\n    # Extract learning rates from history if available\n    if 'learning_rates' in history:\n        plt.plot(history['learning_rates'])\n        plt.title(f'Learning Rate Schedule - Fold {fold}')\n        plt.xlabel('Epoch')\n        plt.ylabel('Learning Rate')\n        plt.grid(True)\n        \n        lr_plot_path = os.path.join(PATHS['PLOTS'], f'lr_curve_fold_{fold}.png')\n        plt.savefig(lr_plot_path)\n        plt.close()\n        print(f\"Saved learning rate curve plot to {lr_plot_path}\")\n\n\ndef train_fold(fold: int, train_data: pd.DataFrame, valid_data: pd.DataFrame, config: Config, data_root: str) -> Dict:\n    \"\"\"Train a single fold with early stopping and weight restoration.\"\"\"\n    device = torch.device(config.device)\n    \n    # Create datasets\n    train_dataset = CassavaDataset(\n        train_data,\n        data_root=os.path.join(data_root, PATHS['TRAIN_IMAGES']),\n        transforms=DataTransforms.get_train_transforms(config)\n    )\n    \n    valid_dataset = CassavaDataset(\n        valid_data,\n        data_root=os.path.join(data_root, PATHS['TRAIN_IMAGES']),\n        transforms=DataTransforms.get_valid_transforms(config)\n    )\n    \n    # Create DataLoaders\n    train_loader = DataLoader(\n        train_dataset,\n        batch_size=config.train_batch_size,\n        shuffle=True,\n        num_workers=config.num_workers,\n        pin_memory=True\n    )\n    \n    valid_loader = DataLoader(\n        valid_dataset,\n        batch_size=config.valid_batch_size,\n        shuffle=False,\n        num_workers=config.num_workers,\n        pin_memory=True\n    )\n    \n    # Initialize the model with custom drop connect rate and custom head\n    model = CustomEfficientNet(\n        config.model_name,\n        train_data.label.nunique(),\n        pretrained=True,\n        drop_connect_rate=config.drop_connect_rate,\n        dropout_rate=config.dropout_rate\n    ).to(device)\n    \n    # Initialize the trainer\n    trainer = Trainer(model, config)\n    \n    best_val_acc = 0\n    best_val_loss = float('inf')\n    model_states = {}\n    history = {\n        'train_loss': [], 'train_acc': [], 'val_loss': [], 'val_acc': [], 'learning_rates': []\n    }\n    \n    # Training loop\n    for epoch in range(config.num_epochs):\n        print(f'Epoch {epoch + 1}/{config.num_epochs}')\n        \n        # Train for one epoch\n        train_loss, train_acc = trainer.train_epoch(train_loader)\n        \n        # Validate the model\n        val_loss, val_acc = trainer.validate(valid_loader)\n        \n        # Get current learning rate\n        current_lr = trainer.optimizer.param_groups[0]['lr']\n        \n        # Update scheduler\n        trainer.scheduler.step()\n        \n        # Update history\n        history['train_loss'].append(train_loss)\n        history['train_acc'].append(train_acc)\n        history['val_loss'].append(val_loss)\n        history['val_acc'].append(val_acc)\n        history['learning_rates'].append(current_lr)\n        \n        # Print metrics\n        print(f'Train Loss: {train_loss:.4f}, Accuracy: {train_acc:.4f}')\n        print(f'Valid Loss: {val_loss:.4f}, Accuracy: {val_acc:.4f}')\n        print(f'Learning Rate: {current_lr:.6f}')\n        \n        # Save checkpoint at specified frequency\n        if (epoch + 1) % config.save_checkpoint_freq == 0 or epoch == config.num_epochs - 1:\n            checkpoint_path = os.path.join(\n                PATHS['WEIGHTS'],\n                f'{config.model_name}_fold_{fold}_epoch_{epoch}_checkpoint.pth'\n            )\n            torch.save({\n                'epoch': epoch,\n                'model_state_dict': model.state_dict(),\n                'optimizer_state_dict': trainer.optimizer.state_dict(),\n                'scheduler_state_dict': trainer.scheduler.state_dict(),\n                'train_loss': train_loss,\n                'val_loss': val_loss,\n                'train_acc': train_acc,\n                'val_acc': val_acc,\n                'config': {k: v for k, v in config.__dict__.items() if not k.startswith('__')}\n            }, checkpoint_path)\n            print(f'Saved checkpoint: {checkpoint_path}')\n        \n        # Update best model\n        if val_acc > best_val_acc:\n            best_val_acc = val_acc\n            best_val_loss = val_loss\n            \n            # Save best model weights\n            if config.save_best_model:\n                best_model_path = os.path.join(\n                    PATHS['WEIGHTS'],\n                    f'{config.model_name}_fold_{fold}_best.pth'\n                )\n                torch.save(model.state_dict(), best_model_path)\n                print(f'Saved best model weights: {best_model_path}')\n            \n            # Save checkpoint for weights loading during inference\n            checkpoint_path = os.path.join(\n                PATHS['WEIGHTS'],\n                f'{config.model_name}_fold_{fold}_{epoch}'\n            )\n            torch.save(model.state_dict(), checkpoint_path)\n            print(f'Saved checkpoint for inference: {checkpoint_path}')\n        \n        # Check early stopping\n        if trainer.early_stopping(model, val_acc, epoch):\n            print(f'Early stopping triggered at epoch {epoch + 1}')\n            # Restore best weights\n            trainer.early_stopping.restore_model(model)\n            break\n        \n        # Save model state for specified epochs\n        if epoch in config.used_epochs:\n            model_states[epoch] = model.state_dict().copy()\n            print(f'Saved model state for epoch {epoch}')\n    \n    # Plot training curves after training is complete\n    plot_training_curves(history, fold)\n    \n    # Save final model\n    final_model_path = os.path.join(\n        PATHS['WEIGHTS'],\n        f'{config.model_name}_fold_{fold}_final.pth'\n    )\n    torch.save(model.state_dict(), final_model_path)\n    print(f'Saved final model weights: {final_model_path}')\n    \n    # Clean up\n    torch.cuda.empty_cache()\n    \n    return {\n        'best_val_acc': best_val_acc,\n        'best_val_loss': best_val_loss,\n        'history': history,\n        'model_states': model_states,\n        'best_epoch': trainer.early_stopping.best_epoch\n    }\n\n\n@torch.no_grad()\ndef process_image_for_tta(\n    model: nn.Module,\n    image: np.ndarray,\n    config: Config,\n    device: torch.device\n) -> torch.Tensor:\n    \"\"\"\n    Process a single image for Test Time Augmentation using the specified patch and augmentation strategy.\n    \n    1. Extract 4 overlapping patches from corners\n    2. Extract 2 center crops\n    3. Apply specified augmentations to each patch\n    4. Average predictions across all patches and augmentations\n    \"\"\"\n    h, w = image.shape[:2]\n    patch_size = config.image_size\n    \n    # Handle images smaller than patch_size\n    if h < patch_size or w < patch_size:\n        # Resize image to be at least patch_size x patch_size\n        new_h = max(patch_size, h)\n        new_w = max(patch_size, w)\n        resized_image = np.array(Image.fromarray(image).resize((new_w, new_h)))\n        image = resized_image\n        h, w = image.shape[:2]\n    \n    # Extract 4 overlapping patches from corners\n    patches = [\n        image[:patch_size, :patch_size],                      # Top-left\n        image[:patch_size, w-patch_size:],                    # Top-right\n        image[h-patch_size:, :patch_size],                    # Bottom-left\n        image[h-patch_size:, w-patch_size:],                  # Bottom-right\n        image[h//2-patch_size//2:h//2+patch_size//2,          # Center\n              w//2-patch_size//2:w//2+patch_size//2],\n        image[h//2-patch_size//2:h//2+patch_size//2,          # Center (duplicate)\n              w//2-patch_size//2:w//2+patch_size//2]\n    ]\n    \n    # Apply different augmentations to each patch\n    predictions = []\n    \n    # Process each patch\n    for i, patch in enumerate(patches):\n        # For the first 4 patches, apply 2 different augmentations\n        if i < 4:\n            for aug_type in range(1, 3):  # Augmentation types 1 and 2\n                transform = DataTransforms.get_specific_tta_transforms(config, aug_type)\n                transformed = transform(image=patch)['image']\n                \n                # Get prediction\n                pred = model(transformed.unsqueeze(0).to(device))\n                predictions.append(pred.softmax(dim=1))\n        else:\n            # For center crops, use no augmentation\n            transform = DataTransforms.get_specific_tta_transforms(config, 0)\n            transformed = transform(image=patch)['image']\n            \n            # Get prediction\n            pred = model(transformed.unsqueeze(0).to(device))\n            predictions.append(pred.softmax(dim=1))\n    \n    # Average all predictions\n    final_pred = torch.mean(torch.cat(predictions, dim=0), dim=0, keepdim=True)\n    return final_pred\n    \ndef main():\n    \"\"\"\n    Main function to run the Cassava Leaf Disease Classification training pipeline.\n    This function:\n    1. Initializes configuration\n    2. Sets random seeds\n    3. Loads and prepares data\n    4. Runs k-fold cross-validation training\n    5. Makes predictions on test set\n    6. Creates submission file\n    \"\"\"\n    # Initialize configuration\n    config = Config()\n    \n    # Set seeds for reproducibility\n    seed_everything(config.seed)\n    \n    # Load disease mapping\n    disease_map = Config.load_disease_map(PATHS['DISEASE_MAP'])\n    print(f\"Disease mapping: {disease_map}\")\n    \n    # Read training data\n    train_df = pd.read_csv(PATHS['TRAIN_CSV'])\n    print(f\"Training data shape: {train_df.shape}\")\n    print(f\"Class distribution:\\n{train_df['label'].value_counts()}\")\n    \n    # Setup k-fold cross-validation\n    skf = StratifiedKFold(\n        n_splits=config.num_folds,\n        shuffle=True,\n        random_state=config.seed\n    )\n    \n    # Dictionary to store results for each fold\n    fold_results = {}\n    \n    # Run k-fold cross-validation\n    for fold, (train_idx, val_idx) in enumerate(skf.split(train_df, train_df['label'])):\n        print(f\"\\n{'='*50}\")\n        print(f\"Fold {fold}\")\n        print(f\"{'='*50}\")\n        \n        # Check if we should process this fold\n        if fold not in config.used_folds:\n            print(f\"Skipping fold {fold} as per configuration.\")\n            continue\n        \n        # Split data for this fold\n        train_data = train_df.iloc[train_idx].reset_index(drop=True)\n        valid_data = train_df.iloc[val_idx].reset_index(drop=True)\n        \n        print(f\"Training on {len(train_data)} samples, validating on {len(valid_data)} samples\")\n        \n        # Train the model for this fold\n        fold_result = train_fold(fold, train_data, valid_data, config, '')\n        \n        # Store results\n        fold_results[fold] = fold_result\n        \n        print(f\"\\nFold {fold} best validation accuracy: {fold_result['best_val_acc']:.4f}\")\n        print(f\"Fold {fold} best validation loss: {fold_result['best_val_loss']:.4f}\")\n        print(f\"Fold {fold} best epoch: {fold_result['best_epoch']}\")\n    \n    # Print overall results\n    print(\"\\n==== Overall Results ====\")\n    if fold_results:\n        # Calculate average metrics across folds\n        avg_val_acc = sum(res['best_val_acc'] for res in fold_results.values()) / len(fold_results)\n        avg_val_loss = sum(res['best_val_loss'] for res in fold_results.values()) / len(fold_results)\n        \n        print(f\"Average best validation accuracy: {avg_val_acc:.4f}\")\n        print(f\"Average best validation loss: {avg_val_loss:.4f}\")\n        \n        # Optionally perform inference on test set\n        test_inference(config, fold_results)\n    else:\n        print(\"No folds were processed.\")\n\ndef test_inference(config, fold_results):\n    \"\"\"\n    Perform inference on the test set using trained models.\n    \n    Args:\n        config: Configuration object\n        fold_results: Dictionary of results from each fold\n    \"\"\"\n    print(\"\\n==== Test Inference ====\")\n    \n    # Read test data\n    test_df = pd.read_csv(PATHS['TEST_CSV'])\n    print(f\"Test data shape: {test_df.shape}\")\n    \n    # Initialize device\n    device = torch.device(config.device)\n    \n    # Create test dataset\n    test_dataset = CassavaDataset(\n        test_df,\n        data_root=PATHS['TEST_IMAGES'],\n        transforms=DataTransforms.get_valid_transforms(config),\n        output_label=False\n    )\n    \n    # Create test dataloader\n    test_loader = DataLoader(\n        test_dataset,\n        batch_size=config.valid_batch_size,\n        shuffle=False,\n        num_workers=config.num_workers,\n        pin_memory=True\n    )\n    \n    # Ensemble predictions from multiple folds\n    all_predictions = []\n    \n    # For each fold\n    for fold in fold_results.keys():\n        print(f\"Inferencing with fold {fold} model...\")\n        \n        # Load the best model for this fold\n        model = CustomEfficientNet(\n            config.model_name,\n            len(Config.load_disease_map(PATHS['DISEASE_MAP'])),\n            pretrained=False,\n            drop_connect_rate=config.drop_connect_rate,\n            dropout_rate=config.dropout_rate\n        ).to(device)\n        \n        best_model_path = os.path.join(\n            PATHS['WEIGHTS'],\n            f'{config.model_name}_fold_{fold}_best.pth'\n        )\n        \n        # Check if best model exists\n        if os.path.exists(best_model_path):\n            model.load_state_dict(torch.load(best_model_path))\n            print(f\"Loaded best model from {best_model_path}\")\n        else:\n            # Try to load from the used_epochs list\n            best_epoch = fold_results[fold]['best_epoch']\n            checkpoint_path = os.path.join(\n                PATHS['WEIGHTS'],\n                f'{config.model_name}_fold_{fold}_{best_epoch}'\n            )\n            \n            if os.path.exists(checkpoint_path):\n                model.load_state_dict(torch.load(checkpoint_path))\n                print(f\"Loaded checkpoint from {checkpoint_path}\")\n            else:\n                print(f\"Warning: Could not find best model for fold {fold}, using final model\")\n                final_path = os.path.join(\n                    PATHS['WEIGHTS'],\n                    f'{config.model_name}_fold_{fold}_final.pth'\n                )\n                model.load_state_dict(torch.load(final_path))\n        \n        # Set model to evaluation mode\n        model.eval()\n        \n        # Get predictions for this fold\n        fold_preds = []\n        \n        with torch.no_grad():\n            for images in tqdm(test_loader, desc=f\"Predicting fold {fold}\"):\n                images = images.to(device)\n                outputs = model(images)\n                probs = outputs.softmax(dim=1)\n                fold_preds.extend(probs.cpu().numpy())\n        \n        all_predictions.append(np.array(fold_preds))\n    \n    # Average predictions across folds\n    if all_predictions:\n        final_predictions = np.mean(all_predictions, axis=0)\n        \n        # Get the most likely class for each sample\n        predicted_classes = np.argmax(final_predictions, axis=1)\n        \n        # Create submission dataframe\n        submission_df = pd.DataFrame({\n            'image_id': test_df['image_id'],\n            'label': predicted_classes\n        })\n        \n        # Save submission file\n        submission_df.to_csv(PATHS['OUTPUT'], index=False)\n        print(f\"Saved submission file to {PATHS['OUTPUT']}\")\n    else:\n        print(\"No predictions were made.\")\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-06T23:23:22.537685Z","iopub.execute_input":"2025-04-06T23:23:22.53796Z"}},"outputs":[],"execution_count":null}]}