150-layout.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Populating the interactive namespace from numpy and matplotlib\n"
     ]
    }
   ],
   "source": [
    "%pylab inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "lines_to_next_cell": 2,
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torchmore import layers, flex"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# SEMANTIC SEGMENTATION"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Task\n",
    "\n",
    "<img src=\"figs/seg-example.png\" style=\"height: 4in\" />"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Image to Image Networks\n",
    "\n",
    "- input: RGB image, output: RGB image\n",
    "    - example: noise removal, image enhancement, upscaling\n",
    "- input: RGB image, output: integer classes\n",
    "    - example: semantic segmentation, binarization\n",
    "- output resolution may be same, higher (upscaling), or lower\n",
    "- often, segmentation / classificaiton is sufficient at lower resolution\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Semantic Segmentation\n",
    "\n",
    "Here is a minimal example of a semantic segmentation network. Note that `NLLLoss` automatically extends to 2D."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.0184, grad_fn=<NllLoss2DBackward>) torch.Size([1, 10, 256, 256]) torch.Size([1, 256, 256])\n"
     ]
    }
   ],
   "source": [
    "model = nn.Sequential(\n",
    "    nn.Conv2d(3, 64, 3, padding=1),\n",
    "    nn.BatchNorm2d(64),\n",
    "    nn.ReLU(),\n",
    "    nn.Conv2d(64, 10, 3, padding=1)\n",
    ")\n",
    "target = torch.zeros((1, 256, 256)).long()\n",
    "input = torch.zeros((1, 3, 256, 256))\n",
    "output = model(input)\n",
    "lossfn = nn.NLLLoss()\n",
    "loss = lossfn(output, target)\n",
    "print(loss, output.size(), target.size())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Image to Image\n",
    "\n",
    "Here is a minimal example of an image to image network. Note that `MSELoss` automatically extends to 2D."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.0008, grad_fn=<MseLossBackward>) torch.Size([1, 3, 256, 256]) torch.Size([1, 3, 256, 256])\n"
     ]
    }
   ],
   "source": [
    "model = nn.Sequential(\n",
    "    nn.Conv2d(3, 64, 3, padding=1),\n",
    "    nn.BatchNorm2d(64),\n",
    "    nn.ReLU(),\n",
    "    nn.Conv2d(64, 3, 3, padding=1)\n",
    ")\n",
    "target = torch.zeros((1, 3, 256, 256))\n",
    "input = torch.zeros((1, 3, 256, 256))\n",
    "output = model(input)\n",
    "lossfn = nn.MSELoss()\n",
    "loss = lossfn(output, target)\n",
    "print(loss, output.size(), target.size())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# GAN-Style Image-to-Image Transformations\n",
    "\n",
    "In addition to the plain image-to-image transformations, we can implement GANs:\n",
    "\n",
    "- add a random noise input in addition to the original input\n",
    "- add a discriminator to the output and propagate back through it\n",
    "\n",
    "GANs can be useful for:\n",
    "\n",
    "- domain transfer\n",
    "- training data generation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# ARCHITECTURES FOR SEMANTIC SEGMENTATION"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# MD-LSTM\n",
    "\n",
    "<img src=\"figs/seg-mdlstm.png\" style=\"height: 4in\" />\n",
    "\n",
    "Byeon, Wonmin, and Thomas M. Breuel. \"Supervised texture segmentation using 2D LSTM networks.\" 2014 IEEE International Conference on Image Processing (ICIP). IEEE, 2014.\n",
    "\n",
    "Byeon, Wonmin, et al. \"Scene labeling with lstm recurrent neural networks.\" Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. 2015."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Fully Convolutional Networks\n",
    "\n",
    "<img src=\"figs/seg-fully-conv.png\" style=\"height: 4in\" />\n",
    "\n",
    "Long, Jonathan, Evan Shelhamer, and Trevor Darrell. \"Fully convolutional networks for semantic segmentation.\" Proceedings of the IEEE conference on computer vision and pattern recognition. 2015."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Issue: Footprints\n",
    "\n",
    "What information can be taken into account in making segmentation decisions?\n",
    "\n",
    "**Convolutional Architectures** $2^k \\cdot r$ with $k$ = number of max-pooling layers and $r$ footprint of the convolutions\n",
    "\n",
    "**MD-LSTM Architectures** Unlimited / global.\n",
    "\n",
    "Consequence:\n",
    "\n",
    "- may need very deep networks to take whole page information into account with convolutional architectures\n",
    "- even then, precise geometry is difficult to keep track of across the whole page"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Issue: Output Resolution\n",
    "\n",
    "Problem:\n",
    "\n",
    "- often, we want pixel-accurate outputs for segmentation\n",
    "- standard architectures with max pooling increase footprints but decrease output resolution"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Issue: Output Resolution\n",
    "\n",
    "Solution:\n",
    "\n",
    "- simply scale up low resolution segmentations to full image size after inference\n",
    "- scale up low resolution segmentations stepwise as part of the network architecture\n",
    "- make the scaling-up kernels themselves learnable (\"transposed convolutions\")\n",
    "- avoid the down scaling by replacing max pooling with strided convolutions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Transposed Convolution\n",
    "\n",
    "(\"Deconvolutional\")\n",
    "\n",
    "<img src=\"figs/seg-deconv.png\" style=\"height: 4in\" />\n",
    "\n",
    "Noh, Hyeonwoo, Seunghoon Hong, and Bohyung Han. \"Learning deconvolution network for semantic segmentation.\" Proceedings of the IEEE international conference on computer vision. 2015."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# U-Net\n",
    "\n",
    "<img src=\"figs/seg-unet.png\" style=\"height: 4in\" />\n",
    "\n",
    "Ronneberger, Olaf, Philipp Fischer, and Thomas Brox. \"U-net: Convolutional networks for biomedical image segmentation.\" International Conference on Medical image computing and computer-assisted intervention. Springer, Cham, 2015."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dilated Convolutions\n",
    "\n",
    "<img src=\"figs/seg-atrous.png\" style=\"height: 4in\" />\n",
    "\n",
    "Chen, Liang-Chieh, et al. \"Deeplab: Semantic image segmentation with deep convolutional nets, atrous convolution, and fully connected crfs.\" IEEE transactions on pattern analysis and machine intelligence 40.4 (2017): 834-848."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "lines_to_next_cell": 0,
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# LAYOUT MODELS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "lines_to_next_cell": 0,
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Layout Analysis\n",
    "\n",
    "Goals of document layout analysis:\n",
    "\n",
    "- identify text/image regions on pages\n",
    "- find text lines\n",
    "- find words within text lines\n",
    "- determine reading order"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Document Segmentation\n",
    "\n",
    "<img src=\"figs/docseg-printed.png\" style=\"height=5in\" />\n",
    "\n",
    "Corbelli, Andrea, et al. \"Historical document digitization through layout analysis and deep content classification.\" 2016 23rd International Conference on Pattern Recognition (ICPR). IEEE, 2016."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Historical Layout Analysis\n",
    "\n",
    "<img src=\"figs/docseg-handwritten.png\" style=\"height=5in\" />\n",
    "\n",
    "Chen, Kai, et al. \"Convolutional neural networks for page segmentation of historical document images.\" 2017 14th IAPR International Conference on Document Analysis and Recognition (ICDAR). Vol. 1. IEEE, 2017."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RCNN for Layout Analysis\n",
    "\n",
    "- RCNN has been used for word bounding boxes in scene text\n",
    "- likely does not work well for whole page segmentation\n",
    "  - regions often not exactly axis-aligned rectangles\n",
    "  - regions wildly differ in scale and shape\n",
    "  - region hypotheses likely difficult to regress from moving windows\n",
    "  \n",
    "Jaderberg, Max, et al. \"Reading text in the wild with convolutional neural networks.\" International Journal of Computer Vision 116.1 (2016): 1-20."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Layout Analysis as Semantic Segmentation\n",
    "\n",
    "- semantic image segmentation = image-to-image transformation\n",
    "- each image is classified into one of several different semantic classes\n",
    "- possible pixel classes for layout analysis:\n",
    "  - text / image / table / figure / background\n",
    "  - baseline / background\n",
    "  - centerline / background\n",
    "  - text / image / text line separator / column separator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "lines_to_next_cell": 0,
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# What do we label?\n",
    "\n",
    "Possibilities:\n",
    "\n",
    "- foreground pixels only\n",
    "- all pixels in a rectangle / region\n",
    "- leave it up to the algorithm\n",
    "\n",
    "Very different results:\n",
    "\n",
    "- foreground/algorithm = we don't know which pixels belong together\n",
    "- all pixels in a region = group pixels together by connected components\n",
    "\n",
    "Choice depends on post-processing."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "lines_to_next_cell": 0,
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Page Level Segmentation\n",
    "\n",
    "Page level segmentation divides images into text regions, image regions, and background.\n",
    "\n",
    "- precise pixel-level segmentations are not usually needed\n",
    "- segmentations can be computed at a lower level of resolution\n",
    "- simple properties like text/image/background can be computed based on local texture / statistics\n",
    "- separating adjacent text columns or images may be difficult, since background gaps may be narrow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Page Level Segmentation\n",
    "\n",
    "Different uses:\n",
    "\n",
    "- simply mask out non-text regions for further processing (basic binary map sufficient)\n",
    "- extract text regions via connected components (requires higher quality segmentation)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Layout Tasks\n",
    "\n",
    "<img src=\"figs/layout-tasks.png\" style=\"width: 10in\" />"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "lines_to_next_cell": 0,
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Simple Approach\n",
    "\n",
    "Word segmentation:\n",
    "\n",
    "- assume training data consists of page images and word bounding boxes\n",
    "- create a binary target that is 1 inside word bounding boxes and 0 elsewhere\n",
    "- learn an image-to-image model predicting the binary map from the input document image\n",
    "\n",
    "Problem:\n",
    "\n",
    "- word bounding boxes are often overlapping\n",
    "- how do we turn the resulting binary image into something we can feed to a word recognizer?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "lines_to_next_cell": 0,
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Centerline / Baseline Approach\n",
    "\n",
    "Word/line segmentation:\n",
    "\n",
    "- assume training data consists of page images and word/line bounding boxes\n",
    "- create a binary image that marks either the center line or the baseline of each bounding box\n",
    "- learn an image-to-image model predicting the centerline/baseline map from the input document image:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "lines_to_next_cell": 0,
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Centerline / Baseline Approach\n",
    "\n",
    "\n",
    "Properties:\n",
    "\n",
    "- works better than the simple approach\n",
    "- still need to use non-DL mechanisms to find the word/line bounding boxes from the seeds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Line Marker Learning\n",
    "\n",
    "<img src=\"figs/layout-line-based.png\" style=\"height: 5in\" />"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "lines_to_next_cell": 0,
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Marker Plus Separator Approach\n",
    "\n",
    "Word/line segmentation:\n",
    "\n",
    "- assume training data consists of page images and word/line bounding boxes\n",
    "- three output classes:\n",
    "  - background\n",
    "  - marker (center of bounding box)\n",
    "  - boundary (outline of bounding box)\n",
    "- train image-to-image segmentation model to output all three classes\n",
    "- recover word/line images via marker morphology"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "lines_to_next_cell": 0,
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Marker Plus Separator Approach\n",
    "\n",
    "\n",
    "Properties:\n",
    "\n",
    "- functions like RCNN, in that it finds both the location and the size of object instances (words/lines)\n",
    "- simpler to understand/tune: we can see the marker/boundary proposals"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Marker-plus-Separator\n",
    "\n",
    "<table>\n",
    "    <tr>\n",
    "        <td><img src=\"figs/word-input.png\" style=\"width: 6in\" /></td>\n",
    "        <td><img src=\"figs/word-segmentation.png\" style=\"width: 6in\" /></td>\n",
    "    </tr>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Word Segmentation\n",
    "\n",
    "<table>\n",
    "    <tr>\n",
    "        <td><img src=\"figs/word-output.png\" style=\"width: 6in\" /></td>\n",
    "        <td><img src=\"figs/word-extracted.png\" style=\"width: 6in\" /></td>\n",
    "    </tr>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Semantic Segmentation into Blocks\n",
    "\n",
    "<img src=\"figs/docseg-convnet.png\" style=\"width: 7in\" />\n",
    "\n",
    "He, Dafang, et al. \"Multi-scale multi-task fcn for semantic page segmentation and table detection.\" 2017 14th IAPR International Conference on Document Analysis and Recognition (ICDAR). Vol. 1. IEEE, 2017. (Also uses CRFs.)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Summary: Layout Analysis as Semantic Segmentation\n",
    "\n",
    "- standard semantic segmentation models from vision\n",
    "    - convolutional, RCNN, LSTM, U-Net, etc.\n",
    "- simple per-pixel region-type classification\n",
    "    - useful for text/image segmentation, masking, etc.\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Summary: Layout Analysis as Semantic Segmentation (2)\n",
    "\n",
    "- text line / word marker learning\n",
    "    - works well on binary images in combination with traditional analysis techniques\n",
    "- marker plus boundary\n",
    "    - permits grayscale-based extraction of elements, multiple hypothesis generation\n",
    "- rectangle (RCNN, regression-like approaches)\n",
    "    - directly yields bounding boxes for extraction; commonly used with scene text\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# END-TO-END SEGMENTATION"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Layout Analysis with U-Net (PyTorch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "from torch import nn\n",
    "class UnetLayer(nn.Module):\n",
    "    def __init__(self, d, r=3, sub=None):\n",
    "        super().__init__()\n",
    "        self.down = nn.MaxPool2d(2)\n",
    "        self.up = flex.ConvTranspose2d(d, r, stride=2, padding=1, output_padding=1)\n",
    "        self.sub = sub\n",
    "    def forward(self, x):\n",
    "        lo = self.down(x)\n",
    "        out = self.sub(lo)\n",
    "        hi = self.up(out)\n",
    "        result = torch.cat([x, hi], dim=1)\n",
    "        return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "def make_seg_unet(noutput=3):\n",
    "    model = nn.Sequential(\n",
    "        layers.Input(\"BHWD\", reorder=\"BDHW\", range=(0, 1), sizes=[None, 1, None, None]),\n",
    "        *conv2d(64, 3, repeat=3),\n",
    "        UnetLayer(64, sub=[\n",
    "            *conv2d(128, 3, repeat=3),\n",
    "            UnetLayer(128, sub=[\n",
    "                *conv2d(256, 3, repeat=3),\n",
    "                UnetLayer(256, sub=[\n",
    "                    *conv2d(512, 3, repeat=3),\n",
    "                ])\n",
    "            ])\n",
    "        ]),\n",
    "        *conv2d(64, 3, repeat=2),\n",
    "        flex.Conv2d(noutput, 5))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Layout Analysis with MDLSTM (PyTorch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "def make_seg_lstm(noutput=3):\n",
    "    model = nn.Sequential(\n",
    "        layers.Input(\"BHWD\", reorder=\"BDHW\", range=(0, 1), sizes=[None, 1, None, None]),\n",
    "        *conv2mp(50, 3, 2, repeat=3),\n",
    "        *conv2mp(100, 3, 2, repeat=3),\n",
    "        flex.Lstm2(200)\n",
    "        *conv2d(400, 5),\n",
    "        flex.Conv2d(noutput, 1)\n",
    "    )\n",
    "    flex.shape_inference(model, (1, 256, 256, 1))\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Semi-Supervised and Weakly Supervised Approaches\n",
    "\n",
    "Problem:\n",
    "\n",
    "- getting training data for layout analysis is difficult\n",
    "- manual markup is a lot of work"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Semi-Supervised and Weakly Supervised Approaches\n",
    "\n",
    "Solution:\n",
    "\n",
    "- text-only / image-only data is easy to get\n",
    "- text regions can be verified via OCR\n",
    "- approach similar to RCNN: generate lots of text boxes and verify\n",
    "- use _soft labels_ and take advantage of posterior probability estimates"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Target Generation\n",
    "\n",
    "- start with `hOCR` output from an existing manually constructed recognizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def target_for_page(image, hocr):\n",
    "    htmlparser = etree.HTMLParser()\n",
    "    doc = etree.parse(io.BytesIO(hocr), htmlparser)\n",
    "    pages = list(doc.xpath('//*[@class=\"ocr_page\"]'))\n",
    "    ...\n",
    "    for word in page.xpath(\"//*[@class='ocrx_word']\"):\n",
    "        ...\n",
    "        target[y0-a:y1+a, x0-a:x1+a] = 1\n",
    "    for word in page.xpath(\"//*[@class='ocrx_word']\"):\n",
    "        ...\n",
    "        target[y0-b:y1+b, x0-b:x1+b] = 0\n",
    "    for word in page.xpath(\"//*[@class='ocrx_word']\"):\n",
    "        ...\n",
    "        target[yc-c:yc+c, x0+d:x1-d] = 2\n",
    "    return target\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- missing words are not that critical\n",
    "- remember: networks estimate posteriors\n",
    "- e.g. 20% missing means score drops from 1.0 to 0.8"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Word Layout Target\n",
    "\n",
    "<img src=\"figs/word-layout-target.png\" style=\"height: 6in\" />"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# The Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_seg_lstm(noutput=3):\n",
    "    model = nn.Sequential(\n",
    "        layers.Input(\"BHWD\", reorder=\"BDHW\", range=(0, 1), sizes=[None, 1, None, None]),\n",
    "        KeepSize(sub=nn.Sequential(\n",
    "                *conv2d(50, 5),\n",
    "                *conv2mp(50, 3, 2, repeat=3),\n",
    "                *conv2mp(100, 3, 2, repeat=3),\n",
    "                *conv2d(200, 3, repeat=2),\n",
    "                flex.BDHW_LSTM(200)\n",
    "            )\n",
    "        ),\n",
    "        *conv2d(400, 5),\n",
    "        flex.Conv2d(noutput, 3)\n",
    "    )\n",
    "    flex.shape_inference(model, (1, 256, 256, 1))\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Word Segmentation Training\n",
    "\n",
    "<img src=\"figs/wordseg-training.png\" style=\"height: 6in\" />"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Markers and Marker Boundaries\n",
    "\n",
    "<table>\n",
    "    <tr>\n",
    "        <td><img src=\"figs/wseg-markerhi.png\" style=\"height: 3in\" /></td>\n",
    "        <td><img src=\"figs/wseg-marker-sep.png\" style=\"height: 3in\" /></td>\n",
    "    </tr>\n",
    "</table>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def step1():\n",
    "    word_markers = (probs[:,:,2] > 0.5)\n",
    "    word_markers = ndi.minimum_filter(ndi.maximum_filter(word_markers, (1, 3)), (1, 3))\n",
    "    _, sources = ndi.distance_transform_edt(1-word_markers, return_indices=True)\n",
    "    word_sources = word_labels[sources[0], sources[1]]\n",
    "    word_boundaries = maximum((roll(word_sources, 1, 0)!=word_sources), roll(word_sources, 1, 1)!=word_sources)\n",
    "    word_boundaries = ndi.minimum_filter(ndi.maximum_filter(word_boundaries, 4), 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Word Regions and Marker-Selected Regions\n",
    "\n",
    "<table>\n",
    "    <tr>\n",
    "        <td><img src=\"figs/wseg-regions-minus-boundaries.png\" style=\"height: 3in\" /></td>\n",
    "        <td><img src=\"figs/wseg-selected-word-regions.png\" style=\"height: 3in\" /></td>\n",
    "    </tr>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def step2():\n",
    "    separators = maximum(probs[:,:,1]>0.3, word_boundaries)\n",
    "    all_components, n = ndi.label(1-separators)\n",
    "    word_markers = (probs[:,:,2] > 0.5) * (1-separators)\n",
    "    word_markers = ndi.minimum_filter(ndi.maximum_filter(word_markers, (1, 3)), (1, 3))\n",
    "    word_labels, n = ndi.label(word_markers)\n",
    "    wordmap = corresponding_labels(all_components, word_labels)\n",
    "    result = wordmap[all_components]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Extracted Word Images\n",
    "\n",
    "<img src=\"figs/wseg-extracted-words.png\" style=\"height: 4in\" />"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_words():\n",
    "    for rect in ndi.find_objects(result):\n",
    "        rect = pad_rect(5)\n",
    "        yield image[rect]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Benefits of Morphology-Based Approaches\n",
    "\n",
    "Compare with R-CNN, rectangle regression, etc. approaches.\n",
    "\n",
    "- very simple code, no special cases\n",
    "- easy to visualize markers / boundaries\n",
    "- can cope with non-rectangular layouts / regions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Semi-Supervised / Unsupervised Training\n",
    "\n",
    "We can now use this word extractor for semi-supervised or unsupervised training:\n",
    "\n",
    "- take unknown documents\n",
    "- run word detector (possibly multiple candidates per marker)\n",
    "- run word OCR over detected words\n",
    "- select words/non-words and optimal rectangle\n",
    "- use as more training data\n",
    "\n",
    "Similar to Fast R-CNN training etc., but without rectangle regression.\n"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Slideshow",
  "jupytext": {
   "cell_metadata_filter": "-all",
   "formats": "ipynb",
   "main_language": "python"
  },
  "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}