wehub-resource-sync
593b94c120
pytest / Unit Tests (push) Has been cancelled
pytest / Integration (integration_tests_a) (push) Has been cancelled
pytest / Integration (integration_tests_b) (push) Has been cancelled
pytest / Integration (integration_tests_c) (push) Has been cancelled
pytest / Integration (integration_tests_d) (push) Has been cancelled
pytest / Integration (integration_tests_e) (push) Has been cancelled
pytest / Integration (integration_tests_f) (push) Has been cancelled
pytest / Integration (integration_tests_g) (push) Has been cancelled
pytest / Integration (integration_tests_h) (push) Has been cancelled
pytest / Integration (integration_tests_i) (push) Has been cancelled
pytest / Integration (integration_tests_j) (push) Has been cancelled
pytest / Distributed (distributed_a) (push) Has been cancelled
pytest / Distributed (distributed_b) (push) Has been cancelled
pytest / Distributed (distributed_c) (push) Has been cancelled
pytest / Distributed (distributed_d) (push) Has been cancelled
pytest / Distributed (distributed_e) (push) Has been cancelled
pytest / Distributed (distributed_f) (push) Has been cancelled
pytest / Minimal Install (push) Has been cancelled
pytest / Event File (push) Has been cancelled
pytest (slow) / py-slow (push) Has been cancelled
Publish JSON Schema / publish-schema (push) Has been cancelled
579 lines
19 KiB
Plaintext
579 lines
19 KiB
Plaintext
{
|
||
"cells": [
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"# Semantic Segmentation with Ludwig\n",
|
||
"\n",
|
||
"[](https://colab.research.google.com/github/ludwig-ai/ludwig/blob/main/examples/semantic_segmentation/semantic_segmentation.ipynb)\n",
|
||
"\n",
|
||
"This notebook demonstrates semantic segmentation on the **CamSeq01** driving dataset using three different image decoders available in Ludwig:\n",
|
||
"\n",
|
||
"| Decoder | Architecture | Encoder pairing | Best for |\n",
|
||
"|---------|-------------|-----------------|----------|\n",
|
||
"| `unet` | Symmetric encoder-decoder with skip connections | Built-in UNet encoder | General purpose, well-studied baseline |\n",
|
||
"| `segformer` | Lightweight MLP head on top of multi-scale features | `dinov2` / any ViT backbone | High accuracy, transformer features |\n",
|
||
"| `fpn` | Feature Pyramid Network top-down pathway | Any CNN (e.g. `efficientnet`) | Fast inference, good at multi-scale objects |\n",
|
||
"\n",
|
||
"**GPU required** — an A100 or similar is recommended for the SegFormer run.\n",
|
||
"\n",
|
||
"CamSeq01 contains 101 urban driving images at 960×720 with 32 semantic classes."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"## Setup"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"!pip install 'ludwig[vision]' --quiet"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"import logging\n",
|
||
"import time\n",
|
||
"\n",
|
||
"import matplotlib.pyplot as plt\n",
|
||
"import numpy as np\n",
|
||
"import pandas as pd\n",
|
||
"import torch\n",
|
||
"\n",
|
||
"from ludwig.api import LudwigModel"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"## Dataset\n",
|
||
"\n",
|
||
"Ludwig ships a built-in downloader for CamSeq01. \n",
|
||
"The call below downloads and caches the dataset, then returns a `pd.DataFrame`\n",
|
||
"with two columns: `image_path` and `mask_path`."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"from ludwig.datasets import camseq\n",
|
||
"\n",
|
||
"df = camseq.load(split=False)\n",
|
||
"print(f\"Total samples: {len(df)}\")\n",
|
||
"df.head()"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"# Reserve first image for visual comparison at the end\n",
|
||
"pred_set = df[0:1].reset_index(drop=True)\n",
|
||
"train_set = df[1:]\n",
|
||
"\n",
|
||
"print(f\"Training samples : {len(train_set)}\")\n",
|
||
"print(f\"Prediction sample: {len(pred_set)}\")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"# Quick look at the raw image and ground-truth mask\n",
|
||
"from PIL import Image as PILImage\n",
|
||
"\n",
|
||
"fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n",
|
||
"axes[0].imshow(PILImage.open(pred_set[\"image_path\"][0]))\n",
|
||
"axes[0].set_title(\"Input image\")\n",
|
||
"axes[0].axis(\"off\")\n",
|
||
"axes[1].imshow(PILImage.open(pred_set[\"mask_path\"][0]))\n",
|
||
"axes[1].set_title(\"Ground-truth mask (32 classes)\")\n",
|
||
"axes[1].axis(\"off\")\n",
|
||
"plt.tight_layout()\n",
|
||
"plt.show()"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"## Option 1: UNet (default)\n",
|
||
"\n",
|
||
"The classic UNet uses a symmetric encoder-decoder with skip connections. \n",
|
||
"Ludwig's implementation supports a configurable number of stages (`num_stages`)\n",
|
||
"so you can trade off capacity against speed."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"unet_config = {\n",
|
||
" \"input_features\": [\n",
|
||
" {\n",
|
||
" \"name\": \"image_path\",\n",
|
||
" \"type\": \"image\",\n",
|
||
" \"preprocessing\": {\"num_processes\": 4, \"height\": 512, \"width\": 512},\n",
|
||
" \"encoder\": {\"type\": \"unet\"},\n",
|
||
" }\n",
|
||
" ],\n",
|
||
" \"output_features\": [\n",
|
||
" {\n",
|
||
" \"name\": \"mask_path\",\n",
|
||
" \"type\": \"image\",\n",
|
||
" \"preprocessing\": {\n",
|
||
" \"num_processes\": 4,\n",
|
||
" \"height\": 512,\n",
|
||
" \"width\": 512,\n",
|
||
" \"num_classes\": 32,\n",
|
||
" },\n",
|
||
" \"decoder\": {\n",
|
||
" \"type\": \"unet\",\n",
|
||
" \"num_stages\": 4, # configurable depth\n",
|
||
" \"num_fc_layers\": 0,\n",
|
||
" \"conv_norm\": \"batch\",\n",
|
||
" },\n",
|
||
" \"loss\": {\"type\": \"softmax_cross_entropy\"},\n",
|
||
" }\n",
|
||
" ],\n",
|
||
" \"combiner\": {\"type\": \"concat\", \"num_fc_layers\": 0},\n",
|
||
" \"trainer\": {\n",
|
||
" \"epochs\": 50,\n",
|
||
" \"early_stop\": 10,\n",
|
||
" \"batch_size\": 4,\n",
|
||
" \"learning_rate\": 0.0001,\n",
|
||
" },\n",
|
||
"}\n",
|
||
"\n",
|
||
"t0 = time.time()\n",
|
||
"unet_model = LudwigModel(unet_config, logging_level=logging.WARNING)\n",
|
||
"unet_stats, _, unet_output_dir = unet_model.train(\n",
|
||
" dataset=train_set,\n",
|
||
" experiment_name=\"seg_comparison\",\n",
|
||
" model_name=\"unet\",\n",
|
||
" skip_save_processed_input=True,\n",
|
||
")\n",
|
||
"unet_time = time.time() - t0\n",
|
||
"print(f\"UNet training time: {unet_time:.1f}s\")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"unet_preds, _ = unet_model.predict(pred_set)\n",
|
||
"if not isinstance(unet_preds, pd.DataFrame):\n",
|
||
" unet_preds = unet_preds.compute()\n",
|
||
"unet_pred_mask = torch.from_numpy(unet_preds[\"mask_path_predictions\"].iloc[0])\n",
|
||
"print(\"UNet prediction mask shape:\", unet_pred_mask.shape)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"## Option 2: SegFormer (transformer backbone)\n",
|
||
"\n",
|
||
"SegFormer pairs a DINOv2 vision-transformer backbone with a lightweight MLP\n",
|
||
"decoder head. The hierarchical features from DINOv2 feed the SegFormer head\n",
|
||
"directly — no upsampling convolutions needed until the final prediction.\n",
|
||
"\n",
|
||
"- Encoder: `dinov2` (`facebook/dinov2-base`, ~86M params pretrained)\n",
|
||
"- Decoder: `segformer` (hidden MLP projection → bilinear upsample)\n",
|
||
"- Best suited to: high-accuracy segmentation where GPU memory is not the bottleneck"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"segformer_config = {\n",
|
||
" \"input_features\": [\n",
|
||
" {\n",
|
||
" \"name\": \"image_path\",\n",
|
||
" \"type\": \"image\",\n",
|
||
" \"preprocessing\": {\"num_processes\": 4, \"height\": 512, \"width\": 512},\n",
|
||
" \"encoder\": {\n",
|
||
" \"type\": \"dinov2\",\n",
|
||
" \"pretrained_model_name_or_path\": \"facebook/dinov2-base\",\n",
|
||
" \"use_pretrained\": True,\n",
|
||
" \"trainable\": True,\n",
|
||
" },\n",
|
||
" }\n",
|
||
" ],\n",
|
||
" \"output_features\": [\n",
|
||
" {\n",
|
||
" \"name\": \"mask_path\",\n",
|
||
" \"type\": \"image\",\n",
|
||
" \"preprocessing\": {\n",
|
||
" \"num_processes\": 4,\n",
|
||
" \"height\": 512,\n",
|
||
" \"width\": 512,\n",
|
||
" \"num_classes\": 32,\n",
|
||
" },\n",
|
||
" \"decoder\": {\n",
|
||
" \"type\": \"segformer\",\n",
|
||
" \"hidden_size\": 256,\n",
|
||
" \"dropout\": 0.1,\n",
|
||
" },\n",
|
||
" \"loss\": {\"type\": \"softmax_cross_entropy\"},\n",
|
||
" }\n",
|
||
" ],\n",
|
||
" \"combiner\": {\"type\": \"concat\", \"num_fc_layers\": 0},\n",
|
||
" \"trainer\": {\n",
|
||
" \"epochs\": 50,\n",
|
||
" \"early_stop\": 10,\n",
|
||
" \"batch_size\": 4,\n",
|
||
" \"learning_rate\": 0.0001,\n",
|
||
" },\n",
|
||
"}\n",
|
||
"\n",
|
||
"t0 = time.time()\n",
|
||
"segformer_model = LudwigModel(segformer_config, logging_level=logging.WARNING)\n",
|
||
"segformer_stats, _, segformer_output_dir = segformer_model.train(\n",
|
||
" dataset=train_set,\n",
|
||
" experiment_name=\"seg_comparison\",\n",
|
||
" model_name=\"segformer\",\n",
|
||
" skip_save_processed_input=True,\n",
|
||
")\n",
|
||
"segformer_time = time.time() - t0\n",
|
||
"print(f\"SegFormer training time: {segformer_time:.1f}s\")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"segformer_preds, _ = segformer_model.predict(pred_set)\n",
|
||
"if not isinstance(segformer_preds, pd.DataFrame):\n",
|
||
" segformer_preds = segformer_preds.compute()\n",
|
||
"segformer_pred_mask = torch.from_numpy(segformer_preds[\"mask_path_predictions\"].iloc[0])\n",
|
||
"print(\"SegFormer prediction mask shape:\", segformer_pred_mask.shape)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"## Option 3: FPN (lightweight)\n",
|
||
"\n",
|
||
"The Feature Pyramid Network decoder builds a top-down pathway over multi-scale\n",
|
||
"feature maps produced by any CNN backbone. Combined with EfficientNet it gives\n",
|
||
"a good accuracy/speed balance and fits comfortably in smaller GPU budgets.\n",
|
||
"\n",
|
||
"- Encoder: `efficientnet` (pretrained on ImageNet, ~5M params)\n",
|
||
"- Decoder: `fpn` (lateral projections + top-down merging over 4 pyramid levels)\n",
|
||
"- Best suited to: production deployments where inference latency matters"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"fpn_config = {\n",
|
||
" \"input_features\": [\n",
|
||
" {\n",
|
||
" \"name\": \"image_path\",\n",
|
||
" \"type\": \"image\",\n",
|
||
" \"preprocessing\": {\"num_processes\": 4, \"height\": 512, \"width\": 512},\n",
|
||
" \"encoder\": {\n",
|
||
" \"type\": \"efficientnet\",\n",
|
||
" \"use_pretrained\": True,\n",
|
||
" \"trainable\": True,\n",
|
||
" },\n",
|
||
" }\n",
|
||
" ],\n",
|
||
" \"output_features\": [\n",
|
||
" {\n",
|
||
" \"name\": \"mask_path\",\n",
|
||
" \"type\": \"image\",\n",
|
||
" \"preprocessing\": {\n",
|
||
" \"num_processes\": 4,\n",
|
||
" \"height\": 512,\n",
|
||
" \"width\": 512,\n",
|
||
" \"num_classes\": 32,\n",
|
||
" },\n",
|
||
" \"decoder\": {\n",
|
||
" \"type\": \"fpn\",\n",
|
||
" \"num_channels\": 256,\n",
|
||
" \"num_levels\": 4,\n",
|
||
" },\n",
|
||
" \"loss\": {\"type\": \"softmax_cross_entropy\"},\n",
|
||
" }\n",
|
||
" ],\n",
|
||
" \"combiner\": {\"type\": \"concat\", \"num_fc_layers\": 0},\n",
|
||
" \"trainer\": {\n",
|
||
" \"epochs\": 100,\n",
|
||
" \"early_stop\": 10,\n",
|
||
" \"batch_size\": 8,\n",
|
||
" \"learning_rate\": 0.0001,\n",
|
||
" },\n",
|
||
"}\n",
|
||
"\n",
|
||
"t0 = time.time()\n",
|
||
"fpn_model = LudwigModel(fpn_config, logging_level=logging.WARNING)\n",
|
||
"fpn_stats, _, fpn_output_dir = fpn_model.train(\n",
|
||
" dataset=train_set,\n",
|
||
" experiment_name=\"seg_comparison\",\n",
|
||
" model_name=\"fpn\",\n",
|
||
" skip_save_processed_input=True,\n",
|
||
")\n",
|
||
"fpn_time = time.time() - t0\n",
|
||
"print(f\"FPN training time: {fpn_time:.1f}s\")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"fpn_preds, _ = fpn_model.predict(pred_set)\n",
|
||
"if not isinstance(fpn_preds, pd.DataFrame):\n",
|
||
" fpn_preds = fpn_preds.compute()\n",
|
||
"fpn_pred_mask = torch.from_numpy(fpn_preds[\"mask_path_predictions\"].iloc[0])\n",
|
||
"print(\"FPN prediction mask shape:\", fpn_pred_mask.shape)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"## UNet depth ablation\n",
|
||
"\n",
|
||
"Ludwig's UNet decoder exposes a `num_stages` parameter that controls how many\n",
|
||
"encoder/decoder stage pairs are stacked. More stages → richer multi-scale\n",
|
||
"representations but also more parameters and longer training times. The input\n",
|
||
"spatial dimensions must be divisible by `2^num_stages`.\n",
|
||
"\n",
|
||
"Run the standalone sweep script to get the full table:\n",
|
||
"\n",
|
||
"```bash\n",
|
||
"python unet_depth_sweep.py\n",
|
||
"```\n",
|
||
"\n",
|
||
"Below we replicate a mini version for illustration."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"import yaml\n",
|
||
"\n",
|
||
"SWEEP_BASE = {\n",
|
||
" \"input_features\": [\n",
|
||
" {\n",
|
||
" \"name\": \"image_path\",\n",
|
||
" \"type\": \"image\",\n",
|
||
" \"preprocessing\": {\"num_processes\": 4, \"height\": 512, \"width\": 512},\n",
|
||
" \"encoder\": {\"type\": \"unet\"},\n",
|
||
" }\n",
|
||
" ],\n",
|
||
" \"output_features\": [\n",
|
||
" {\n",
|
||
" \"name\": \"mask_path\",\n",
|
||
" \"type\": \"image\",\n",
|
||
" \"preprocessing\": {\"num_processes\": 4, \"height\": 512, \"width\": 512, \"num_classes\": 32},\n",
|
||
" \"decoder\": {\"type\": \"unet\", \"num_fc_layers\": 0, \"conv_norm\": \"batch\"},\n",
|
||
" \"loss\": {\"type\": \"softmax_cross_entropy\"},\n",
|
||
" }\n",
|
||
" ],\n",
|
||
" \"combiner\": {\"type\": \"concat\", \"num_fc_layers\": 0},\n",
|
||
" \"trainer\": {\"epochs\": 20, \"early_stop\": 5, \"batch_size\": 4, \"learning_rate\": 0.0001},\n",
|
||
"}\n",
|
||
"\n",
|
||
"sweep_results = []\n",
|
||
"\n",
|
||
"for depth in [2, 3, 4, 5]:\n",
|
||
" cfg = yaml.safe_load(yaml.dump(SWEEP_BASE))\n",
|
||
" cfg[\"output_features\"][0][\"decoder\"][\"num_stages\"] = depth\n",
|
||
"\n",
|
||
" m = LudwigModel(cfg, logging_level=logging.WARNING)\n",
|
||
" t0 = time.time()\n",
|
||
" stats, _, _ = m.train(\n",
|
||
" dataset=train_set,\n",
|
||
" experiment_name=\"depth_sweep\",\n",
|
||
" model_name=f\"unet_depth_{depth}\",\n",
|
||
" skip_save_processed_input=True,\n",
|
||
" )\n",
|
||
" elapsed = time.time() - t0\n",
|
||
"\n",
|
||
" n_params = sum(p.numel() for p in m.model.parameters() if p.requires_grad)\n",
|
||
"\n",
|
||
" val_loss = None\n",
|
||
" try:\n",
|
||
" val_loss = min(stats[\"validation\"][\"combined\"][\"loss\"])\n",
|
||
" except (KeyError, TypeError):\n",
|
||
" pass\n",
|
||
"\n",
|
||
" sweep_results.append(\n",
|
||
" {\n",
|
||
" \"num_stages\": depth,\n",
|
||
" \"trainable_params\": f\"{n_params:,}\",\n",
|
||
" \"best_val_loss\": round(val_loss, 4) if val_loss is not None else \"n/a\",\n",
|
||
" \"training_time_s\": round(elapsed, 1),\n",
|
||
" }\n",
|
||
" )\n",
|
||
" print(f\"depth={depth} params={n_params:,} val_loss={val_loss} time={elapsed:.1f}s\")\n",
|
||
"\n",
|
||
"sweep_df = pd.DataFrame(sweep_results)\n",
|
||
"sweep_df"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"fig, ax1 = plt.subplots(figsize=(8, 4))\n",
|
||
"depths = [r[\"num_stages\"] for r in sweep_results]\n",
|
||
"times = [r[\"training_time_s\"] for r in sweep_results]\n",
|
||
"\n",
|
||
"ax1.bar(depths, times, color=\"steelblue\", alpha=0.7, label=\"Training time (s)\")\n",
|
||
"ax1.set_xlabel(\"UNet num_stages\")\n",
|
||
"ax1.set_ylabel(\"Training time (s)\", color=\"steelblue\")\n",
|
||
"ax1.tick_params(axis=\"y\", labelcolor=\"steelblue\")\n",
|
||
"ax1.set_xticks(depths)\n",
|
||
"\n",
|
||
"val_losses = [r[\"best_val_loss\"] for r in sweep_results if isinstance(r[\"best_val_loss\"], float)]\n",
|
||
"if len(val_losses) == len(depths):\n",
|
||
" ax2 = ax1.twinx()\n",
|
||
" ax2.plot(depths, val_losses, \"o-\", color=\"tomato\", label=\"Best val loss\")\n",
|
||
" ax2.set_ylabel(\"Best val loss\", color=\"tomato\")\n",
|
||
" ax2.tick_params(axis=\"y\", labelcolor=\"tomato\")\n",
|
||
"\n",
|
||
"plt.title(\"UNet depth: training time vs validation loss\")\n",
|
||
"plt.tight_layout()\n",
|
||
"plt.show()"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"## Visualisation\n",
|
||
"\n",
|
||
"Compare the predicted segmentation maps from the three models side by side."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"def to_rgb(mask_tensor: torch.Tensor) -> np.ndarray:\n",
|
||
" \"\"\"Convert a CHW or HW class-index tensor to a displayable RGB image.\"\"\"\n",
|
||
" t = mask_tensor\n",
|
||
" if t.dim() == 3:\n",
|
||
" # (C, H, W) logits or class probabilities -> argmax\n",
|
||
" t = t.argmax(dim=0)\n",
|
||
" arr = t.cpu().numpy().astype(np.float32)\n",
|
||
" arr = (arr - arr.min()) / max(arr.max() - arr.min(), 1e-5)\n",
|
||
" return arr\n",
|
||
"\n",
|
||
"\n",
|
||
"input_img = PILImage.open(pred_set[\"image_path\"][0])\n",
|
||
"gt_mask = PILImage.open(pred_set[\"mask_path\"][0])\n",
|
||
"\n",
|
||
"fig, axes = plt.subplots(1, 5, figsize=(22, 5))\n",
|
||
"\n",
|
||
"axes[0].imshow(input_img)\n",
|
||
"axes[0].set_title(\"Input image\")\n",
|
||
"\n",
|
||
"axes[1].imshow(gt_mask)\n",
|
||
"axes[1].set_title(\"Ground truth\")\n",
|
||
"\n",
|
||
"axes[2].imshow(to_rgb(unet_pred_mask), cmap=\"tab20\")\n",
|
||
"axes[2].set_title(\"UNet (depth=4)\")\n",
|
||
"\n",
|
||
"axes[3].imshow(to_rgb(segformer_pred_mask), cmap=\"tab20\")\n",
|
||
"axes[3].set_title(\"SegFormer + DINOv2\")\n",
|
||
"\n",
|
||
"axes[4].imshow(to_rgb(fpn_pred_mask), cmap=\"tab20\")\n",
|
||
"axes[4].set_title(\"FPN + EfficientNet\")\n",
|
||
"\n",
|
||
"for ax in axes:\n",
|
||
" ax.axis(\"off\")\n",
|
||
"\n",
|
||
"plt.suptitle(\"Segmentation map comparison\", fontsize=14)\n",
|
||
"plt.tight_layout()\n",
|
||
"plt.show()"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"# Print a summary comparison table\n",
|
||
"comparison = pd.DataFrame(\n",
|
||
" [\n",
|
||
" {\"model\": \"UNet (num_stages=4)\", \"encoder\": \"unet\", \"decoder\": \"unet\", \"training_time_s\": round(unet_time, 1)},\n",
|
||
" {\n",
|
||
" \"model\": \"SegFormer + DINOv2\",\n",
|
||
" \"encoder\": \"dinov2\",\n",
|
||
" \"decoder\": \"segformer\",\n",
|
||
" \"training_time_s\": round(segformer_time, 1),\n",
|
||
" },\n",
|
||
" {\n",
|
||
" \"model\": \"FPN + EfficientNet\",\n",
|
||
" \"encoder\": \"efficientnet\",\n",
|
||
" \"decoder\": \"fpn\",\n",
|
||
" \"training_time_s\": round(fpn_time, 1),\n",
|
||
" },\n",
|
||
" ]\n",
|
||
")\n",
|
||
"comparison"
|
||
]
|
||
}
|
||
],
|
||
"metadata": {
|
||
"kernelspec": {
|
||
"display_name": "Python 3",
|
||
"language": "python",
|
||
"name": "python3"
|
||
},
|
||
"language_info": {
|
||
"name": "python",
|
||
"version": "3.12.0"
|
||
}
|
||
},
|
||
"nbformat": 4,
|
||
"nbformat_minor": 4
|
||
}
|