workshop_15_03_2025_segmentation_fixed.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "965aba33-a7f7-4c4f-ad70-110e218e3922",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9bfeb12-d3af-488e-891b-49c77e38dc7e",
   "metadata": {},
   "source": [
    "## Install Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7ea3765d-fe0e-4cb1-983e-43339bb88198",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ! pip install -U lightning ultralytics"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ff79e71-b83d-405f-8e7e-5d8eec7a00c8",
   "metadata": {},
   "source": [
    "## Organize Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "abf6fb48-0a1c-447c-907b-719ccd8ffa04",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import gc\n",
    "from pathlib import Path\n",
    "import requests\n",
    "import zipfile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1a8979b8-0cfb-491c-b2f0-87fa47c6962b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "37eaff67-a4d8-4be5-a976-3efca5487b5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d52e87f4-e8a6-4f00-99c8-e25c00c5bf1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "136e7146-d005-4eda-80be-9e757ddcf86c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import Dataset, DataLoader, random_split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b7a87569-7cbb-4070-8186-eb9e4c877e1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchmetrics.classification import Dice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "8dff2759-e577-4017-9450-e837e0ab9d21",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision.models import resnet34, resnet\n",
    "from torchvision.models.detection import mask_rcnn, maskrcnn_resnet50_fpn, maskrcnn_resnet50_fpn_v2\n",
    "from torchvision.models.detection.faster_rcnn import FastRCNNPredictor\n",
    "from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor\n",
    "from torchvision.datasets import OxfordIIITPet, Cityscapes, wrap_dataset_for_transforms_v2\n",
    "from torchvision.transforms import v2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "bd733e83-e4f1-4f32-a1b6-7f28bdb6885a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import lightning as pl\n",
    "from lightning.pytorch.callbacks import ModelCheckpoint\n",
    "from lightning.pytorch.loggers import TensorBoardLogger"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "d3a4ccf1-3fd8-4842-b6c0-28f7b20dfbd6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ultralytics import YOLO"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce67380f-76b5-4763-9cb9-099d96e47f92",
   "metadata": {},
   "source": [
    "## Initialize Folders for Model and Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "89b1ab04-805c-49d4-851b-ae16cf462b85",
   "metadata": {},
   "outputs": [],
   "source": [
    "def makedirs(dir_path: Path):\n",
    "    dir_path.mkdir(exist_ok=True, parents=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "f035353e-5682-42fb-95b5-4e52ff8bac0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA = Path('data')\n",
    "MODELS = Path('models')\n",
    "unet_path = MODELS / 'unset'\n",
    "yolo_path = MODELS / 'yolo11'\n",
    "pets_path = DATA / 'pets'\n",
    "carparts_path = DATA / 'carparts'\n",
    "cracks_path = DATA / 'cracks'\n",
    "packages_path = DATA / 'packages'\n",
    "makedirs(unet_path)\n",
    "makedirs(yolo_path)\n",
    "makedirs(pets_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de0bae89-1f6b-44e5-bfd6-ab618779c888",
   "metadata": {},
   "source": [
    "## Initialize Device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "b511aa17-7dfb-407e-86b7-bacdb83f7406",
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_device():\n",
    "    # For the most part I'll try to import functions and classes near\n",
    "    # where they are used\n",
    "    # to make it clear where they come from.\n",
    "    if torch.backends.mps.is_available():\n",
    "        device = 'mps'\n",
    "    else:\n",
    "        device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "\n",
    "    print(f'Device: {device}')\n",
    "\n",
    "    return device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "a6644c11-583c-4965-89ae-eab3bb17150a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Device: mps\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'mps'"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "device = init_device()\n",
    "device"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "caf0802e",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Section A: UNet Segmentation Training with Pretrained ResNet34 Backbone\n",
    "\n",
    "In this section we will train a UNet model for semantic segmentation. The encoder uses a pretrained ResNet34 backbone. We will use the Oxford-IIIT Pet dataset (downloaded via TorchVision) and train our UNet model using PyTorch Lightning."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1a701e8-9b24-403e-98b3-4280438c18df",
   "metadata": {},
   "source": [
    "#### Prepare the Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "a57aedcf",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ConvBlock(nn.Module):\n",
    "    def __init__(self, in_ch, out_ch):\n",
    "        super().__init__()\n",
    "        self.conv = nn.Sequential(\n",
    "            nn.Conv2d(in_ch, out_ch, 3, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(out_ch, out_ch, 3, padding=1),\n",
    "            nn.ReLU()\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.conv(x)\n",
    "\n",
    "class UNetResNet34(nn.Module):\n",
    "    def __init__(self, n_classes=1):\n",
    "        super().__init__()\n",
    "        base = resnet34(weights=resnet.ResNet34_Weights.DEFAULT)\n",
    "        self.enc0 = nn.Sequential(base.conv1, base.bn1, base.relu)\n",
    "        self.enc1 = nn.Sequential(base.maxpool, base.layer1)\n",
    "        self.enc2 = base.layer2\n",
    "        self.enc3 = base.layer3\n",
    "        self.enc4 = base.layer4\n",
    "\n",
    "        self.dec3 = ConvBlock(512+256, 256)\n",
    "        self.dec2 = ConvBlock(256+128, 128)\n",
    "        self.dec1 = ConvBlock(128+64, 64)\n",
    "        self.dec0 = ConvBlock(64+64, 32)\n",
    "        self.final = nn.Conv2d(32, n_classes, kernel_size=1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        enc0 = self.enc0(x)\n",
    "        enc1 = self.enc1(enc0)\n",
    "        enc2 = self.enc2(enc1)\n",
    "        enc3 = self.enc3(enc2)\n",
    "        enc4 = self.enc4(enc3)\n",
    "\n",
    "        dec3 = self.dec3(torch.cat([F.interpolate(enc4, enc3.size()[2:]), enc3], dim=1))\n",
    "        dec2 = self.dec2(torch.cat([F.interpolate(dec3, enc2.size()[2:]), enc2], dim=1))\n",
    "        dec1 = self.dec1(torch.cat([F.interpolate(dec2, enc1.size()[2:]), enc1], dim=1))\n",
    "        dec0 = self.dec0(torch.cat([F.interpolate(dec1, enc0.size()[2:]), enc0], dim=1))\n",
    "\n",
    "        out = self.final(dec0)   # [B, 1, 59, 59]\n",
    "\n",
    "        # Upsample explicitly to match target (118,118)\n",
    "        out = F.interpolate(out, size=x.shape[2:], mode='bilinear', align_corners=False)\n",
    "    \n",
    "        return torch.sigmoid(out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "3b43ef81",
   "metadata": {},
   "outputs": [],
   "source": [
    "class LitUNet(pl.LightningModule):\n",
    "    def __init__(self, model):\n",
    "        super().__init__()\n",
    "        self.model = model\n",
    "        self.loss_fn = nn.BCELoss()\n",
    "        self.dice = Dice(num_classes=2, average='micro')\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.model(x)\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        x, y = batch\n",
    "        y_hat = self(x)\n",
    "        loss = self.loss_fn(y_hat, y)\n",
    "        dice_score = self.dice((y_hat>0.5).int(), y.int())\n",
    "        self.log(\"train_loss\", loss, prog_bar=True, on_step=False, on_epoch=True)\n",
    "        self.log(\"train_dice\", dice_score, prog_bar=True, on_step=False, on_epoch=True)\n",
    "        \n",
    "        return loss\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        x, y = batch\n",
    "        y_hat = self(x)\n",
    "        loss = self.loss_fn(y_hat, y)\n",
    "        dice_score = self.dice((y_hat>0.5).int(), y.int())\n",
    "        self.log(\"val_loss\", loss, prog_bar=True, on_step=False, on_epoch=True)\n",
    "        self.log(\"val_dice\", dice_score, prog_bar=True, on_step=False, on_epoch=True)\n",
    "\n",
    "        return loss\n",
    "\n",
    "    def test_step(self, batch, batch_idx):\n",
    "        x, y = batch\n",
    "        y_hat = self(x)\n",
    "        loss = self.loss_fn(y_hat, y)\n",
    "        dice_score = self.dice((y_hat>0.5).int(), y.int())\n",
    "        self.log(\"test_loss\", loss, prog_bar=True, on_step=False, on_epoch=True)\n",
    "        self.log(\"test_dice\", dice_score, prog_bar=True, on_step=False, on_epoch=True)\n",
    "\n",
    "        return loss\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        optimizer = torch.optim.AdamW(self.parameters(), lr=1e-3, weight_decay=1e-4)\n",
    "        scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=5, T_mult=2)\n",
    "        \n",
    "        return [optimizer], [scheduler]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e60d10a-d75e-4d4d-9375-77192bc877b9",
   "metadata": {},
   "source": [
    "#### Initialize Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "441bc413-af64-4cc6-8fe2-836ceb35e2e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TargetTransform(object):\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "    def forward(self, x):\n",
    "        return (x - 1).clamp(0, 1).float()\n",
    "\n",
    "    def __call__(self, *args, **kwargs):\n",
    "        return self.forward(*args, **kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "5dd94757-7d69-4d3d-8214-52768dee7742",
   "metadata": {},
   "outputs": [],
   "source": [
    "workers=0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "687e0d8d-1ad6-4d9a-959b-9e8fd706ea03",
   "metadata": {},
   "outputs": [],
   "source": [
    "class JointTransform:\n",
    "    def __init__(self, base_size, crop_size):\n",
    "        self.transforms = v2.Compose([\n",
    "            v2.Resize(base_size),\n",
    "            v2.RandomHorizontalFlip(p=0.1),\n",
    "            v2.RandomVerticalFlip(p=0.1),\n",
    "            v2.RandomRotation(30),\n",
    "            v2.RandomCrop(crop_size),\n",
    "            v2.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),\n",
    "            v2.RandomGrayscale(p=0.1),\n",
    "            v2.ToTensor(),\n",
    "        ])\n",
    "\n",
    "    def __call__(self, image, target):\n",
    "        image = self.transforms(image, target)\n",
    "        target = v2.ToTensor()(target).float()\n",
    "        return image, target"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "c56b6fe5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define outside PetDataModule at the global level\n",
    "class PetDataModule(pl.LightningDataModule):\n",
    "    def setup(self, stage=None):\n",
    "        transforms_img = v2.Compose([\n",
    "            v2.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),\n",
    "            v2.RandomGrayscale(p=0.1),\n",
    "            v2.Resize((118, 118)),\n",
    "            v2.ToTensor(),\n",
    "        ])\n",
    "        transforms_target = v2.Compose([\n",
    "            v2.Resize((118, 118)),\n",
    "            v2.PILToTensor(),\n",
    "            TargetTransform(),\n",
    "        ])\n",
    "        joint_transforms = v2.Compose([\n",
    "            v2.Resize(118),\n",
    "            v2.RandomHorizontalFlip(p=0.1),\n",
    "            v2.RandomVerticalFlip(p=0.1),\n",
    "            v2.RandomRotation(30),\n",
    "            v2.RandomCrop(100),\n",
    "            # v2.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),\n",
    "            # v2.RandomGrayscale(p=0.1),\n",
    "            # v2.ToTensor(),\n",
    "            v2.PILToTensor(),\n",
    "            TargetTransform(),\n",
    "        ])\n",
    "        self.train_ds = OxfordIIITPet(\n",
    "            root=pets_path, \n",
    "            split='trainval', \n",
    "            target_types='segmentation',\n",
    "            # transform=transforms_img, \n",
    "            # target_transform=transforms_target,\n",
    "            transforms=joint_transforms,\n",
    "            download=True\n",
    "        )\n",
    "        self.val_ds = OxfordIIITPet(\n",
    "            root=pets_path, \n",
    "            split='test', \n",
    "            target_types='segmentation',\n",
    "            transform=transforms_img, \n",
    "            target_transform=transforms_target, \n",
    "            download=True)\n",
    "        self.batch_size = 32\n",
    "\n",
    "    def train_dataloader(self):\n",
    "        return DataLoader(\n",
    "            self.train_ds, \n",
    "            batch_size=self.batch_size, \n",
    "            shuffle=True, \n",
    "            num_workers=workers\n",
    "        )\n",
    "\n",
    "    def val_dataloader(self):\n",
    "        return DataLoader(\n",
    "            self.val_ds, \n",
    "            batch_size=self.batch_size, \n",
    "            shuffle=False, \n",
    "            num_workers=workers\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5655b91-cc85-48d0-a5ea-f86b6af25e7b",
   "metadata": {},
   "source": [
    "#### Prepare Checkpoint Callback"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "186a4010-e6a8-48db-b61e-f8a6de7dab3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint_callback = ModelCheckpoint(\n",
    "    monitor='val_loss',\n",
    "    dirpath=unet_path / 'checkpoints/',\n",
    "    filename='unet-{epoch:02d}-{val_loss:.2f}',\n",
    "    save_top_k=1,\n",
    "    mode='min',\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "616ae8a7-8b99-4178-8f22-10a526f9f74f",
   "metadata": {},
   "source": [
    "#### Setup Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "e0cbbcb8",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = UNetResNet34()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "63242cfc-b9ba-42c9-9881-7366a2056a51",
   "metadata": {},
   "outputs": [],
   "source": [
    "lit_model = LitUNet(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "ef1b574b-81b9-44eb-b0ad-bcbb12600fae",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using 16bit Automatic Mixed Precision (AMP)\n",
      "GPU available: True (mps), used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "HPU available: False, using: 0 HPUs\n"
     ]
    }
   ],
   "source": [
    "data_module = PetDataModule()\n",
    "\n",
    "trainer = pl.Trainer(\n",
    "    max_epochs=10, \n",
    "    callbacks=[checkpoint_callback],\n",
    "    accelerator='auto',\n",
    "    devices=1,\n",
    "    precision='16-mixed',\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c27e1f2d-7f05-469c-b837-dde87928c137",
   "metadata": {},
   "source": [
    "#### Train the Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91a0d95b-efb9-4dbb-9b68-ec55d81d900d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  5%|█████████▊                                                                                                                                                                                            | 39.2M/792M [01:25<13:14, 948kB/s]"
     ]
    }
   ],
   "source": [
    "trainer.fit(\n",
    "    lit_model, \n",
    "    datamodule=data_module\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4da3422e-08a2-43cb-945d-d8c7073b1b22",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.save_checkpoint(\"last_checkpoint_aug.ckpt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61dd6330-f679-4e92-bf14-85fb394dea48",
   "metadata": {},
   "source": [
    "#### Visualize Predicted Masks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1ad92ee",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "data_module.setup()\n",
    "x, y = next(iter(data_module.val_dataloader()))\n",
    "preds = lit_model(x).detach().cpu().numpy()\n",
    "plt.figure(figsize=(10,5))\n",
    "for i in range(3):\n",
    "    plt.subplot(3,3,i*3+1)\n",
    "    plt.imshow(x[i].permute(1,2,0).cpu().numpy())\n",
    "    plt.title('Input')\n",
    "    plt.subplot(3,3,i*3+2)\n",
    "    plt.imshow(y[i][0], cmap='gray')\n",
    "    plt.title('Ground Truth')\n",
    "    plt.subplot(3,3,i*3+3)\n",
    "    plt.imshow(preds[i][0]>0.5, cmap='gray')\n",
    "    plt.title('Prediction')\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47b14d90-19e7-404c-912a-ed4431f4622e",
   "metadata": {},
   "source": [
    "Check [torchvision model-hub for UNet model](https://pytorch.org/hub/mateuszbuda_brain-segmentation-pytorch_unet/) source "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50e87420-45f4-465c-b16c-2944f6588f0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "gc.collect()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60d28de5",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Section B: Train YOLO11 for Instance Segmentation\n",
    "\n",
    "In this section we will train a YOLO11 using `ultralytics` models on the different provided datasets for instance segmentation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b1fe6b5-539f-4d49-8ea5-ff8591839632",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = YOLO(\"yolo11n-seg.pt\") "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2329b3d9-d3f7-44f4-9a3f-dadb15c32f19",
   "metadata": {},
   "source": [
    "#### Download YAML Metadata File"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dad46ed3-05f5-4195-a688-afde99ee9aad",
   "metadata": {},
   "outputs": [],
   "source": [
    "def download_yaml(url: str, dest: Path) -> Path:\n",
    "    dest.mkdir(exist_ok=True, parents=True)\n",
    "    file_name = url.split('/')[-1]\n",
    "    response = requests.get(url)\n",
    "    file_path = dest / file_name\n",
    "    if response.status_code == 200:\n",
    "        with file_path.open('wb') as file:\n",
    "            file.write(response.content)\n",
    "        print(f'{file_name} has been downloaded successfully in {file_path}.')\n",
    "    else:\n",
    "        print(f'Failed to download file. Status code: {response.status_code}')\n",
    "\n",
    "    return file_path"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58cb9f79-f49e-45ff-94ca-83ab78576043",
   "metadata": {},
   "source": [
    "#### Prepare Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7891276-a738-4583-90ea-39440d3d1928",
   "metadata": {},
   "outputs": [],
   "source": [
    "url = 'https://github.com/ultralytics/ultralytics/raw/main/ultralytics/cfg/datasets/carparts-seg.yaml'\n",
    "yaml_path = download_yaml(url, carparts_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5518a25a-7941-4f5a-bbde-5c098f54494a",
   "metadata": {},
   "source": [
    "#### Train the Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0476080-81d0-4905-8df3-82d546f9edcb",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "results = model.train(\n",
    "    data=yaml_path.name, \n",
    "    epochs=100, \n",
    "    imgsz=640,\n",
    "    device=device,\n",
    "    workers=11,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06325ed2-52c8-4a2d-abe1-fd90a027c80d",
   "metadata": {},
   "source": [
    "#### Visualize Prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a811fe59-99c3-4a09-a5f2-e2d0a237fd86",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Predict with the model\n",
    "results = model(source=\"https://ultralytics.com/images/bus.jpg\", show=True)  # predict on an image\n",
    "\n",
    "# Access the results\n",
    "# for result in results:\n",
    "#     xy = result.masks.xy  # mask in polygon format\n",
    "#     xyn = result.masks.xyn  # normalized\n",
    "#     masks = result.masks.data  # mask in matrix format (num_objects x H x W)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69761033-dca8-4970-b0a2-616986c8075d",
   "metadata": {},
   "outputs": [],
   "source": [
    "img = result.plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "415b3745-d830-4921-ada2-c11f22fa110e",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(img)\n",
    "plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a94f129f-326d-47a9-a1d8-90ff52eea2fe",
   "metadata": {},
   "source": [
    "#### Train on Crack Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2645ea3b-fe4d-4fb1-b2ab-4073e2fbfc08",
   "metadata": {},
   "outputs": [],
   "source": [
    "url = 'https://github.com/ultralytics/ultralytics/blob/main/ultralytics/cfg/datasets/crack-seg.yaml'\n",
    "yaml_path = download_yaml(url, cracks_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d7464e0-eb3f-40cb-8851-0af583c858c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = model.train(\n",
    "    data=yaml_path.name, \n",
    "    epochs=100, \n",
    "    imgsz=640,\n",
    "    device=device,\n",
    "    workers=11,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ddeff365-2f82-4a8c-9a79-ad9bb535ea85",
   "metadata": {},
   "source": [
    "#### Train on Packages Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbb7a92e-d0b7-4df6-842f-e21b6e2c5374",
   "metadata": {},
   "outputs": [],
   "source": [
    "url = 'https://github.com/ultralytics/ultralytics/blob/main/ultralytics/cfg/datasets/package-seg.yaml'\n",
    "yaml_path = download_yaml(url, packages_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "946a3306-e887-4f05-93c1-1cfe9e9380b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = model.train(\n",
    "    data=yaml_path.name, \n",
    "    epochs=100, \n",
    "    imgsz=640,\n",
    "    device=device,\n",
    "    workers=11,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ab27be5-cc7b-424e-a03a-0196c332028c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "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.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}