{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","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":"none","dataSources":[{"sourceId":112509,"databundleVersionId":14254895,"sourceType":"competition"},{"sourceId":14327780,"sourceType":"datasetVersion","datasetId":9147040},{"sourceId":14348504,"sourceType":"datasetVersion","datasetId":9161315}],"dockerImageVersionId":31236,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install segmentation-models-pytorch albumentations timm scikit-learn opencv-python\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-30T22:57:47.962420Z","iopub.execute_input":"2025-12-30T22:57:47.962977Z","iopub.status.idle":"2025-12-30T22:57:51.418728Z","shell.execute_reply.started":"2025-12-30T22:57:47.962942Z","shell.execute_reply":"2025-12-30T22:57:51.417922Z"},"_kg_hide-output":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import warnings\nwarnings.filterwarnings('ignore')\nimport os\nimport cv2\nimport torch\nimport timm\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport random\nfrom glob import glob\nfrom sklearn.cluster import KMeans\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch import nn\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nimport segmentation_models_pytorch as smp\n\n\ndef seeding(SEED):\n    np.random.seed(SEED)\n    random.seed(SEED)\n    os.environ['PYTHONHASHSEED'] = str(SEED)\n    torch.manual_seed(SEED)\n    if torch.cuda.is_available(): \n        torch.cuda.manual_seed(SEED)\n        torch.cuda.manual_seed_all(SEED)\n        torch.backends.cudnn.deterministic = True\n        torch.backends.cudnn.benchmark = False\n\n\nseeding(32)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-30T22:57:51.420763Z","iopub.execute_input":"2025-12-30T22:57:51.421090Z","iopub.status.idle":"2025-12-30T22:57:57.069360Z","shell.execute_reply.started":"2025-12-30T22:57:51.421058Z","shell.execute_reply":"2025-12-30T22:57:57.068523Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport cv2\nimport numpy as np\nimport joblib\nfrom sklearn.mixture import GaussianMixture\nfrom sklearn.preprocessing import StandardScaler\n\ndef extract_features(img):\n    lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)\n    feats = np.concatenate([img.reshape(-1, 3), lab.reshape(-1, 3)], axis=1)\n    return feats.astype(np.float32)\n\n# --- CONFIGURATION ---\ntraining_images = [\n'/kaggle/input/representative-images/image2.png',\n'/kaggle/input/representative-images/image.png'\n\n]\nn_components = 4\n\n# 1. Collect Data\nall_train_feats = []\nfor p in training_images:\n    img = cv2.imread(p)\n    if img is not None:\n        img = cv2.resize(img, (512, 512)) # Standardize training size\n        all_train_feats.append(extract_features(img))\n\ntrain_matrix = np.vstack(all_train_feats)\n\n# 2. Fit Scaler and GMM\nscaler = StandardScaler()\ntrain_matrix_scaled = scaler.fit_transform(train_matrix)\n\ngmm = GaussianMixture(n_components=n_components, covariance_type=\"tied\", random_state=42)\ngmm.fit(train_matrix_scaled)\n\n# 3. Save models to disk\njoblib.dump(gmm, 'gmm_model.pkl')\njoblib.dump(scaler, 'scaler.pkl')\nprint(\"Training complete. Model and Scaler saved.\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport cv2\nimport numpy as np\nimport joblib\n\ndef run_inference(input_dir, output_dir, confidence_thresh=0.6):\n    # 1. Load the saved models\n    if not os.path.exists('gmm_model.pkl') or not os.path.exists('scaler.pkl'):\n        print(\"Error: Model files not found. Run training script first.\")\n        return\n\n    gmm = joblib.load('/kaggle/working/gmm_model.pkl')\n    scaler = joblib.load('/kaggle/working/scaler.pkl')\n    \n    os.makedirs(output_dir, exist_ok=True)\n    kernel = np.ones((5, 5), np.uint8)\n\n    # 2. Process images\n    for name in sorted(os.listdir(input_dir)):\n        if not name.lower().endswith(('.png', '.jpg', '.jpeg')):\n            continue\n\n        img_path = os.path.join(input_dir, name)\n        img = cv2.imread(img_path)\n        h, w, _ = img.shape\n\n        # Preprocess\n        lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)\n        feats = np.concatenate([img.reshape(-1, 3), lab.reshape(-1, 3)], axis=1).astype(np.float32)\n        \n        # Scale using the TRAINING scaler\n        feats_scaled = scaler.transform(feats)\n\n        # Predict\n        probs = gmm.predict_proba(feats_scaled)\n        labels = probs.argmax(axis=1)\n        confidence = probs.max(axis=1)\n\n        # Create Mask\n        mask = labels.reshape(h, w).astype(np.uint8)\n        conf = confidence.reshape(h, w)\n\n        # Clean up: Mark low confidence as 255 (Background/Unknown)\n        mask[conf < confidence_thresh] = 255\n\n        # Morphological smoothing to remove noise\n        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)\n        \n        # Save\n        out_path = os.path.join(output_dir, f\"{name.split('.')[0]}_mask.{name.split('.')[1]}\")\n        cv2.imwrite(out_path, mask)\n        print(f\"Processed: {name}\")\n\n# --- RUN INFERENCE ---\nrun_inference(\n    input_dir=\"/kaggle/input/csiro-biomass/train\", \n    output_dir=\"/kaggle/working/GMM_masks\"\n)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import shutil\nimport os\n\n# Define the folder you want to zip and the name of the output file\nfolder_to_zip = '/kaggle/working/GMM_masks'\noutput_filename = 'Image2Biomass_masks_dataset'\n\n# This creates 'my_images_dataset.zip' in /kaggle/working/\nshutil.make_archive(output_filename, 'zip', folder_to_zip)\n\nprint(f\"Zip file created: {output_filename}.zip\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import matplotlib.pyplot as plt\n\nsample_name = \"ID1011485656.jpg\"\nIMAGE_DIR = \"/kaggle/input/csiro-biomass/train\"\nMASK_DIR = \"/kaggle/working/GMM_masks\"\n\nimg = cv2.imread(f\"{IMAGE_DIR}/{sample_name}\")\nmask = cv2.imread(\n    f\"{MASK_DIR}/{sample_name.replace('.jpg', '_mask.jpg')}\",\n    cv2.IMREAD_GRAYSCALE\n)\n\n\nimg_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\nmask = np.array(mask, dtype=np.uint8)\n\nplt.figure(figsize=(12,4))\n\nplt.subplot(1,2,1)\nplt.imshow(img_rgb)\nplt.title(\"Original Image\")\nplt.axis(\"off\")\n\nplt.subplot(1,2,2)\nplt.imshow(mask, cmap=\"tab10\")\nplt.title(\"GMM Segmentation Mask\")\nplt.axis(\"off\")\n\nplt.show()\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n# Semantic classes\nCLASS_NAMES = {\n    0: \"Green matter\",\n    1: \"Dead matter\" ,\n    2: \"Clover\",\n    3: \"Background\"\n    \n}\n\n# RGB colors (visible & distinct)\nCLASS_COLORS = {\n    0: (0, 255, 0),       # dark green\n    1: (0, 180, 180),    # cyan   \n    2: (165, 42, 42),      # brown\n    3: (0, 0, 0)    # black\n}\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"CLUSTER_TO_CLASS = {\n    0: 0,  # Background\n    1: 1,  # Clover\n    2: 2,  # Green matter\n    3: 3,  # Dead matter\n}\n\ndef convert_clusters(mask):\n    out = np.full_like(mask, 255)  # ignore by default\n    for k, v in CLUSTER_TO_CLASS.items():\n        out[mask == k] = v\n    return out\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def colorize_mask(mask):\n    h, w = mask.shape\n    color_mask = np.zeros((h, w, 3), dtype=np.uint8)\n\n    for cls, color in CLASS_COLORS.items():\n        color_mask[mask == cls] = color\n\n    return color_mask\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def overlay_mask(image_rgb, mask, alpha=0.45):\n    color_mask = colorize_mask(mask)\n    overlay = cv2.addWeighted(image_rgb, 1 - alpha, color_mask, alpha, 0)\n    return overlay\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from matplotlib.patches import Patch\n\ndef visualize_sample(image_path, gmm_mask_path):\n    # Load image\n    img_bgr = cv2.imread(image_path)\n    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)\n\n    # Load GMM mask\n    gmm_mask = cv2.imread(gmm_mask_path, cv2.IMREAD_GRAYSCALE)\n\n    # Convert to semantic classes\n    sem_mask = convert_clusters(gmm_mask)\n\n    # Create overlay\n    overlay = overlay_mask(img_rgb, sem_mask)\n\n    # -------- Plot --------\n    fig, axes = plt.subplots(1, 4, figsize=(26, 6))\n\n    axes[0].imshow(img_rgb)\n    axes[0].set_title(\"Original Image\", fontsize=14)\n    axes[0].axis(\"off\")\n\n    axes[1].imshow(gmm_mask, cmap=\"tab10\")\n    axes[1].set_title(\"Raw GMM Clusters\", fontsize=14)\n    axes[1].axis(\"off\")\n\n    axes[2].imshow(colorize_mask(sem_mask))\n    axes[2].set_title(\"Semantic Mask\", fontsize=14)\n    axes[2].axis(\"off\")\n\n    axes[3].imshow(overlay)\n    axes[3].set_title(\"Overlay\", fontsize=14)\n    axes[3].axis(\"off\")\n\n    # Legend\n    legend_patches = [\n        Patch(color=np.array(CLASS_COLORS[k])/255, label=v)\n        for k, v in CLASS_NAMES.items()\n    ]\n\n    plt.legend(\n        handles=legend_patches,\n        loc=\"lower center\",\n        bbox_to_anchor=(-1.3, -0.25),\n        ncol=4,\n        fontsize=12\n    )\n\n    plt.tight_layout()\n    plt.show()\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"visualize_sample(\n    image_path=\"/kaggle/input/csiro-biomass/train/ID1035947949.jpg\",\n    gmm_mask_path=\"/kaggle/input/representative-images/Image2Biomass_masks_dataset/ID1035947949_mask.jpg\"\n)\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class GrassDataset(Dataset):\n    def __init__(self, image_dir, mask_dir):\n        self.images = sorted(os.listdir(image_dir))[:1]\n        self.image_dir = image_dir\n        self.mask_dir = mask_dir\n\n    def __len__(self):\n        return len(self.images)\n\n    def __getitem__(self, idx):\n        name = self.images[idx]\n\n        img = cv2.imread(os.path.join(self.image_dir, name))\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n        img = torch.from_numpy(img).permute(2,0,1).float() / 255.0\n\n        mask = cv2.imread(\n            os.path.join(self.mask_dir, name.replace(\".jpg\", \"_mask.jpg\")),\n            cv2.IMREAD_GRAYSCALE\n        )\n        mask = torch.from_numpy(mask).long()\n\n        return img, mask\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-30T22:59:07.068432Z","iopub.execute_input":"2025-12-30T22:59:07.068762Z","iopub.status.idle":"2025-12-30T22:59:07.075196Z","shell.execute_reply.started":"2025-12-30T22:59:07.068735Z","shell.execute_reply":"2025-12-30T22:59:07.074462Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nmodel = smp.Unet(\n    encoder_name=\"resnet101\",\n    encoder_weights=\"imagenet\",\n    in_channels=3,\n    classes=6  # MULTI-CLASS\n).to(device)\n\ncriterion = nn.CrossEntropyLoss(ignore_index = 255)\noptimizer = torch.optim.Adam(model.parameters(), lr=1e-4)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-30T23:07:48.114608Z","iopub.execute_input":"2025-12-30T23:07:48.114915Z","iopub.status.idle":"2025-12-30T23:07:52.264450Z","shell.execute_reply.started":"2025-12-30T23:07:48.114875Z","shell.execute_reply":"2025-12-30T23:07:52.263628Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from PIL import Image\nimg = Image.open('/kaggle/working/GMM_masks/ID1011485656_mask.jpg')\nimg_ = Image.open('/kaggle/input/csiro-biomass/train/ID1011485656.jpg')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-30T22:59:12.009211Z","iopub.execute_input":"2025-12-30T22:59:12.009939Z","iopub.status.idle":"2025-12-30T22:59:12.015532Z","shell.execute_reply.started":"2025-12-30T22:59:12.009907Z","shell.execute_reply":"2025-12-30T22:59:12.014774Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"plt.imshow(img)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-30T22:59:12.487136Z","iopub.execute_input":"2025-12-30T22:59:12.487821Z","iopub.status.idle":"2025-12-30T22:59:13.100496Z","shell.execute_reply.started":"2025-12-30T22:59:12.487791Z","shell.execute_reply":"2025-12-30T22:59:13.099761Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"plt.imshow(img_)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-30T22:58:20.855995Z","iopub.execute_input":"2025-12-30T22:58:20.856843Z","iopub.status.idle":"2025-12-30T22:58:21.195519Z","shell.execute_reply.started":"2025-12-30T22:58:20.856811Z","shell.execute_reply":"2025-12-30T22:58:21.194736Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nimport torch\n\n# Load a sample mask from your dataset4\ndataset = GrassDataset(\n    '/kaggle/input/csiro-biomass/train', \n    '/kaggle/working/GMM_masks'# Explicitly name it\n)\nloader = DataLoader(dataset, batch_size=8, shuffle=True)\n\nfor _, masks in loader:\n    sample_mask = masks\n    print(f\"Unique values in mask: {torch.unique(sample_mask)}\")\n    print(f\"Mask shape: {sample_mask.shape}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-30T22:59:15.794588Z","iopub.execute_input":"2025-12-30T22:59:15.794917Z","iopub.status.idle":"2025-12-30T22:59:15.949648Z","shell.execute_reply.started":"2025-12-30T22:59:15.794878Z","shell.execute_reply":"2025-12-30T22:59:15.948977Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n\nmodel.to(device)\nEPOCHS = 60\n\nfor epoch in range(EPOCHS):\n    model.train()\n    total_loss = 0\n\n    for imgs, masks in loader:\n        imgs, masks = imgs.to(device), masks.to(device)\n\n        optimizer.zero_grad()\n        outputs = model(imgs)\n        loss = criterion(outputs, masks)\n        loss.backward()\n        optimizer.step()\n\n        total_loss += loss.item()\n    torch.save(model.state_dict(), 'best_model_mask.pth')\n    print(f\"Epoch {epoch+1}/{EPOCHS} | Loss: {total_loss/len(loader):.4f}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-30T23:13:14.493714Z","iopub.execute_input":"2025-12-30T23:13:14.494287Z","iopub.status.idle":"2025-12-30T23:15:49.110110Z","shell.execute_reply.started":"2025-12-30T23:13:14.494255Z","shell.execute_reply":"2025-12-30T23:15:49.109104Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def mean_iou(pred, target, num_classes=4):\n    ious = []\n    for cls in range(num_classes):\n        pred_i = pred == cls\n        target_i = target == cls\n        inter = (pred_i & target_i).sum()\n        union = (pred_i | target_i).sum()\n        if union == 0:\n            continue\n        ious.append(inter / union)\n    return np.mean(ious)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-30T23:00:26.748007Z","iopub.execute_input":"2025-12-30T23:00:26.748848Z","iopub.status.idle":"2025-12-30T23:00:26.753201Z","shell.execute_reply.started":"2025-12-30T23:00:26.748816Z","shell.execute_reply":"2025-12-30T23:00:26.752600Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model.eval()\nfinal_masks_dir = '/kaggle/input/representative-images/Image2Biomass_masks_dataset'\nimage = cv2.imread('/kaggle/input/csiro-biomass/train/ID1012260530.jpg')\nimage_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\nresized = cv2.resize(image_rgb, (256,256))\ntensor = torch.tensor(resized/255.).permute(2,0,1).unsqueeze(0).float().to(device)\n\nwith torch.no_grad():\n    pred = model(tensor)\n    pred = torch.argmax(pred, dim=1)[0].cpu().numpy()\n\npred = cv2.resize(pred, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST)\n\nmask_name = os.path.basename('/kaggle/input/csiro-biomass/train/ID1012260530.jpg').replace('.jpg', '_mask.jpg')  # assuming masks are .png\ngt_mask_path = os.path.join(final_masks_dir, mask_name)\ngt_mask = cv2.imread(gt_mask_path, cv2.IMREAD_UNCHANGED)  # load mask as-is\n\n# Check dimensions\nif gt_mask.shape != pred.shape:\n    gt_mask = cv2.resize(gt_mask, (pred.shape[1], pred.shape[0]), interpolation=cv2.INTER_NEAREST)\n\n# Plot\nplt.figure(figsize=(15,5))\nplt.subplot(1,3,1); plt.title(\"Image\"); plt.imshow(image_rgb); plt.axis(\"off\")\nplt.subplot(1,3,2); plt.title(\"Prediction\"); plt.imshow(pred); plt.axis(\"off\")\nplt.subplot(1,3,3); plt.title(\"Ground Truth\"); plt.imshow(gt_mask); plt.axis(\"off\")\nplt.show()\n\n# Compute Mean IoU\niou_score = mean_iou(pred, gt_mask)\nprint(\"Mean IoU:\", iou_score)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-30T23:17:39.406444Z","iopub.execute_input":"2025-12-30T23:17:39.407223Z","iopub.status.idle":"2025-12-30T23:17:40.382403Z","shell.execute_reply.started":"2025-12-30T23:17:39.407191Z","shell.execute_reply":"2025-12-30T23:17:40.381558Z"},"_kg_hide-output":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## **Don't worry this is a base model trained on 1 image 1 epoch**","metadata":{}},{"cell_type":"code","source":"# BiomassMultiTaskDataset\n'''\nclass BiomassMultiTaskDataset(Dataset):\n    def __init__(self, csv_file, image_dir, mask_dir, transform=None):\n        \"\"\"\n        Args:\n            csv_file (string): Path to the Kaggle train.csv (with IDs and biomass targets).\n            image_dir (string): Directory with all the .jpg images.\n            mask_dir (string): Directory with your generated GMM .png masks.\n            transform (callable, optional): Optional transform to be applied on a sample.\n        \"\"\"\n        self.data_frame = pd.read_csv(csv_file)\n        self.image_dir = image_dir\n        self.mask_dir = mask_dir\n        self.transform = transform\n\n    def __len__(self):\n        return len(self.data_frame)\n\n    def __getitem__(self, idx):\n        # 1. Get Image ID from CSV (assuming column name is 'image_id' or 'ID')\n        # Adjust 'image_id' to match your specific CSV header\n        img_id = str(self.data_frame.iloc[idx, 0]) \n        \n        # 2. Get the Biomass Target (assuming it's in the second column)\n        label = torch.tensor(self.data_frame.iloc[idx, 1]).long()\n\n        # 3. Load Image\n        img_name = f\"{img_id}.jpg\"\n        img_path = os.path.join(self.image_dir, img_name)\n        image = cv2.imread(img_path)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n\n        # 4. Load Mask (GMM)\n        mask_name = f\"{img_id}_mask.png\"\n        mask_path = os.path.join(self.mask_dir, mask_name)\n        \n        # Read as grayscale\n        mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)\n        \n        if mask is None:\n            # Fallback if a mask is missing: create a blank one or raise error\n            mask = np.zeros((image.shape[0], image.shape[1]), dtype=np.uint8)\n\n        # 5. Apply Transformations\n        # If using Albumentations, apply to both image and mask simultaneously\n        if self.transform:\n            augmented = self.transform(image=image, mask=mask)\n            image = augmented['image']\n            mask = augmented['mask']\n        else:\n            # Basic conversion to Tensors if no transforms provided\n            image = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0\n            mask = torch.from_numpy(mask).long()\n\n        return image, mask, label\n2. Loading and Testing the Data\nThis snippet shows how to initialize the loader and verify that the shapes are correct before you start training.\n\nPython\n\n# Define paths\nCSV_PATH = \"/kaggle/input/csiro-biomass/train.csv\"\nIMG_DIR = \"/kaggle/input/csiro-biomass/train\"\nMASK_DIR = \"/kaggle/working/final_masks\"\n\n# Initialize Dataset\ndataset = BiomassMultiTaskDataset(\n    csv_file=CSV_PATH,\n    image_dir=IMG_DIR,\n    mask_dir=MASK_DIR\n)\n\n# Initialize DataLoader\ntrain_loader = DataLoader(dataset, batch_size=8, shuffle=True, num_workers=2)\n\n# --- Verification Step ---\nimages, masks, labels = next(iter(train_loader))\n\nprint(f\"Batch Image Shape: {images.shape}\") # Expect: [8, 3, H, W]\nprint(f\"Batch Mask Shape:  {masks.shape}\")  # Expect: [8, H, W]\nprint(f\"Batch Label Shape: {labels.shape}\") # Expect: [8]\n'''","metadata":{"trusted":true,"jupyter":{"source_hidden":true},"_kg_hide-output":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# BiomassMultiTaskModel\n\n'''\nimport torch\nimport torch.nn as nn\nimport torchvision.models as models\n\nclass BiomassMultiTaskModel(nn.Module):\n    def __init__(self, n_seg_classes, n_biomass_classes):\n        super(BiomassMultiTaskModel, self).__init__()\n        \n        # 1. Shared Backbone (using ResNet18 for speed)\n        # We take everything except the last pooling/fc layer\n        backbone = models.resnet18(pretrained=True)\n        self.features = nn.Sequential(*list(backbone.children())[:-2]) # Output shape: [Batch, 512, H/32, W/32]\n\n        # 2. Segmentation Head (Upsampling back to original image size)\n        # Using simple bilinear upsampling + convolution\n        self.seg_head = nn.Sequential(\n            nn.Conv2d(512, 256, kernel_size=3, padding=1),\n            nn.ReLU(),\n            nn.Conv2d(256, n_seg_classes, kernel_size=1),\n            nn.Upsample(scale_factor=32, mode='bilinear', align_corners=True) \n        )\n\n        # 3. Classification Head (Global Pool + Linear)\n        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))\n        self.cls_head = nn.Linear(512, n_biomass_classes)\n\n    def forward(self, x):\n        x = self.features(x)\n        \n        # Segmentation output\n        seg_out = self.seg_head(x)\n        \n        # Classification output\n        flat_x = self.avgpool(x).view(x.size(0), -1)\n        cls_out = self.cls_head(flat_x)\n        \n        return seg_out, cls_out\n'''","metadata":{"trusted":true,"jupyter":{"source_hidden":true},"_kg_hide-output":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Training a MultiModel Segmentaion teacher and image input with Biomass target\n'''\nimport torch\n\n# Configuration\nbest_val_loss = float('inf')\npatience = 5\ncounter = 0\nmodel_path = \"best_biomass_mtl_model.pth\"\n\n# Weights\nw_seg, w_cls = 0.4, 1.0\n\nfor epoch in range(num_epochs):\n    model.train()\n    train_loss = 0\n    \n    for imgs, masks, labels in train_loader:\n        imgs, masks, labels = imgs.to(device), masks.to(device), labels.to(device)\n        \n        optimizer.zero_grad()\n        pred_masks, pred_labels = model(imgs)\n        \n        loss_seg = criterion_seg(pred_masks, masks)\n        loss_cls = criterion_cls(pred_labels, labels)\n        total_loss = (w_seg * loss_seg) + (w_cls * loss_cls)\n        \n        total_loss.backward()\n        optimizer.step()\n        train_loss += total_loss.item()\n\n    # Validation Phase\n    model.eval()\n    val_loss = 0\n    with torch.no_grad():\n        for imgs, masks, labels in val_loader:\n            imgs, masks, labels = imgs.to(device), masks.to(device), labels.to(device)\n            p_masks, p_labels = model(imgs)\n            v_loss = (w_seg * criterion_seg(p_masks, masks)) + (w_cls * criterion_cls(p_labels, labels))\n            val_loss += v_loss.item()\n    \n    avg_val_loss = val_loss / len(val_loader)\n    print(f\"Epoch {epoch}: Val Loss: {avg_val_loss:.4f}\")\n\n    # --- Early Stopping & Saving ---\n    if avg_val_loss < best_val_loss:\n        best_val_loss = avg_val_loss\n        torch.save(model.state_dict(), model_path)\n        print(\" Model Saved\")\n        counter = 0\n    else:\n        counter += 1\n        if counter >= patience:\n            print(\" Early stopping triggered\")\n            break\n\n  '''          ","metadata":{"trusted":true,"jupyter":{"source_hidden":true},"_kg_hide-output":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# r2 and miou metrics\n'''\nimport numpy as np\nfrom sklearn.metrics import r2_score\n\ndef calculate_segmentation_stats(pred_masks, true_masks, n_classes=6):\n    \"\"\"Calculates Mean IoU for the GMM masks.\"\"\"\n    preds = torch.argmax(pred_masks, dim=1).cpu().numpy()\n    targets = true_masks.cpu().numpy()\n    \n    ious = []\n    for cls in range(n_classes):\n        intersection = np.logical_and(preds == cls, targets == cls).sum()\n        union = np.logical_or(preds == cls, targets == cls).sum()\n        if union == 0:\n            ious.append(float('nan')) # Ignore classes not present\n        else:\n            ious.append(intersection / union)\n    return np.nanmean(ious)\n\ndef calculate_biomass_metrics(all_preds, all_labels):\n    \"\"\"\n    Calculates the weighted R2 score as per CSIRO competition rules.\n    Weights: Green(0.1), Dead(0.1), Clover(0.1), GDM(0.2), Total(0.5)\n    \"\"\"\n    weights = np.array([0.1, 0.1, 0.1, 0.2, 0.5])\n    \n    # Log-transform targets as per competition requirements: log(y + 1)\n    # Ensure all_preds and all_labels are [N, 5]\n    r2_scores = []\n    for i in range(5):\n        score = r2_score(all_labels[:, i], all_preds[:, i])\n        r2_scores.append(max(0, score)) # R2 can be negative; competition usually clips\n        \n    weighted_score = np.sum(np.array(r2_scores) * weights)\n    return weighted_score, r2_scores\n\n  '''  ","metadata":{"trusted":true,"jupyter":{"source_hidden":true},"_kg_hide-output":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# evaluation function\n'''\nimport torch\nfrom tqdm import tqdm\n\ndef evaluate_model(model, val_loader, device):\n    model.eval()\n    \n    all_bio_preds = []\n    all_bio_targets = []\n    total_miou = 0\n    \n    print(\"Running Evaluation...\")\n    with torch.no_grad():\n        for imgs, masks, labels in tqdm(val_loader):\n            imgs = imgs.to(device)\n            masks = masks.to(device)\n            labels = labels.to(device)\n            \n            # 1. Forward Pass\n            pred_seg, pred_bio = model(imgs)\n            \n            # 2. Collect Biomass Predictions (Classification/Regression)\n            # We move them to CPU and convert to numpy for sklearn metrics\n            all_bio_preds.append(pred_bio.cpu().numpy())\n            all_bio_targets.append(labels.cpu().numpy())\n            \n            # 3. Calculate Segmentation Metric for this batch\n            batch_miou = calculate_segmentation_stats(pred_seg, masks)\n            total_miou += batch_miou\n\n    # Flatten lists into large numpy arrays\n    all_bio_preds = np.vstack(all_bio_preds)\n    all_bio_targets = np.vstack(all_bio_targets)\n    \n    # 4. Final Metric Calculations\n    weighted_r2, r2_per_class = calculate_biomass_metrics(all_bio_preds, all_bio_targets)\n    mean_miou = total_miou / len(val_loader)\n    \n    # --- Print Report ---\n    print(\"\\n\" + \"=\"*30)\n    print(f\"FINAL WEIGHTED R2: {weighted_r2:.4f}\")\n    print(f\"MEAN SEGMENTATION IoU: {mean_miou:.4f}\")\n    print(\"-\" * 30)\n    categories = ['Green', 'Dead', 'Clover', 'GDM', 'Total']\n    for cat, score in zip(categories, r2_per_class):\n        print(f\"{cat} R2 Score: {score:.4f}\")\n    print(\"=\"*30)\n    \n    return weighted_r2, mean_miou\n\n# How to use it:\n# final_r2, final_miou = evaluate_model(model, val_loader, device)\n'''","metadata":{"trusted":true,"jupyter":{"source_hidden":true},"_kg_hide-output":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# performing Inference\n'''\n# Load the best version found during training\nmodel.load_state_dict(torch.load(\"best_biomass_mtl_model.pth\"))\n\n# Run the full evaluation\nfinal_r2, final_miou = evaluate_model(model, val_loader, device)\n'''","metadata":{"trusted":true,"jupyter":{"source_hidden":true},"_kg_hide-output":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Final score\n'''\nprint(f\"the final r2 score is : {final_r2}\")\nprint (f\"the final mIoU is : {final_miou}\")\n'''","metadata":{"trusted":true,"jupyter":{"source_hidden":true},"_kg_hide-output":true},"outputs":[],"execution_count":null}]}