notebooks/vis_experiment.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eaade18c",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import json\n",
    "import cv2\n",
    "import matplotlib.cm as cm\n",
    "import math\n",
    "from PIL import Image\n",
    "\n",
    "\n",
    "# load test frame,\n",
    "def load_test_frame_files(file):\n",
    "    with open(file, 'r') as f:\n",
    "        meta = json.load(f)\n",
    "    fs = meta['frames']\n",
    "    fs = sorted(fs, key=lambda d: d['file_path'])\n",
    "    \n",
    "    frames = []\n",
    "    for frame in fs:\n",
    "        frames.append(frame['file_path'])\n",
    "    return frames\n",
    "\n",
    "def format_axes(axes):\n",
    "    for ax in axes:\n",
    "        if type(ax) is np.ndarray:\n",
    "            format_axes(ax)\n",
    "        else:\n",
    "            ax.set_xticks([])\n",
    "            ax.set_yticks([])\n",
    "            \n",
    "def weighted_percentile(x, w, ps, assume_sorted=False):\n",
    "    \"\"\"Compute the weighted percentile(s) of a single vector.\"\"\"\n",
    "    x = x.reshape([-1])\n",
    "    w = w.reshape([-1])\n",
    "    if not assume_sorted:\n",
    "        sortidx = np.argsort(x)\n",
    "    x, w = x[sortidx], w[sortidx]\n",
    "    acc_w = np.cumsum(w)\n",
    "    return np.interp(np.array(ps) * (acc_w[-1] / 100), acc_w, x)\n",
    "            \n",
    "def visualize_cmap(value,\n",
    "                   weight,\n",
    "                   colormap,\n",
    "                   lo=None,\n",
    "                   hi=None,\n",
    "                   percentile=99.,\n",
    "                   curve_fn=lambda x: x,\n",
    "                   modulus=None,\n",
    "                   matte_background=True):\n",
    "    \"\"\"Visualize a 1D image and a 1D weighting according to some colormap.\n",
    "\n",
    "    Args:\n",
    "    value: A 1D image.\n",
    "    weight: A weight map, in [0, 1].\n",
    "    colormap: A colormap function.\n",
    "    lo: The lower bound to use when rendering, if None then use a percentile.\n",
    "    hi: The upper bound to use when rendering, if None then use a percentile.\n",
    "    percentile: What percentile of the value map to crop to when automatically\n",
    "      generating `lo` and `hi`. Depends on `weight` as well as `value'.\n",
    "    curve_fn: A curve function that gets applied to `value`, `lo`, and `hi`\n",
    "      before the rest of visualization. Good choices: x, 1/(x+eps), log(x+eps).\n",
    "    modulus: If not None, mod the normalized value by `modulus`. Use (0, 1]. If\n",
    "      `modulus` is not None, `lo`, `hi` and `percentile` will have no effect.\n",
    "    matte_background: If True, matte the image over a checkerboard.\n",
    "\n",
    "    Returns:\n",
    "    A colormap rendering.\n",
    "    \"\"\"\n",
    "    # Identify the values that bound the middle of `value' according to `weight`.\n",
    "    lo_auto, hi_auto = weighted_percentile(\n",
    "      value, weight, [50 - percentile / 2, 50 + percentile / 2])\n",
    "\n",
    "    # If `lo` or `hi` are None, use the automatically-computed bounds above.\n",
    "    eps = np.finfo(np.float32).eps\n",
    "    lo = lo or (lo_auto - eps)\n",
    "    hi = hi or (hi_auto + eps)\n",
    "\n",
    "    # Curve all values.\n",
    "    value, lo, hi = [curve_fn(x) for x in [value, lo, hi]]\n",
    "\n",
    "    # Wrap the values around if requested.\n",
    "    if modulus:\n",
    "        value = np.mod(value, modulus) / modulus\n",
    "    else:\n",
    "        # Otherwise, just scale to [0, 1].\n",
    "        value = np.nan_to_num(\n",
    "        np.clip((value - np.minimum(lo, hi)) / np.abs(hi - lo), 0, 1))\n",
    "\n",
    "    if colormap:\n",
    "        colorized = colormap(value)[:, :, :3]\n",
    "    else:\n",
    "        assert len(value.shape) == 3 and value.shape[-1] == 3\n",
    "        colorized = value\n",
    "\n",
    "    return colorized\n",
    "\n",
    "depth_curve_fn = lambda x: -np.log(x + np.finfo(np.float32).eps)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b54099c",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# set path to experiments and groundtruth data here\n",
    "EXPERIMENT_DIR = '/mnt/res_nas/silvanweder/experiments' # set this to where you saved your experiments to\n",
    "GROUNDTRUTH_DIR = '/mnt/res_nas/silvanweder/datasets/object-removal-custom-clean' # set this to where you saved your data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "846d3cf1",
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Available Experiments:')\n",
    "for exp in os.listdir(EXPERIMENT_DIR):\n",
    "    if exp.startswith('.'):\n",
    "        continue\n",
    "    print('\\t -', exp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3650b92a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# select experiments from available experiments above\n",
    "experiment = 'final_tests_real'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd8474bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "suffix = '_real' # set to either _real or _synthetic depending on what masks you used"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3d0c648",
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Available Sequences:')\n",
    "for sc in sorted(os.listdir(os.path.join(EXPERIMENT_DIR, experiment))):\n",
    "    print(f'\\t- {sc}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eeb2220d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set available sequence here\n",
    "sequence = '002'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c75c2889",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# set options here\n",
    "rotate = True # do we need to rotate the renderings\n",
    "eval_run = 'train_test_preds' # visualizing test or training images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6e2361b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "experiment_path = os.path.join(EXPERIMENT_DIR, experiment, sequence)\n",
    "groundtruth_path = os.path.join(GROUNDTRUTH_DIR, sequence)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c55b2ee",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "if eval_run == 'test_preds':\n",
    "    test_frame_files = load_test_frame_files(os.path.join(experiment_path, 'transforms_test.json'))\n",
    "elif eval_run == 'train_test_preds':\n",
    "    test_frame_files = load_test_frame_files(os.path.join(experiment_path, 'transforms_train.json'))\n",
    "else:\n",
    "    raise ValueError(f'Invalid eval run {eval_run}')\n",
    "\n",
    "\n",
    "for i, frame in enumerate(test_frame_files):\n",
    "    \n",
    "    # load groundtruth image\n",
    "    image_gt = np.asarray(Image.open(os.path.join(groundtruth_path, 'images', frame.split('/')[-1])))\n",
    "    image_gt = cv2.resize(image_gt, (256, 192))\n",
    "\n",
    "    # load input image\n",
    "    input_mask = np.load(os.path.join(groundtruth_path, frame.replace('images', f'masks{suffix}').replace('jpg', 'npy'))) \n",
    "    image_input = image_gt.copy()\n",
    "    image_input[input_mask == 1] = (255, 255, 255)\n",
    "\n",
    "    image_inpainted = np.asarray(Image.open(os.path.join(experiment_path, *frame.replace('images', f'lama_images_output{suffix}').replace('.jpg', '_mask001.png').split('/')[-2:])))\n",
    "\n",
    "    image_est = np.asarray(Image.open(os.path.join(experiment_path, eval_run, f'color_{str(i).zfill(3)}.png')))\n",
    "    depth_est = np.asarray(Image.open(os.path.join(experiment_path, eval_run, f'distance_mean_{str(i).zfill(3)}.tiff')))\n",
    "\n",
    "    # colorize uncertainty and depth map\n",
    "    depth_est = (visualize_cmap(depth_est, np.ones_like(depth_est), cm.get_cmap('turbo'), curve_fn=depth_curve_fn).copy() * 255).astype(np.uint8)\n",
    "    \n",
    "    # rotate all images\n",
    "    if rotate:\n",
    "        image_gt = cv2.rotate(image_gt, cv2.ROTATE_90_CLOCKWISE)\n",
    "        image_inpainted = cv2.rotate(image_inpainted,  cv2.ROTATE_90_CLOCKWISE)\n",
    "        image_input = cv2.rotate(image_input,  cv2.ROTATE_90_CLOCKWISE)\n",
    "        image_est = cv2.rotate(image_est,  cv2.ROTATE_90_CLOCKWISE)\n",
    "        depth_est = cv2.rotate(depth_est,  cv2.ROTATE_90_CLOCKWISE)\n",
    "\n",
    "\n",
    "    fig, ax = plt.subplots(1, 5, figsize=(5 * 12, 16))\n",
    "    ax[0].imshow(image_gt)\n",
    "    ax[0].set_title('Groundtruth Image', fontsize=60)\n",
    "    ax[1].imshow(image_input)\n",
    "    ax[1].set_title('Masked Image', fontsize=60)\n",
    "    ax[2].imshow(image_inpainted)\n",
    "    ax[2].set_title('Inpainted Image', fontsize=60)\n",
    "    ax[3].imshow(image_est)\n",
    "    ax[3].set_title('Rendered Image', fontsize=60)\n",
    "    ax[4].imshow(depth_est)\n",
    "    ax[4].set_title('Rendered Depth', fontsize=60)\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "    plt.close('all')\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "object-remoival",
   "language": "python",
   "name": "object-removal"
  },
  "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}