Upload layout-fine-tune.ipynb

layout-fine-tune.ipynb

@@ -0,0 +1,187 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Loading Packages"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.optim as optim\n",
+    "from torch.utils.data import DataLoader\n",
+    "# from transformers import SegformerConfig\n",
+    "# from surya.model.detection.segformer import SegformerForRegressionMask\n",
+    "from surya.input.processing import prepare_image_detection\n",
+    "from surya.model.detection.segformer import load_processor , load_model\n",
+    "from datasets import load_dataset\n",
+    "from tqdm import tqdm\n",
+    "from torch.utils.tensorboard import SummaryWriter\n",
+    "import torch.nn.functional as F\n",
+    "import numpy as np \n",
+    "from surya.layout import parallel_get_regions"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Initializing The Dataset And Model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+    "dataset = load_dataset(\"vikp/publaynet_bench\", split=\"train[:100]\") # You can choose you own dataset\n",
+    "model = load_model(\"vikp/surya_layout2\") "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Helper Functions, Loss Function And Optimizer"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "optimizer = optim.Adam(model.parameters(), lr=0.00001)\n",
+    "log_dir = \"logs\"\n",
+    "checkpoint_dir = \"checkpoints\"\n",
+    "os.makedirs(log_dir, exist_ok=True)\n",
+    "os.makedirs(checkpoint_dir, exist_ok=True)\n",
+    "writer = SummaryWriter(log_dir=log_dir)\n",
+    "\n",
+    "def logits_to_bboxes(logits,image) : # This function is useful for converting the mask into bounding boxes.(The model does not provide bounding boxes.)\n",
+    "    correct_shape = (300, 300)  \n",
+    "    logits_temp = F.interpolate(logits, size=correct_shape, mode='bilinear', align_corners=False)\n",
+    "    logits_temp = logits_temp.cpu().detach().numpy().astype(np.float32)\n",
+    "\n",
+    "    heatmap_count = logits_temp.shape[1]\n",
+    "    heatmaps = [logits_temp[i][k] for i in range(logits_temp.shape[0]) for k in range(heatmap_count)]\n",
+    "    regions = parallel_get_regions(heatmaps=heatmaps, orig_size=image.size, id2label=model.config.id2label)\n",
+    "\n",
+    "    final_bboxes = []\n",
+    "    for i in regions.bboxes :\n",
+    "        final_bboxes.append(i.bbox)\n",
+    "    return final_bboxes\n",
+    "\n",
+    "\n",
+    "def loss_function(): # This model does not have inbuild loss function, So we have to define it according to our dataset and the Requirements.\n",
+    "    pass"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Fine-Tuning Process"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "num_epochs = 5\n",
+    "for epoch in range(num_epochs):\n",
+    "    model.train()\n",
+    "    running_loss = 0.0\n",
+    "    avg_loss = 0.0\n",
+    "\n",
+    "    for idx, item in enumerate(tqdm(dataset, desc=f\"Epoch {epoch + 1}/{num_epochs}\")):\n",
+    "\n",
+    "        images = [prepare_image_detection(img=item['image'], processor=load_processor())]\n",
+    "        images = torch.stack(images, dim=0).to(model.dtype).to(model.device)\n",
+    "        \n",
+    "        optimizer.zero_grad()\n",
+    "        outputs = model(pixel_values=images)\n",
+    "\n",
+    "        predicted_boxes = logits_to_bboxes(outputs.logits, item['image'])\n",
+    "        target_boxes = item['bboxes']\n",
+    "\n",
+    "        loss = loss_function(predicted_boxes,target_boxes)\n",
+    "\n",
+    "        loss.backward()\n",
+    "        optimizer.step()\n",
+    "        running_loss += loss.item()\n",
+    "\n",
+    "        avg_loss = 0.9 * avg_loss + 0.1 * loss.item() if idx > 0 else loss.item()\n",
+    "\n",
+    "    avg_loss = running_loss / len(dataset)\n",
+    "    writer.add_scalar('Training Loss', avg_loss, epoch + 1)\n",
+    "    print(f\"Average Loss for Epoch {epoch + 1}: {avg_loss:.4f}\")\n",
+    "\n",
+    "    torch.save(model.state_dict(), os.path.join(checkpoint_dir, f\"model_epoch_{epoch + 1}.pth\"))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Loading The Checkpoint "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "checkpoint_path = 'checkpoints/model_epoch_350.pth'  \n",
+    "state_dict = torch.load(checkpoint_path,weights_only=True)\n",
+    "\n",
+    "model.load_state_dict(state_dict)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.to('cpu')\n",
+    "model.save_pretrained(\"fine-tuned-surya-model-layout\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "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.10.14"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}