{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","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"},{"sourceId":13701079,"sourceType":"datasetVersion","datasetId":8715435}],"dockerImageVersionId":31193,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import warnings\nwarnings.filterwarnings('ignore')\n\nimport os\nimport pandas as pd\nimport numpy as np\nfrom sklearn.preprocessing import LabelEncoder\nfrom sklearn.model_selection import StratifiedKFold\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nimport torchvision.transforms as transforms\nfrom torchvision.models import efficientnet_b4\nimport cv2\nfrom tqdm import tqdm\nimport timm\n\n# Configuration\nclass Config:\n    image_size = (500, 250)  # width, height\n    batch_size = 16\n    num_epochs = 15\n    learning_rate = 1e-4\n    num_folds = 5\n    num_workers = 4\n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    \nconfig = Config()\n\n# Data Preprocessing\nclass CSIRODataset(Dataset):\n    def __init__(self, df, image_dir, transform=None, is_test=False):\n        self.df = df\n        self.image_dir = image_dir\n        self.transform = transform\n        self.is_test = is_test\n        \n    def __len__(self):\n        return len(self.df)\n    \n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n        image_path = os.path.join(self.image_dir, row['image_path'])\n        \n        # Load image\n        image = cv2.imread(image_path)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        \n        if self.transform:\n            image = self.transform(image)\n            \n        if self.is_test:\n            return image, row['image_path']\n        else:\n            # Multi-target output\n            targets = {\n                'state': row['state_encoded'],\n                'species': row['species_encoded'],\n                'ndvi': row['Pre_GSHH_NDVI'],\n                'height': row['Height_Ave_cm'],\n                'clover': row['Dry_Clover_g'],\n                'dead': row['Dry_Dead_g'],\n                'green': row['Dry_Green_g']\n            }\n            return image, targets\n\n# Multi-target Model\nclass CSIROModel(nn.Module):\n    def __init__(self, num_states, num_species):\n        super(CSIROModel, self).__init__()\n        \n        # Backbone\n        self.backbone = efficientnet_b4(pretrained=True)\n        in_features = self.backbone.classifier[1].in_features\n        self.backbone.classifier = nn.Identity()\n        \n        # Multi-task heads\n        self.state_head = nn.Linear(in_features, num_states)\n        self.species_head = nn.Linear(in_features, num_species)\n        \n        # Regression heads\n        self.ndvi_head = nn.Sequential(\n            nn.Linear(in_features, 512),\n            nn.ReLU(),\n            nn.Dropout(0.3),\n            nn.Linear(512, 1)\n        )\n        \n        self.height_head = nn.Sequential(\n            nn.Linear(in_features, 512),\n            nn.ReLU(),\n            nn.Dropout(0.3),\n            nn.Linear(512, 1)\n        )\n        \n        self.clover_head = nn.Sequential(\n            nn.Linear(in_features, 512),\n            nn.ReLU(),\n            nn.Dropout(0.3),\n            nn.Linear(512, 1)\n        )\n        \n        self.dead_head = nn.Sequential(\n            nn.Linear(in_features, 512),\n            nn.ReLU(),\n            nn.Dropout(0.3),\n            nn.Linear(512, 1)\n        )\n        \n        self.green_head = nn.Sequential(\n            nn.Linear(in_features, 512),\n            nn.ReLU(),\n            nn.Dropout(0.3),\n            nn.Linear(512, 1)\n        )\n        \n    def forward(self, x):\n        features = self.backbone(x)\n        \n        outputs = {\n            'state': self.state_head(features),\n            'species': self.species_head(features),\n            'ndvi': self.ndvi_head(features).squeeze(-1),\n            'height': self.height_head(features).squeeze(-1),\n            'clover': self.clover_head(features).squeeze(-1),\n            'dead': self.dead_head(features).squeeze(-1),\n            'green': self.green_head(features).squeeze(-1)\n        }\n        \n        return outputs\n\n# Loss function\nclass MultiTaskLoss(nn.Module):\n    def __init__(self, weight_state=1.0, weight_species=1.0, weight_regression=1.0):\n        super(MultiTaskLoss, self).__init__()\n        self.ce_loss = nn.CrossEntropyLoss()\n        self.mse_loss = nn.MSELoss()\n        self.weights = {\n            'state': weight_state,\n            'species': weight_species,\n            'ndvi': weight_regression,\n            'height': weight_regression,\n            'clover': weight_regression,\n            'dead': weight_regression,\n            'green': weight_regression\n        }\n    \n    def forward(self, outputs, targets):\n        total_loss = 0\n        \n        # Classification losses\n        state_loss = self.ce_loss(outputs['state'], targets['state'])\n        species_loss = self.ce_loss(outputs['species'], targets['species'])\n        \n        # Regression losses\n        ndvi_loss = self.mse_loss(outputs['ndvi'], targets['ndvi'])\n        height_loss = self.mse_loss(outputs['height'], targets['height'])\n        clover_loss = self.mse_loss(outputs['clover'], targets['clover'])\n        dead_loss = self.mse_loss(outputs['dead'], targets['dead'])\n        green_loss = self.mse_loss(outputs['green'], targets['green'])\n        \n        # Weighted sum\n        total_loss = (self.weights['state'] * state_loss +\n                     self.weights['species'] * species_loss +\n                     self.weights['ndvi'] * ndvi_loss +\n                     self.weights['height'] * height_loss +\n                     self.weights['clover'] * clover_loss +\n                     self.weights['dead'] * dead_loss +\n                     self.weights['green'] * green_loss)\n        \n        losses = {\n            'total': total_loss,\n            'state': state_loss,\n            'species': species_loss,\n            'ndvi': ndvi_loss,\n            'height': height_loss,\n            'clover': clover_loss,\n            'dead': dead_loss,\n            'green': green_loss\n        }\n        \n        return total_loss, losses\n\n# Training function\ndef train_epoch(model, dataloader, optimizer, criterion, device):\n    model.train()\n    running_loss = 0.0\n    \n    for batch_idx, (images, targets) in enumerate(tqdm(dataloader)):\n        images = images.to(device)\n        \n        # Move targets to device\n        batch_targets = {}\n        for key in targets:\n            if key in ['state', 'species']:\n                batch_targets[key] = targets[key].to(device)\n            else:\n                batch_targets[key] = targets[key].float().to(device)\n        \n        optimizer.zero_grad()\n        outputs = model(images)\n        loss, loss_dict = criterion(outputs, batch_targets)\n        loss.backward()\n        optimizer.step()\n        \n        running_loss += loss.item()\n    \n    return running_loss / len(dataloader)\n\n# Validation function\ndef validate_epoch(model, dataloader, criterion, device):\n    model.eval()\n    running_loss = 0.0\n    \n    with torch.no_grad():\n        for images, targets in dataloader:\n            images = images.to(device)\n            \n            # Move targets to device\n            batch_targets = {}\n            for key in targets:\n                if key in ['state', 'species']:\n                    batch_targets[key] = targets[key].to(device)\n                else:\n                    batch_targets[key] = targets[key].float().to(device)\n            \n            outputs = model(images)\n            loss, loss_dict = criterion(outputs, batch_targets)\n            running_loss += loss.item()\n    \n    return running_loss / len(dataloader)\n\n# Main training code\ndef main():\n    # Load data\n    train_df = pd.read_csv('/kaggle/input/csiro-stratified-train/state_train_pivoted.csv')\n    \n    # Encoding categorical variables\n    state_encoder = LabelEncoder()\n    species_encoder = LabelEncoder()\n    \n    train_df['state_encoded'] = state_encoder.fit_transform(train_df['State'])\n    train_df['species_encoded'] = species_encoder.fit_transform(train_df['Species'])\n    \n    # Image transforms\n    train_transform = transforms.Compose([\n        transforms.ToPILImage(),\n        transforms.Resize(config.image_size[::-1]),  # (height, width)\n        transforms.RandomHorizontalFlip(p=0.5),\n        transforms.RandomVerticalFlip(p=0.5),\n        transforms.RandomRotation(10),\n        transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),\n        transforms.ToTensor(),\n        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n    ])\n    \n    val_transform = transforms.Compose([\n        transforms.ToPILImage(),\n        transforms.Resize(config.image_size[::-1]),\n        transforms.ToTensor(),\n        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n    ])\n    \n    # Stratified K-Fold\n    skf = StratifiedKFold(n_splits=config.num_folds, shuffle=True, random_state=42)\n    \n    # Training loop for each fold\n    for fold, (train_idx, val_idx) in enumerate(skf.split(train_df, train_df['State'])):\n        print(f\"\\n=== Fold {fold + 1}/{config.num_folds} ===\")\n        \n        train_fold_df = train_df.iloc[train_idx].reset_index(drop=True)\n        val_fold_df = train_df.iloc[val_idx].reset_index(drop=True)\n        \n        # Datasets\n        train_dataset = CSIRODataset(train_fold_df, '/kaggle/input/csiro-biomass/', train_transform)\n        val_dataset = CSIRODataset(val_fold_df, '/kaggle/input/csiro-biomass/', val_transform)\n        \n        # DataLoaders\n        train_loader = DataLoader(train_dataset, batch_size=config.batch_size, \n                                shuffle=True, num_workers=config.num_workers)\n        val_loader = DataLoader(val_dataset, batch_size=config.batch_size, \n                              shuffle=False, num_workers=config.num_workers)\n        \n        # Model\n        model = CSIROModel(\n            num_states=len(state_encoder.classes_),\n            num_species=len(species_encoder.classes_)\n        ).to(config.device)\n        \n        # Multi-GPU setup\n        if torch.cuda.device_count() > 1:\n            print(f\"Using {torch.cuda.device_count()} GPUs!\")\n            model = nn.DataParallel(model)\n        \n        # Optimizer and loss\n        optimizer = optim.AdamW(model.parameters(), lr=config.learning_rate, weight_decay=1e-4)\n        scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=config.num_epochs)\n        criterion = MultiTaskLoss()\n        \n        best_val_loss = float('inf')\n        \n        for epoch in range(config.num_epochs):\n            print(f\"\\nEpoch {epoch+1}/{config.num_epochs}\")\n            \n            # Training\n            train_loss = train_epoch(model, train_loader, optimizer, criterion, config.device)\n            print(f\"Train Loss: {train_loss:.4f}\")\n            \n            # Validation\n            val_loss = validate_epoch(model, val_loader, criterion, config.device)\n            print(f\"Val Loss: {val_loss:.4f}\")\n            \n            # Save best model\n            if val_loss < best_val_loss:\n                best_val_loss = val_loss\n                torch.save({\n                    'model_state_dict': model.module.state_dict() if torch.cuda.device_count() > 1 else model.state_dict(),\n                    'optimizer_state_dict': optimizer.state_dict(),\n                    'state_encoder': state_encoder,\n                    'species_encoder': species_encoder,\n                    'epoch': epoch,\n                    'loss': val_loss\n                }, f'best_model_fold_{fold}.pth')\n            \n            scheduler.step()\n        \n        print(f\"Fold {fold+1} completed. Best Val Loss: {best_val_loss:.4f}\")\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-12T05:14:56.646382Z","iopub.execute_input":"2025-11-12T05:14:56.646780Z","iopub.status.idle":"2025-11-12T05:34:59.505545Z","shell.execute_reply.started":"2025-11-12T05:14:56.646730Z","shell.execute_reply":"2025-11-12T05:34:59.504628Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}