finetune-maskformer-pytorch.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import evaluate\n",
    "import json\n",
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "from datasets import load_dataset\n",
    "from huggingface_hub import hf_hub_download\n",
    "from os import path\n",
    "from PIL import Image as PImage\n",
    "from tqdm.auto import tqdm\n",
    "from torch import nn\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "from torchvision import transforms as T\n",
    "from transformers import MaskFormerForInstanceSegmentation, MaskFormerImageProcessor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ade_mean=[0.485, 0.456, 0.406]\n",
    "ade_std=[0.229, 0.224, 0.225]\n",
    "\n",
    "palette = [\n",
    "    [120, 120, 120], [4, 200, 4], [4, 4, 250], [6, 230, 230],\n",
    "    [80, 50, 50], [120, 120, 80], [140, 140, 140], [204, 5, 255]\n",
    "]\n",
    "\n",
    "\n",
    "def np_from_tensor(img_t, mean=[0.,0.,0.], std=[1.,1.,1.]):\n",
    "    img_t = (img_t * np.array(std)[:, None, None]) + np.array(mean)[:, None, None]\n",
    "    return np.moveaxis((255 * img_t).numpy().astype(np.uint8), 0, -1)\n",
    "\n",
    "\n",
    "def add_segmentations_to_image(img, segs):\n",
    "    color_segmentation_map = np.zeros((segs.shape[0], segs.shape[1], 3), dtype=np.uint8)\n",
    "    for label, color in enumerate(palette):\n",
    "        color_segmentation_map[segs == label, :] = color\n",
    "    img_mask = np.array(img) * 0.5 + color_segmentation_map * 0.5\n",
    "    return img_mask.astype(np.uint8)\n",
    "\n",
    "\n",
    "def mask_from_label(masks, labels, label_name):\n",
    "  print(\"Label:\", label_name)\n",
    "  idx = labels.index(label_name)\n",
    "\n",
    "  visual_mask = (masks[idx].bool().numpy() * 255).astype(np.uint8)\n",
    "  return visual_mask\n",
    "\n",
    "\n",
    "def add_mask_label_to_image(img, mask_label, label_idx):\n",
    "    img_mask_label = np.zeros((mask_label.shape[0], mask_label.shape[1], 3), dtype=np.uint8)\n",
    "    img_mask_label[mask_label == 255, :] = palette[label_idx]\n",
    "    img_mask_label = 0.5 * img + 0.5 * img_mask_label\n",
    "    return img_mask_label.astype(np.uint8)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_id = f\"thiagohersan/satellite-trees\"\n",
    "base_model_id = f\"facebook/maskformer-swin-base-ade\"\n",
    "result_model_id = f\"maskformer-satellite-trees\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "id2label = json.load(open(hf_hub_download(dataset_id, \"id2label.json\", repo_type=\"dataset\"), \"r\"))\n",
    "id2label = {int(k):v for k,v in id2label.items()}\n",
    "label2id = {v:int(k) for k,v in id2label.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = MaskFormerForInstanceSegmentation.from_pretrained(\n",
    "    base_model_id,\n",
    "    id2label=id2label,\n",
    "    label2id=label2id,\n",
    "    ignore_mismatched_sizes=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Download and Check Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = load_dataset(dataset_id)\n",
    "dataset = dataset[\"train\"].train_test_split(test_size=0.2, shuffle=True, seed=1010)\n",
    "\n",
    "train_ds = dataset[\"train\"]\n",
    "test_ds = dataset[\"test\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create PyTorch Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_transform = T.Compose([\n",
    "    T.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.25),\n",
    "    T.RandomPosterize(bits=2, p=0.2),\n",
    "    T.RandomAdjustSharpness(sharpness_factor=3, p=0.2),\n",
    "    T.RandomAutocontrast(p=0.3),\n",
    "    T.RandomEqualize(p=0.3),\n",
    "    T.ToTensor(),\n",
    "    T.Normalize(mean=ade_mean, std=ade_std)\n",
    "])\n",
    "\n",
    "test_transform = T.Compose([\n",
    "    T.ToTensor(),\n",
    "    T.Normalize(mean=ade_mean, std=ade_std)\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ImageSegmentationDataset(Dataset):\n",
    "    def __init__(self, dataset, transform):\n",
    "        self.dataset = dataset\n",
    "        self.transform = transform\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.dataset)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        original_image = np.array(self.dataset[idx]['pixel_values'])\n",
    "        original_segmentation_map = np.array(self.dataset[idx]['label'])\n",
    "\n",
    "        image = self.transform(self.dataset[idx]['pixel_values'])\n",
    "        segmentation_map = original_segmentation_map\n",
    "\n",
    "        return image, segmentation_map, original_image, original_segmentation_map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = ImageSegmentationDataset(train_ds, transform=train_transform)\n",
    "test_dataset = ImageSegmentationDataset(test_ds, transform=test_transform)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create PyTorch DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "preprocessor = MaskFormerImageProcessor(\n",
    "    do_resize=False,\n",
    "    do_normalize=False,\n",
    "    do_rescale=False,\n",
    "    ignore_index=255,\n",
    "    reduce_labels=False\n",
    ")\n",
    "\n",
    "def collate_fn(batch):\n",
    "    inputs = list(zip(*batch))\n",
    "    images = inputs[0]\n",
    "    segmentation_maps = inputs[1]\n",
    "\n",
    "    batch = preprocessor(\n",
    "        images=images,\n",
    "        segmentation_maps=segmentation_maps,\n",
    "        return_tensors=\"pt\",\n",
    "    )\n",
    "\n",
    "    batch[\"original_images\"] = list(inputs[2])\n",
    "    batch[\"original_segmentation_maps\"] = list(inputs[3])\n",
    "\n",
    "    return batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataloader = DataLoader(train_dataset, batch_size=4, shuffle=True, collate_fn=collate_fn)\n",
    "test_dataloader = DataLoader(test_dataset, batch_size=4, shuffle=False, collate_fn=collate_fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch = next(iter(train_dataloader))\n",
    "b_idx = 0\n",
    "\n",
    "for k,v in batch.items():\n",
    "  try:\n",
    "    print(k,v.shape, v.dtype)\n",
    "  except:\n",
    "    print(k,v[b_idx].shape, v[b_idx].dtype)\n",
    "  \n",
    "b_labels = [id2label[label] for label in batch[\"class_labels\"][b_idx].tolist()]\n",
    "\n",
    "print(b_labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PImage.fromarray(np_from_tensor(batch['pixel_values'][b_idx]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PImage.fromarray(mask_from_label(batch[\"mask_labels\"][b_idx], b_labels, 'vegetation'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PImage.fromarray(\n",
    "    add_mask_label_to_image(\n",
    "        np_from_tensor(batch['pixel_values'][b_idx]),\n",
    "        mask_from_label(batch[\"mask_labels\"][b_idx], b_labels, 'water'),\n",
    "        b_labels.index('water')\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PImage.fromarray(batch['original_images'][b_idx])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PImage.fromarray(\n",
    "    add_segmentations_to_image(\n",
    "        batch['original_images'][b_idx],\n",
    "        batch[\"original_segmentation_maps\"][b_idx]\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "outputs = model(\n",
    "    batch[\"pixel_values\"],\n",
    "    class_labels=batch[\"class_labels\"],\n",
    "    mask_labels=batch[\"mask_labels\"]\n",
    ")\n",
    "\n",
    "outputs.loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metric = evaluate.load(\"mean_iou\")\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "model.to(device)\n",
    "\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=5e-5)\n",
    "\n",
    "running_loss = 0.0\n",
    "num_samples = 0\n",
    "\n",
    "for epoch in range(16):\n",
    "    print(\"Epoch:\", epoch)\n",
    "\n",
    "    model.train()\n",
    "    for idx, batch in enumerate(tqdm(train_dataloader)):\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        outputs = model(\n",
    "            pixel_values=batch[\"pixel_values\"].to(device),\n",
    "            mask_labels=[labels.to(device) for labels in batch[\"mask_labels\"]],\n",
    "            class_labels=[labels.to(device) for labels in batch[\"class_labels\"]],\n",
    "        )\n",
    "\n",
    "        loss = outputs.loss\n",
    "        loss.backward()\n",
    "\n",
    "        batch_size = batch[\"pixel_values\"].size(0)\n",
    "        running_loss += loss.item()\n",
    "        num_samples += batch_size\n",
    "\n",
    "        if idx % 10 == 0:\n",
    "            print(\"Loss: \", running_loss/num_samples)\n",
    "\n",
    "        optimizer.step()\n",
    "\n",
    "    model.eval()\n",
    "    for idx, batch in enumerate(tqdm(test_dataloader)):\n",
    "        if idx > 7:\n",
    "            break\n",
    "\n",
    "        with torch.no_grad():\n",
    "            outputs = model(pixel_values=batch[\"pixel_values\"].to(device))\n",
    "\n",
    "        original_images = batch[\"original_images\"]\n",
    "        target_sizes = [(image.shape[0], image.shape[1]) for image in original_images]\n",
    "\n",
    "        predicted_segmentation_maps = preprocessor.post_process_semantic_segmentation(\n",
    "            outputs,\n",
    "            target_sizes=target_sizes\n",
    "        )\n",
    "\n",
    "        predicted_segmentation_maps = [psm.cpu() for psm in predicted_segmentation_maps]\n",
    "        ground_truth_segmentation_maps = batch[\"original_segmentation_maps\"]\n",
    "\n",
    "        metric.add_batch(references=ground_truth_segmentation_maps, predictions=predicted_segmentation_maps)\n",
    "\n",
    "    test_metrics = metric.compute(num_labels=len(id2label), ignore_index=255, reduce_labels=False)\n",
    "    print(\"Mean IoU:\", test_metrics['mean_iou'], \"Vegetation IoU:\", test_metrics['per_category_iou'][label2id['vegetation']])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Push Model to Hub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save_pretrained(path.join(\"models\", result_model_id))\n",
    "preprocessor.save_pretrained(path.join(\"models\", result_model_id))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hub_preprocessor = MaskFormerImageProcessor(\n",
    "    do_resize=False,\n",
    "    do_normalize=False,\n",
    "    do_rescale=True,\n",
    "    ignore_index=255,\n",
    "    reduce_labels=False\n",
    ")\n",
    "\n",
    "model.push_to_hub(result_model_id)\n",
    "hub_preprocessor.push_to_hub(result_model_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch = next(iter(test_dataloader))\n",
    "b_idx = 0\n",
    "\n",
    "for k,v in batch.items():\n",
    "  try:\n",
    "    print(k,v.shape, v.dtype)\n",
    "  except:\n",
    "    print(k,v[b_idx].shape, v[b_idx].dtype)\n",
    "  \n",
    "b_labels = [id2label[label] for label in batch[\"class_labels\"][b_idx].tolist()]\n",
    "\n",
    "print(b_labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "  outputs = model(batch[\"pixel_values\"].to(device))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "original_images = batch[\"original_images\"]\n",
    "target_sizes = [(image.shape[0], image.shape[1]) for image in original_images]\n",
    "predicted_segmentation_maps = preprocessor.post_process_semantic_segmentation(outputs, target_sizes=target_sizes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PImage.fromarray(batch[\"original_images\"][b_idx])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PImage.fromarray(\n",
    "    add_segmentations_to_image(\n",
    "        batch[\"original_images\"][b_idx],\n",
    "        batch[\"original_segmentation_maps\"][b_idx]\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PImage.fromarray(\n",
    "    add_segmentations_to_image(\n",
    "        batch[\"original_images\"][b_idx],\n",
    "        predicted_segmentation_maps[b_idx].cpu().numpy()\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.15 ('hf-model')",
   "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.8.15"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "89e384cab7c47fb35ec95d2248b519cf922ee174880eed636c26cdfb6c4df768"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}