train-segs-v4.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ccd43042",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pylab inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0fa13382",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "\n",
    "import os\n",
    "import scipy.ndimage as ndi\n",
    "import torch\n",
    "\n",
    "# import torchtrainers as tt\n",
    "from webdataset import WebDataset, WebLoader\n",
    "from ocropus4train import ocrhelpers as helpers\n",
    "from ocropus4train import ocrmodels2 as models\n",
    "import matplotlib.pyplot as plt\n",
    "from random import randrange\n",
    "from ocropus4train import nlbin\n",
    "\n",
    "plt.rc(\"image\", cmap=\"gray\")\n",
    "plt.rc(\"image\", interpolation=\"nearest\")\n",
    "\n",
    "\n",
    "def RUN(x):\n",
    "    print(x, \":\", os.popen(x).read().strip())\n",
    "\n",
    "\n",
    "RUN(\"date\")\n",
    "RUN(\"hostname\")\n",
    "RUN(\"nvidia-smi -L\")\n",
    "RUN(\"pwd\")\n",
    "\n",
    "if \"GOPEN_VERBOSE\" in os.environ:\n",
    "    del os.environ[\"GOPEN_VERBOSE\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e2d8d28",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_urls = \"gs://ocro-iaa/segs/books-{000000..000653}-segs.tar\"\n",
    "batch_size = 4\n",
    "mname = \"seg_unet_v4\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4aaa494",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ocropus4train.ocraugment import maybe, normalize, noisify, threshold, distort\n",
    "\n",
    "\n",
    "def affine_transform(image, matrix, output_shape=None, **kw):\n",
    "    if image.ndim == 2:\n",
    "        return ndi.affine_transform(image, matrix, output_shape=output_shape, **kw)\n",
    "    elif image.ndim == 3:\n",
    "        result = np.zeros(output_shape, dtype=image.dtype)\n",
    "        for i in range(image.shape[2]):\n",
    "            ndi.affine_transform(image[:, :, i], matrix, output_shape=output_shape[:2], output=result[:, :, i], **kw)\n",
    "        return result\n",
    "\n",
    "\n",
    "def random_patch(images, orders, size=512):\n",
    "    h, w = images[0].shape[:2]\n",
    "    for image in images:\n",
    "        assert image.shape[:2] == (h, w)\n",
    "    matrix = np.eye(3)\n",
    "    matrix[:2, :2] += np.random.randn(2, 2) * 0.03\n",
    "    # pick the translation so that we get a patch from anywhere within the image\n",
    "    matrix[0, 2] = -random.uniform(-size / 2, h - size / 2)\n",
    "    matrix[1, 2] = -random.uniform(-size / 2, w - size / 2)\n",
    "    # compute the transform from the image to the patch\n",
    "    matrix = np.linalg.inv(matrix)\n",
    "    # apply the affine transform\n",
    "    patches = []\n",
    "    global last_xform \n",
    "    last_xform = locals()\n",
    "    for image, order in zip(images, orders):\n",
    "        output_shape = (size, size) if image.ndim == 2 else (size, size, image.shape[2])\n",
    "        patch = affine_transform(\n",
    "            image,\n",
    "            matrix,\n",
    "            output_shape=output_shape,\n",
    "            order=order,\n",
    "            mode=\"constant\",\n",
    "            cval=0,\n",
    "        )\n",
    "        patches.append(patch)\n",
    "    return patches\n",
    "\n",
    "\n",
    "def random_patch_simple(images, size=512):\n",
    "    h, w = images[0].shape[:2]\n",
    "    for image in images:\n",
    "        assert image.shape[:2] == (h, w)\n",
    "    y = randrange(0, h - size)\n",
    "    x = randrange(0, w - size)\n",
    "    patches = []\n",
    "    for image in images:\n",
    "        patches.append(image[y : y + size, x : x + size])\n",
    "    return patches\n",
    "\n",
    "\n",
    "def get_patches(sample, size=512):\n",
    "    global last_img, last_seg\n",
    "    image1 = sample[\"jpg\"]\n",
    "    assert image1.ndim == 2\n",
    "    assert image1.dtype == np.uint8\n",
    "    image1 = image1.astype(np.float32) / 255.0\n",
    "    if maybe(0.1):  # FIXME only for testing\n",
    "        binarized = nlbin.nlbin(image1)\n",
    "        if maybe(0.2):\n",
    "            binarized = (binarized > 0.5).astype(np.float32)\n",
    "        # binarization sometimes changes image size slightly\n",
    "        assert np.amax(np.abs(np.array(binarized.shape) - np.array(image1.shape))) <= 2\n",
    "        # use ndi.affine_transform to make binarized exactly the same shape as image1\n",
    "        image1 = ndi.affine_transform(\n",
    "            binarized, np.eye(3), output_shape=image1.shape, order=0, mode=\"nearest\"\n",
    "        )\n",
    "    if image1.shape[0] < size or image1.shape[1] < size:\n",
    "        return\n",
    "    wseg = sample[\"words.jpg\"].clip(0, 3)\n",
    "    lseg = sample[\"lines.jpg\"].clip(0, 3)\n",
    "    assert wseg.shape == lseg.shape, (wseg.shape, lseg.shape)\n",
    "    # print(\"words, lines\", wseg.shape, np.sum(wseg >= 1), np.sum(lseg >= 1))\n",
    "    seg1 = np.zeros((wseg.shape[0], wseg.shape[1], 7), dtype=np.uint8)\n",
    "    seg1[:, :, 1] = ndi.maximum_filter((wseg >= 1), 12)\n",
    "    seg1[:, :, 2] = (wseg >= 2)\n",
    "    seg1[:, :, 3] = (wseg >= 3)\n",
    "    seg1[:, :, 4] = ndi.maximum_filter((lseg >= 1), 12)\n",
    "    seg1[:, :, 5] = (lseg >= 2)\n",
    "    seg1[:, :, 6] = (lseg >= 3)\n",
    "    seg1[:, :, 0] = 1 - np.amax(seg1[:, :, 1:], axis=-1)\n",
    "    if np.sum(seg1[:, :, 3]) == 0:\n",
    "        # no words present\n",
    "        print(\"no words\", np.sum(seg1[:, :, 3]))\n",
    "        return\n",
    "    if np.sum(seg1[:, :, 6]) == 0:\n",
    "        # no lines present\n",
    "        print(\"no lines\", np.sum(seg1[:, :, 3]))\n",
    "        return\n",
    "    for i in range(32):\n",
    "        # TODO: add 90 degree rotations to augmentations\n",
    "        # TODO: do something about isolated characters\n",
    "        last_img, last_seg = img, seg = random_patch([image1, seg1], [1, 0], size=size)\n",
    "        if np.sum(seg[:, :, 3]) == 0:\n",
    "            # print(\"no words in patch\", np.sum(seg[:, :, 3]))\n",
    "            # global hack; hack = locals(); raise ValueError()\n",
    "            continue\n",
    "        if maybe(0.05):\n",
    "            img = threshold(img)\n",
    "        if maybe(0.5):\n",
    "            img = distort(img)\n",
    "        if maybe(0.5):\n",
    "            img = normalize(img)\n",
    "        if maybe(0.5):\n",
    "            img = noisify(img)\n",
    "        if maybe(0.1):\n",
    "            sigma = random.uniform(0.5, 2.0)\n",
    "            img = ndi.gaussian_filter(img, sigma)\n",
    "        if np.amin(img) < 0 or np.amax(img) > 1:\n",
    "            print(\"bad patch\", np.amin(img), np.amax(img))\n",
    "            continue\n",
    "        yield (img, seg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71b11fb5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def expand_patches(src):\n",
    "    for sample in src:\n",
    "        # print(sample[\"__key__\"])\n",
    "        for result in get_patches(sample):\n",
    "            yield result\n",
    "\n",
    "\n",
    "def usm_patch(img):\n",
    "    assert img.ndim == 2, img.shape\n",
    "    assert img.dtype == np.float32, img.dtype\n",
    "    assert np.amin(img) >= 0.0 and np.amax(img) <= 1.0\n",
    "    img = img - ndi.gaussian_filter(img, 16.0, mode=\"nearest\")\n",
    "    if np.amin(img) >= 0.0:\n",
    "        print(\"usm_patch: no negative values???\")\n",
    "        # fix this to make subsequent processing tages happy\n",
    "        img = img - np.amin(img) + 0.001 * np.random.normal(size=img.shape)\n",
    "        img = img.astype(np.float32)\n",
    "    img = img.clip(-5.0, 5.0)\n",
    "    return img\n",
    "\n",
    "\n",
    "def img_tensor(img):\n",
    "    assert img.ndim == 2, img.shape\n",
    "    assert img.dtype == np.float32, img.dtype\n",
    "    assert np.amax(img) < 10.0  # make sure it already got normalied somewhere\n",
    "    return torch.tensor(img).unsqueeze(0)\n",
    "\n",
    "\n",
    "def label_tensor(label):\n",
    "    assert label.ndim == 3, label.shape\n",
    "    assert label.shape[2] == 7, label.shape\n",
    "    assert label.dtype == np.uint8, label.dtype\n",
    "    return torch.tensor(label.transpose(2, 0, 1)).float()\n",
    "\n",
    "\n",
    "def pipeline(ds):\n",
    "    return (\n",
    "        ds.decode(\"l8\")\n",
    "        .compose(expand_patches)\n",
    "        .map_tuple(usm_patch)\n",
    "        .map_tuple(img_tensor, label_tensor)\n",
    "    )\n",
    "\n",
    "\n",
    "training = (\n",
    "    pipeline(WebDataset(training_urls, resampled=True).shuffle(10))\n",
    "    .shuffle(1000)\n",
    "    .batched(batch_size)\n",
    ")\n",
    "\n",
    "batch = next(iter(training))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca5d1294",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(batch[0].shape, batch[1].shape)\n",
    "subplot(121); imshow(batch[0][0, 0].numpy())\n",
    "r = batch[1][0, 1:4].sum(axis=0).numpy().astype(np.float32) / 3.0\n",
    "g = batch[1][0, 4:7].sum(axis=0).numpy().astype(np.float32) / 3.0\n",
    "b = batch[1][0, 0].numpy().astype(np.float32())\n",
    "color = np.stack([r, g, b], axis=-1)\n",
    "subplot(122); imshow(color)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e840fc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_dl = WebLoader(\n",
    "    training, batch_size=None, num_workers=4, pin_memory=True\n",
    ").with_epoch(10000 // batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62ba9dcc",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = \"cuda:0\"\n",
    "model = models.make(mname)\n",
    "model.to(device)\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d865412e",
   "metadata": {},
   "outputs": [],
   "source": [
    "gpus = !nvidia-smi -L\n",
    "if len(gpus) > 1:\n",
    "    model.to(\"cuda:0\");\n",
    "    model.model[10].to(\"cuda:1\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e2fd7d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "lossfn = torch.nn.BCEWithLogitsLoss()\n",
    "trainer = helpers.SegTrainer(model, masked=-1, device=None, lossfn=lossfn)\n",
    "trainer.load_best()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce08cce0",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.train(training_dl, learning_rates=[1e-2] + [1e-3] * 15 + [3e-4] * 50 + [1e-4] * 200, every=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1afc3e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "next(model.model[10].module.parameters()).device, model.model[10].device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98c1277d",
   "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.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}