{"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.11.13"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":98498,"databundleVersionId":12148564,"sourceType":"competition"}],"dockerImageVersionId":31153,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"id":"bcbcb0c9","cell_type":"code","source":"#pip install segmentation-models-pytorch","metadata":{},"outputs":[],"execution_count":null},{"id":"4f43cd02","cell_type":"code","source":"\n# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nimport json\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\nfrom PIL import Image, ImageDraw\nimport os\nfrom torch.utils.data import Dataset, DataLoader\nfrom tqdm import tqdm\n#for dirname, _, filenames in os.walk('/kaggle/input'):\n#    for filename in filenames:\n#        print(os.path.join(dirname, filename))\nfrom torchvision.transforms import v2 as T\nfrom sklearn.model_selection import train_test_split\nfrom torchmetrics.classification import BinaryJaccardIndex, BinaryF1Score\nimport torchvision \nimport torch.nn as nn\nimport torch\nfrom matplotlib import pyplot as plt\nimport seaborn as sns\nimport segmentation_models_pytorch as smp\nimport base64\nimport zlib\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{},"outputs":[],"execution_count":null},{"id":"e1d836a9","cell_type":"code","source":"\ndef polygon_to_mask(segmentation, height, width):\n    mask = Image.new(\"L\", (width, height), 0)   # empty black image\n    for poly in segmentation:                   # some annotations contain multiple polygons\n        xy = np.array(poly).reshape((-1, 2))\n        ImageDraw.Draw(mask).polygon(xy.flatten().tolist(), outline=1, fill=1)\n    return np.array(mask)\n","metadata":{},"outputs":[],"execution_count":null},{"id":"f1d15488","cell_type":"code","source":"path = ('/kaggle/input/garbage-segmentation/garbage_segmentation_dataset/train')\nraw_data = json.load(open('/kaggle/input/garbage-segmentation/garbage_segmentation_dataset/train_json.json'))","metadata":{},"outputs":[],"execution_count":null},{"id":"cfa9b2ce","cell_type":"code","source":"#raw_data","metadata":{},"outputs":[],"execution_count":null},{"id":"e9854a59","cell_type":"code","source":"a = os.listdir(path)","metadata":{},"outputs":[],"execution_count":null},{"id":"d9830f10","cell_type":"code","source":"a[-1]","metadata":{},"outputs":[],"execution_count":null},{"id":"1d68328d","cell_type":"code","source":"len([(f.startswith('image') or f.startswith('batch')) for f in os.listdir(path)])","metadata":{},"outputs":[],"execution_count":null},{"id":"51ca4454","cell_type":"code","source":"","metadata":{},"outputs":[],"execution_count":null},{"id":"b6bad8d5","cell_type":"code","source":"e = Image.open(os.path.join(path, a[-1])).convert('RGB')\ndef img_show(e):\n    to_tensor = T.ToTensor()\n    e = to_tensor(e)#.numpy()\n    e = e.permute(1, 2, 0)\n    plt.imshow(e)\nimg_show(e)","metadata":{},"outputs":[],"execution_count":null},{"id":"1eb4c444","cell_type":"code","source":"images = pd.DataFrame(raw_data['images'])\nannotations = pd.DataFrame(raw_data['annotations'])\ncategories = pd.DataFrame(raw_data['categories'])","metadata":{},"outputs":[],"execution_count":null},{"id":"fd97730f","cell_type":"code","source":"train = annotations.merge(images, left_on=\"image_id\", right_on=\"id\", how=\"left\")\ntrain.head()\ntrain = train.drop(columns=['ignore', 'iscrowd', 'license', 'flickr_url', 'coco_url', 'date_captured', 'flickr_640_url'])\ntrain.head()","metadata":{},"outputs":[],"execution_count":null},{"id":"8892551f","cell_type":"code","source":"transforms = T.Compose([\n    T.ToImage(), \n    #T.Resize((256,256)),\n    T.RandomRotation(20),\n    T.ToDtype(torch.float32, scale=True),\n])","metadata":{},"outputs":[],"execution_count":null},{"id":"1e1c701d","cell_type":"code","source":"class GarbageDataset(Dataset):\n    def __init__(self, df, img_path, transform=None):\n        self.df = df\n        self.img_path = img_path\n        self.transform = transform\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        item = self.df.iloc[idx]\n\n        # --- Load image as PIL ---\n        img = Image.open(os.path.join(self.img_path, item['file_name'])).convert(\"RGB\")\n\n        # --- Convert polygon → mask (H, W) ---\n        mask = polygon_to_mask(item.segmentation, item.height, item.width)  # numpy H,W\n        mask = torch.from_numpy(mask).float()  # [H,W]\n        mask = mask.unsqueeze(0)               # [1,H,W]\n\n        # --- Resize both ---\n        img = T.functional.resize(img, (256, 256))\n        mask = T.functional.resize(\n            mask, (256, 256),\n            interpolation=T.InterpolationMode.NEAREST\n        )\n\n        # --- For transforms, align formats ---\n        if self.transform:\n            img, mask = self.transform(img, mask)\n\n        return img, mask  # img:[3,512,512], mask:[1,512,512]\n","metadata":{},"outputs":[],"execution_count":null},{"id":"d87010c0","cell_type":"code","source":"item = train.iloc[0]\nimg = Image.open(os.path.join('/kaggle/input/garbage-segmentation/garbage_segmentation_dataset/train', item['file_name'])).convert(\"RGB\")\n      \nmask = polygon_to_mask(item.segmentation, item.height, item.width)\nimg, mask = transforms(img, mask)\nmask\n            ","metadata":{},"outputs":[],"execution_count":null},{"id":"bb2bce1c","cell_type":"code","source":"img_show(img) \n#img_show(mask)\nplt.show()","metadata":{},"outputs":[],"execution_count":null},{"id":"6b2d2c56","cell_type":"code","source":"img_show(mask)","metadata":{},"outputs":[],"execution_count":null},{"id":"4850bac2","cell_type":"code","source":"unique_ids = train['image_id'].unique()\n","metadata":{},"outputs":[],"execution_count":null},{"id":"443a97f1","cell_type":"code","source":"train_ids, val_ids = train_test_split(unique_ids[:300], test_size=0.2, random_state=42)\ntrain_df = train[train['image_id'].isin(train_ids)].reset_index(drop=True)\nval_df = train[train['image_id'].isin(val_ids)].reset_index(drop=True)\n","metadata":{},"outputs":[],"execution_count":null},{"id":"1f92b059","cell_type":"code","source":"batch_size = 16","metadata":{},"outputs":[],"execution_count":null},{"id":"a9add5c8","cell_type":"code","source":"train_dataset = GarbageDataset(train_df, path, transforms)\nval_dataset = GarbageDataset(val_df, path, transforms)\ntrain_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\nval_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)\n","metadata":{},"outputs":[],"execution_count":null},{"id":"28c8a9dc","cell_type":"code","source":"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\ndevice","metadata":{},"outputs":[],"execution_count":null},{"id":"5c5f1684","cell_type":"code","source":"\nmodel = smp.Unet(\n    encoder_name=\"resnet18\",      # choose encoder/backbone\n    encoder_weights=None,   # pretrained on ImageNet\n    in_channels=3,                # number of input channels (RGB)\n    classes=5,                    # number of output classes (1 mask)\n).to(device)\n","metadata":{},"outputs":[],"execution_count":null},{"id":"f91a0179","cell_type":"code","source":"from torchmetrics.classification import BinaryJaccardIndex, BinaryF1Score\nfrom torchmetrics.classification import MulticlassJaccardIndex, MulticlassF1Score\n\nmetric_iou = MulticlassJaccardIndex(num_classes=5).to(device)\nmetric_dice = MulticlassF1Score(num_classes=5).to(device)  # Dice = F1 for segmentation\n\n#metric_iou = BinaryJaccardIndex(threshold=0.5)\n#metric_dice = BinaryF1Score(threshold=0.5)\n\ndice_loss = smp.losses.DiceLoss(mode=\"multiclass\").to(device)\n#criterion = nn.CrossEntropyLoss()\nce_loss = nn.CrossEntropyLoss().to(device)\n\ndef criterion(logits, masks):\n    return 0.5 * ce_loss(logits, masks) + 0.5 * dice_loss(logits, masks)\n\noptimizer = torch.optim.AdamW(model.parameters(), lr=0.001, weight_decay=1e-4)","metadata":{},"outputs":[],"execution_count":null},{"id":"1741c16e","cell_type":"code","source":"scaler = torch.cuda.amp.GradScaler()      # for mixed precision\nscheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n    optimizer, mode=\"max\", patience=3, verbose=True\n)\n\nbest_iou = 0.0\nsave_path = \"./best_unet.pth\"\n","metadata":{},"outputs":[],"execution_count":null},{"id":"1281138f","cell_type":"code","source":"EPOCHS = 1\n","metadata":{},"outputs":[],"execution_count":null},{"id":"bc5858a4","cell_type":"code","source":"epoch = 0","metadata":{},"outputs":[],"execution_count":null},{"id":"cb3ff8ac","cell_type":"code","source":"\n\nfor i in range(EPOCHS):\n    running_loss = 0.0\n    train_iou = 0.0\n    train_dice = 0.0\n\n    model.train()\n    loop = tqdm(train_loader, desc=f\"Epoch {epoch+1} [TRAIN]\", dynamic_ncols=True)\n\n    running_loss = 0.0\n\n    for images, masks in loop:\n        images, masks = images.to(device), masks.to(device)\n\n        optimizer.zero_grad()\n\n        with torch.amp.autocast('cuda'):\n            logits = model(images)\n            masks = masks.squeeze(1).long()          # (B, 1, H, W) -> (B, H, W)\n            loss = criterion(logits, masks)\n\n\n        scaler.scale(loss).backward()\n        scaler.step(optimizer)\n        scaler.update()\n\n        running_loss += loss.item()\n        preds = torch.argmax(logits, dim=1)\n        train_iou += metric_iou(preds, masks).item()\n        train_dice += metric_dice(preds, masks).item()\n\n        loop.set_postfix({\n            \"loss\": f\"{running_loss / (loop.n+1):.4f}\",\n            \"IoU\": f\"{train_iou / (loop.n+1):.4f}\",\n            \"Dice\": f\"{train_dice / (loop.n+1):.4f}\",\n        })\n\n\n\n\n    model.eval()\n    val_loss = 0.0\n    val_iou = 0.0\n    val_dice = 0.0\n\n    loop = tqdm(val_loader, desc=f\"         [VAL]\", dynamic_ncols=True)\n\n    with torch.no_grad():\n        for images, masks in loop:\n            images, masks = images.to(device), masks.to(device)\n\n            logits = model(images)\n            masks = masks.squeeze(1).long()\n            loss = criterion(logits, masks)\n            val_loss += loss.item()\n\n            preds = torch.argmax(logits, dim=1)     # (B, H, W)\n\n            val_iou += metric_iou(preds, masks).item()\n            val_dice += metric_dice(preds, masks).item()\n\n            loop.set_postfix({\n            \"loss\": f\"{val_loss / (loop.n+1):.4f}\",\n            \"IoU\": f\"{val_iou / (loop.n+1):.4f}\",\n            \"Dice\": f\"{val_dice / (loop.n+1):.4f}\",\n        })\n\n    val_loss /= len(val_loader)\n    val_iou /= len(val_loader)\n    val_dice /= len(val_loader)\n    epoch += 1\n\n\n    # Save best model\n    if val_iou > best_iou:\n        best_iou = val_iou\n        torch.save(model.state_dict(), save_path)\n        print(\"✅ Saved best model!\")\n\n    scheduler.step(val_iou)  # scheduler uses IoU (not loss)\n\n\n\n\nprint(f\"\\nTraining finished. Best model saved at: {save_path}\")","metadata":{},"outputs":[],"execution_count":null},{"id":"0410947d","cell_type":"code","source":"np.array([i != 0 for i in preds.cpu()]).sum()","metadata":{},"outputs":[],"execution_count":null},{"id":"5e746330","cell_type":"code","source":"preds.shape","metadata":{},"outputs":[],"execution_count":null},{"id":"33594164","cell_type":"code","source":"pred_mask = torch.argmax(preds.detach().cpu(), dim=0)\npred_mask","metadata":{},"outputs":[],"execution_count":null},{"id":"a432d7a6","cell_type":"code","source":"pred_mask.shape","metadata":{},"outputs":[],"execution_count":null},{"id":"4d2335e7","cell_type":"code","source":"sns.heatmap(pred_mask)","metadata":{},"outputs":[],"execution_count":null},{"id":"3d3ed884","cell_type":"code","source":"def encode_mask(mask_binary: np.ndarray, class_id: int):\n    mask_binary = mask_binary.astype(np.uint8)\n\n    compressed = zlib.compress(mask_binary.tobytes())\n    encoded = base64.b64encode(compressed).decode(\"utf-8\")\n\n    return f\"{class_id} {encoded}\"","metadata":{},"outputs":[],"execution_count":null},{"id":"0d722c55","cell_type":"code","source":"pred_mask = pred_mask.cpu().numpy().astype(np.uint8)","metadata":{},"outputs":[],"execution_count":null},{"id":"0aba59c0","cell_type":"code","source":"\n\nencoded_strings = []\n\nfor class_id in range(9):\n    binary_mask = (pred_mask == class_id).astype(np.uint8)  # (H, W)\n    encoded = encode_mask(binary_mask, class_id)\n    encoded_strings.append(encoded)\n\nsubmission_string = \" \".join(encoded_strings)\n\n","metadata":{},"outputs":[],"execution_count":null},{"id":"06ec0db5","cell_type":"code","source":"submission_string","metadata":{},"outputs":[],"execution_count":null},{"id":"3331e031","cell_type":"code","source":"sample = pd.read_csv('/kaggle/input/garbage-segmentation/garbage_segmentation_dataset/sample_submission.csv')","metadata":{},"outputs":[],"execution_count":null},{"id":"c97d75e7","cell_type":"code","source":"sample.pred_str[1]","metadata":{},"outputs":[],"execution_count":null},{"id":"fc2dec90","cell_type":"markdown","source":"test","metadata":{}},{"id":"eef26846","cell_type":"code","source":"test_path = '/kaggle/input/garbage-segmentation/garbage_segmentation_dataset/test'\ntest_data = json.load(open('/kaggle/input/garbage-segmentation/garbage_segmentation_dataset/test-without_segmentation.json'))","metadata":{},"outputs":[],"execution_count":null},{"id":"c5f33415","cell_type":"code","source":"test = pd.DataFrame(test_data[\"images\"])","metadata":{},"outputs":[],"execution_count":null},{"id":"87d65d51","cell_type":"code","source":"test = test.drop(columns=['license', 'flickr_url', 'coco_url', 'date_captured', 'flickr_640_url'])","metadata":{},"outputs":[],"execution_count":null},{"id":"a794e7d3","cell_type":"code","source":"test","metadata":{},"outputs":[],"execution_count":null},{"id":"7c677274","cell_type":"code","source":"class TestGarbageDataset(Dataset):\n    def __init__(self, df, img_path, transform=None):\n        self.df = df\n        self.img_path = img_path\n        self.transform = transform\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        item = self.df.iloc[idx]\n\n        # --- Load image as PIL ---\n        img = Image.open(os.path.join(self.img_path, item['file_name'])).convert(\"RGB\")\n\n        # --- Resize both ---\n        img = T.functional.resize(img, (256, 256))\n\n        # --- For transforms, align formats ---\n        if self.transform:\n            img = self.transform(img)\n\n        return img  # img:[3,512,512], mask:[1,512,512]\n","metadata":{},"outputs":[],"execution_count":null},{"id":"aefe8782","cell_type":"code","source":"test_dataset = TestGarbageDataset(test, test_path, transform=transforms)\ntest_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)","metadata":{},"outputs":[],"execution_count":null},{"id":"bdf0e31d","cell_type":"code","source":"all_preds = []\nfor image in test_loader:\n    model.to(torch.device('cpu'))\n    output = model(image)\n    all_preds.append(output)\n    print(\"1\\n\")\n","metadata":{},"outputs":[],"execution_count":null},{"id":"773e94c4","cell_type":"code","source":"len(all_preds)","metadata":{},"outputs":[],"execution_count":null},{"id":"e8ec4106","cell_type":"code","source":"all_mask = []\nfor pred in all_preds:\n    pred_mask = torch.argmax(pred[0].detach().cpu(), dim=0)\n    all_mask.append(pred_mask)\n","metadata":{},"outputs":[],"execution_count":null},{"id":"f4188858","cell_type":"code","source":"pred.shape","metadata":{},"outputs":[],"execution_count":null},{"id":"d3bd0f15","cell_type":"code","source":"sns.heatmap(all_mask[7])","metadata":{},"outputs":[],"execution_count":null}]}