notebooks/demo.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import copy\n",
    "import matplotlib.pyplot as plt\n",
    "import math\n",
    "import os\n",
    "from omegaconf import OmegaConf\n",
    "\n",
    "import torch\n",
    "from torch_geometric.utils import to_dense_adj\n",
    "import torchvision.transforms as T\n",
    "from fsspec.core import url_to_fs\n",
    "from hydra.utils import instantiate\n",
    "from trainer.data.util import AddCanvasElement, AddRelationConstraints, sparse_to_dense\n",
    "from trainer.global_configs import DATASET_DIR, JOB_DIR\n",
    "from trainer.helpers.layout_tokenizer import LayoutSequenceTokenizer\n",
    "from trainer.helpers.sampling import SAMPLING_CONFIG_DICT\n",
    "from trainer.helpers.task import get_cond, filter_canvas\n",
    "from trainer.helpers.visualization import save_gif, save_image, save_label, save_label_with_size, save_relation\n",
    "from trainer.hydra_configs import TestConfig\n",
    "\n",
    "SIZE = (360, 240)\n",
    "\n",
    "# user tunable parameters\n",
    "# cond_type, W_CANVAS = \"relation\", True  # uncomment this line if you want to try relation task\n",
    "cond_type, W_CANVAS = \"cwh\", False  # choices: unconditional, c, cwh, partial, refinement\n",
    "n_samples = 4  # num. of samples to generate at once\n",
    "target_index = 0  # index of real data, partial fields in it are used for conditional generation\n",
    "\n",
    "job_dir = os.path.join(JOB_DIR, \"layoutdm_publaynet/0\")\n",
    "\n",
    "config_path = os.path.join(job_dir, \"config.yaml\")\n",
    "fs, _ = url_to_fs(config_path)\n",
    "if fs.exists(config_path):\n",
    "    with fs.open(config_path, \"rb\") as file_obj:\n",
    "        train_cfg = OmegaConf.load(file_obj)\n",
    "else:\n",
    "    raise FileNotFoundError\n",
    "train_cfg.dataset.dir = DATASET_DIR\n",
    "\n",
    "test_cfg = OmegaConf.structured(TestConfig)\n",
    "test_cfg.cond = cond_type\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "sampling_cfg = OmegaConf.structured(SAMPLING_CONFIG_DICT[test_cfg.sampling])  # NOTE: you may change sampling algorithm\n",
    "OmegaConf.set_struct(sampling_cfg, False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize data and model\n",
    "tokenizer = LayoutSequenceTokenizer(\n",
    "    data_cfg=train_cfg.data, dataset_cfg=train_cfg.dataset\n",
    ")\n",
    "model = instantiate(train_cfg.model)(\n",
    "    backbone_cfg=train_cfg.backbone, tokenizer=tokenizer\n",
    ").to(device)\n",
    "model_path = os.path.join(job_dir, \"best_model.pt\")\n",
    "with fs.open(model_path, \"rb\") as file_obj:\n",
    "    model.load_state_dict(torch.load(file_obj))\n",
    "model = model.to(device)\n",
    "model.eval()\n",
    "sampling_cfg = model.aggregate_sampling_settings(sampling_cfg, test_cfg)\n",
    "\n",
    "if W_CANVAS:\n",
    "    # add canvas and shift label id to load relation gts\n",
    "    assert cond_type == \"relation\"\n",
    "    transform = T.Compose([\n",
    "        AddCanvasElement(),\n",
    "        AddRelationConstraints(edge_ratio=0.1),\n",
    "    ])\n",
    "else:\n",
    "    assert cond_type != \"relation\"\n",
    "    transform = None\n",
    "dataset = instantiate(train_cfg.dataset)(split=\"test\", transform=transform)\n",
    "save_kwargs = {\n",
    "    \"colors\": dataset.colors, \"names\": dataset.labels,\n",
    "    \"canvas_size\": SIZE, \"use_grid\": True,\n",
    "    # \"draw_label\": True,\n",
    "}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Real data visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load target data and visualize GT\n",
    "bbox, label, _, mask = sparse_to_dense(dataset[target_index])\n",
    "gt_cond = model.tokenizer.encode(\n",
    "    {\"label\": label, \"mask\": mask, \"bbox\": bbox}\n",
    ")\n",
    "if \"bos\" in tokenizer.special_tokens:\n",
    "    gt = model.tokenizer.decode(gt_cond[\"seq\"][:, 1:])\n",
    "else:\n",
    "    gt = model.tokenizer.decode(gt_cond[\"seq\"])\n",
    "if W_CANVAS:\n",
    "    gt = filter_canvas(gt)  # remove canvas attributes before visualization\n",
    "plt.axis(\"off\")\n",
    "plt.imshow(save_image(gt[\"bbox\"], gt[\"label\"], gt[\"mask\"], **save_kwargs))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Unconditional Generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert cond_type == \"unconditional\"\n",
    "pred = model.sample(batch_size=n_samples, cond=None, sampling_cfg=sampling_cfg)\n",
    "plt.axis(\"off\")\n",
    "plt.imshow(save_image(pred[\"bbox\"], pred[\"label\"], pred[\"mask\"], **save_kwargs))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Conditional Generation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cond = get_cond(\n",
    "    batch=dataset[target_index],\n",
    "    tokenizer=model.tokenizer,\n",
    "    cond_type=cond_type,\n",
    "    model_type=type(model).__name__,\n",
    ")\n",
    "pred = model.sample(batch_size=n_samples, cond=cond, sampling_cfg=sampling_cfg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Visualization of conditional inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.axis(\"off\")\n",
    "input_ = model.tokenizer.decode(cond[\"seq\"].cpu())\n",
    "mask = pred[\"mask\"][0]\n",
    "label, bbox = pred[\"label\"][0][mask], pred[\"bbox\"][0][mask]\n",
    "if cond_type == \"c\":\n",
    "    plt.imshow(save_label(label, **save_kwargs))\n",
    "elif cond_type == \"cwh\":\n",
    "    plt.imshow(save_label_with_size(label, bbox, **save_kwargs))\n",
    "elif cond_type == \"relation\":\n",
    "    data = cond[\"batch_w_canvas\"]\n",
    "    edge_attr = to_dense_adj(data.edge_index, data.batch, data.edge_attr)\n",
    "    plt.imshow(save_relation(label_with_canvas=data.y.cpu(), edge_attr=edge_attr.cpu()[0], **save_kwargs))\n",
    "elif cond_type  == \"partial\":\n",
    "    plt.imshow(save_image(input_[\"bbox\"], input_[\"label\"], input_[\"mask\"], **save_kwargs))\n",
    "elif cond_type == \"refinement\":\n",
    "    noisy_input = model.tokenizer.decode(cond[\"seq_orig\"].cpu())\n",
    "    plt.imshow(save_image(noisy_input[\"bbox\"][0:1], noisy_input[\"label\"][0:1], noisy_input[\"mask\"][0:1], **save_kwargs))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Visualization of outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(15, 5))\n",
    "ax.set_axis_off()\n",
    "ax.imshow(save_image(pred[\"bbox\"], pred[\"label\"], pred[\"mask\"], **save_kwargs, nrow=int(math.sqrt(n_samples) * 2)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Make GIF for Unconditional Generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_save_kwargs = copy.deepcopy(save_kwargs)\n",
    "new_save_kwargs.pop(\"use_grid\")\n",
    "ids_list = model.model.sample(\n",
    "    batch_size=4,\n",
    "    sampling_cfg=sampling_cfg,\n",
    "    get_intermediate_results=True,\n",
    ")\n",
    "images = []\n",
    "for ids in ids_list:\n",
    "    layouts = model.tokenizer.decode(ids)\n",
    "    image = save_image(\n",
    "        layouts[\"bbox\"],\n",
    "        layouts[\"label\"],\n",
    "        layouts[\"mask\"],\n",
    "        **new_save_kwargs\n",
    "    )\n",
    "    images.append(image)\n",
    "N_step = len(images)\n",
    "images = images[int(0.5*N_step):]\n",
    "save_gif(images, \"../tmp/animation/{}.gif\")\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Dump colors of all labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "labels = []\n",
    "for i, name in enumerate(save_kwargs[\"names\"]):\n",
    "    if \"_cutout\" in name:\n",
    "        continue\n",
    "    else:\n",
    "        labels.append(i)\n",
    "plt.axis(\"off\")\n",
    "plt.imshow(save_label(labels, **save_kwargs))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.12 ('trainer-GSN1huuF-py3.7')",
   "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.12"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "df1bfdd73842a1319c146cb5c112d3818824eaeb9e2048caf661736e74000887"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}