visualizer.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "from pyrender import IntrinsicsCamera, Mesh, Node, Scene, Viewer, OffscreenRenderer, DirectionalLight, SpotLight, PointLight, RenderFlags\n",
    "from pathlib import Path\n",
    "import json\n",
    "import trimesh\n",
    "import cv2\n",
    "os.environ.update(\n",
    "    PYOPENGL_PLATFORM = 'egl',\n",
    "    PYOPENGL_FULL_LOGGING = '0'\n",
    ")\n",
    "import json\n",
    "import matplotlib.patches as patches\n",
    "from pycocotools.coco import COCO\n",
    "\n",
    "# Function to load an image\n",
    "def load_image(image_path):\n",
    "    image = cv2.imread(image_path)\n",
    "    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "    return image\n",
    "\n",
    "# Function to visualize COCO detections\n",
    "def visualize_coco_detections(coco_json, image_dir, tgt_img_id=0):\n",
    "    # Load COCO annotations\n",
    "    coco = COCO(coco_json)\n",
    "\n",
    "    # Load image ids\n",
    "    img_ids = coco.getImgIds()\n",
    "    \n",
    "    for img_id in img_ids:\n",
    "        if img_id != tgt_img_id:\n",
    "            continue\n",
    "        # Load image info\n",
    "        img_info = coco.loadImgs(img_id)[0]\n",
    "        image_path = f\"{image_dir}/{img_info['file_name']}\"\n",
    "        image = load_image(image_path)\n",
    "\n",
    "        # Create a plot\n",
    "        fig, ax = plt.subplots(1, figsize=(12, 9))\n",
    "        ax.imshow(image)\n",
    "        \n",
    "        # Load annotations\n",
    "        ann_ids = coco.getAnnIds(imgIds=img_info['id'])\n",
    "        anns = coco.loadAnns(ann_ids)\n",
    "        \n",
    "        for ann in anns:\n",
    "            # Draw bounding box\n",
    "            bbox = ann['bbox']\n",
    "            rect = patches.Rectangle(\n",
    "                (bbox[0], bbox[1]), bbox[2], bbox[3], linewidth=2, edgecolor='r', facecolor='none'\n",
    "            )\n",
    "            ax.add_patch(rect)\n",
    "            \n",
    "            # Draw category name\n",
    "            category_id = ann['category_id']\n",
    "            category_name = coco.loadCats(category_id)[0]['name']\n",
    "            plt.text(\n",
    "                bbox[0], bbox[1] - 2, category_name, fontsize=12, color='red',\n",
    "                bbox=dict(facecolor='yellow', alpha=0.5)\n",
    "            )\n",
    "        plt.axis('off')\n",
    "        plt.show()\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_info = json.load(open('sample_data/models/models_info.json', 'r'))\n",
    "from seaborn import color_palette\n",
    "palette = color_palette('hls', 30)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# borrowed from https://github.com/hughw19/NOCS_CVPR2019/blob/master/utils.py\n",
    "def get_3d_bbox(scale, shift = 0):\n",
    "    \"\"\"\n",
    "    Input: \n",
    "        scale: [3] or scalar\n",
    "        shift: [3] or scalar\n",
    "    Return \n",
    "        bbox_3d: [3, N]\n",
    "    \"\"\"\n",
    "    if hasattr(scale, \"__iter__\"):\n",
    "        bbox_3d = np.array([[scale[0] / 2, +scale[1] / 2, scale[2] / 2],\n",
    "                  [scale[0] / 2, +scale[1] / 2, -scale[2] / 2],\n",
    "                  [-scale[0] / 2, +scale[1] / 2, scale[2] / 2],\n",
    "                  [-scale[0] / 2, +scale[1] / 2, -scale[2] / 2],\n",
    "                  [+scale[0] / 2, -scale[1] / 2, scale[2] / 2],\n",
    "                  [+scale[0] / 2, -scale[1] / 2, -scale[2] / 2],\n",
    "                  [-scale[0] / 2, -scale[1] / 2, scale[2] / 2],\n",
    "                  [-scale[0] / 2, -scale[1] / 2, -scale[2] / 2]]) + shift\n",
    "    else:\n",
    "        bbox_3d = np.array([[scale / 2, +scale / 2, scale / 2],\n",
    "                  [scale / 2, +scale / 2, -scale / 2],\n",
    "                  [-scale / 2, +scale / 2, scale / 2],\n",
    "                  [-scale / 2, +scale / 2, -scale / 2],\n",
    "                  [+scale / 2, -scale / 2, scale / 2],\n",
    "                  [+scale / 2, -scale / 2, -scale / 2],\n",
    "                  [-scale / 2, -scale / 2, scale / 2],\n",
    "                  [-scale / 2, -scale / 2, -scale / 2]]) +shift\n",
    "\n",
    "    bbox_3d = bbox_3d.transpose()\n",
    "    return bbox_3d\n",
    "\n",
    "\n",
    "def draw(img, imgpts, axes, color):\n",
    "    imgpts = np.int32(imgpts).reshape(-1, 2)\n",
    "\n",
    "\n",
    "    # draw ground layer in darker color\n",
    "    color_ground = (int(color[0]), int(color[1]), int(color[2]))\n",
    "    for i, j in zip([4, 5, 6, 7],[5, 7, 4, 6]):\n",
    "        img = cv2.line(img, tuple(imgpts[i]), tuple(imgpts[j]), color_ground, 3)\n",
    "\n",
    "\n",
    "    # draw pillars in blue color\n",
    "    color_pillar = (int(color[0]), int(color[1]), int(color[2]))\n",
    "    for i, j in zip(range(4),range(4,8)):\n",
    "        img = cv2.line(img, tuple(imgpts[i]), tuple(imgpts[j]), color_pillar, 3)\n",
    "\n",
    "    \n",
    "    # finally, draw top layer in color\n",
    "    for i, j in zip([0, 1, 2, 3],[1, 3, 0, 2]):\n",
    "        img = cv2.line(img, tuple(imgpts[i]), tuple(imgpts[j]), color, 3)\n",
    "\n",
    "\n",
    "    # draw axes\n",
    "    img = cv2.line(img, tuple(axes[0]), tuple(axes[1]), (0, 0, 255), 3)  # z\n",
    "    img = cv2.line(img, tuple(axes[0]), tuple(axes[3]), (255, 0, 0), 3)  # x\n",
    "    img = cv2.line(img, tuple(axes[0]), tuple(axes[2]), (0, 255, 0), 3) ## y last\n",
    "\n",
    "\n",
    "    return img\n",
    "\n",
    "\n",
    "def transform_coordinates_3d(coordinates, RT):\n",
    "    \"\"\"\n",
    "    Input: \n",
    "        coordinates: [3, N]\n",
    "        RT: [4, 4]\n",
    "    Return \n",
    "        new_coordinates: [3, N]\n",
    "    \"\"\"\n",
    "    assert coordinates.shape[0] == 3\n",
    "    coordinates = np.vstack([coordinates, np.ones((1, coordinates.shape[1]), dtype=np.float32)])\n",
    "    new_coordinates = RT @ coordinates\n",
    "    new_coordinates = new_coordinates[:3, :]/new_coordinates[3, :]\n",
    "    return new_coordinates\n",
    "\n",
    "\n",
    "def calculate_2d_projections(coordinates_3d, intrinsics):\n",
    "    \"\"\"\n",
    "    Input: \n",
    "        coordinates: [3, N]\n",
    "        intrinsics: [3, 3]\n",
    "    Return \n",
    "        projected_coordinates: [N, 2]\n",
    "    \"\"\"\n",
    "    projected_coordinates = intrinsics @ coordinates_3d\n",
    "    projected_coordinates = projected_coordinates[:2, :] / projected_coordinates[2, :]\n",
    "    projected_coordinates = projected_coordinates.transpose()\n",
    "    projected_coordinates = np.array(projected_coordinates, dtype=np.int32)\n",
    "\n",
    "    return projected_coordinates\n",
    "\n",
    "def draw_pose(img, intrinsics, rot, center, scale, color=(255, 0, 0)):\n",
    "    mat = np.eye(4)\n",
    "    scale_norm = np.linalg.norm(scale)\n",
    "    mat[:3, :3] = rot * scale_norm\n",
    "    mat[:3, -1] = center\n",
    "    \n",
    "    xyz_axis = 0.3 * np.array([[0, 0, 0], [0, 0, 1], [0, 1, 0], [1, 0, 0]]).transpose()\n",
    "    transformed_axes = transform_coordinates_3d(xyz_axis, mat)\n",
    "    projected_axes = calculate_2d_projections(transformed_axes, intrinsics)\n",
    "\n",
    "    bbox_3d = get_3d_bbox(scale / scale_norm, 0)\n",
    "    transformed_bbox_3d = transform_coordinates_3d(bbox_3d, mat)\n",
    "    projected_bbox = calculate_2d_projections(transformed_bbox_3d, intrinsics)\n",
    "    draw_image_bbox = draw(img, projected_bbox, projected_axes, color)\n",
    "    return draw_image_bbox"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm.notebook import tqdm\n",
    "from scipy import stats\n",
    "root = 'sample_data'\n",
    "scene_id = 99\n",
    "split = 'val_inst'\n",
    "rgb_suffix = 'jpg' if 'pbr' in split else 'png'\n",
    "img_fns = list(Path(f'{root}/{split}/{scene_id:06d}/rgb/').glob(f'*.{rgb_suffix}'))\n",
    "for img_id in sorted(img_fns):\n",
    "    img_id = int(img_id.stem)\n",
    "    r = OffscreenRenderer(viewport_width=1280, viewport_height=720)\n",
    "    \n",
    "    img = cv2.imread(f'{root}/{split}/{scene_id:06d}/rgb/{img_id:06d}.{rgb_suffix}')[..., ::-1]\n",
    "    nocs = cv2.imread(f'{root}/{split}/{scene_id:06d}/rgb_nocs/{img_id:06d}.png')[..., ::-1]\n",
    "    depth = cv2.imread(f'{root}/{split}/{scene_id:06d}/depth/{img_id:06d}.png', cv2.IMREAD_UNCHANGED) / (10000. if 'pbr' in split else 1000.0)\n",
    "\n",
    "    scene_gts = json.load(open(f'{root}/{split}/{scene_id:06d}/scene_gt.json'))\n",
    "    scene_camera = json.load(open(f'{root}/{split}/{scene_id:06d}/scene_camera.json'))\n",
    "    scene_gt = scene_gts[str(img_id)]\n",
    "    intrinsic = np.array(scene_camera[str(img_id)]['cam_K']).reshape(3, 3)\n",
    "\n",
    "    scene = Scene(ambient_light=np.array([1, 1, 1, 1.0]), bg_color=np.zeros((3,)))\n",
    "    direc_l = DirectionalLight(color=np.ones(3), intensity=5)\n",
    "    spot_l = SpotLight(color=np.ones(3), intensity=5, innerConeAngle=np.pi/16, outerConeAngle=np.pi/6)\n",
    "\n",
    "    cam_pose = np.eye(4)\n",
    "    cam = IntrinsicsCamera(intrinsic[0, 0], intrinsic[1, 1], intrinsic[0, 2], intrinsic[1, 2])\n",
    "    scene.add(cam, pose=cam_pose)\n",
    "    scene.add(direc_l, pose=cam_pose)\n",
    "    scene.add(spot_l, pose=cam_pose)\n",
    "\n",
    "    rots = []\n",
    "    centers = []\n",
    "    center_offsets = []\n",
    "    scales = []\n",
    "    for gt in scene_gt:\n",
    "        obj_id = gt['obj_id']\n",
    "        print(obj_id)\n",
    "        R = np.array(gt['cam_R_m2c']).reshape(3, 3)\n",
    "        t = np.array(gt['cam_t_m2c']) / 1000\n",
    "\n",
    "        pose = np.eye(4)\n",
    "        pose[:3, :3] = R\n",
    "        pose[:3, 3] = t\n",
    "        pose[1:3] *= -1\n",
    "\n",
    "        # if 545 <= obj_id <= 692:\n",
    "        #     mesh = trimesh.load(f'{root}/models_artic_recentered/obj_{obj_id:06d}.ply')\n",
    "        # else:\n",
    "        mesh = trimesh.load(f'{root}/models/obj_{obj_id:06d}.ply')\n",
    "        mesh.apply_scale(1e-3)\n",
    "        vertices = mesh.vertices\n",
    "        bounds = np.stack([np.min(vertices, 0), np.max(vertices, 0)])\n",
    "        center = bounds.mean(0)\n",
    "        center_offsets.append(center)\n",
    "        mesh.apply_transform(pose)\n",
    "        \n",
    "        scale = np.array([model_info[str(obj_id)]['size_x'], model_info[str(obj_id)]['size_y'], model_info[str(obj_id)]['size_z']]) / 1000.\n",
    "        rots.append(R)\n",
    "        centers.append(t)\n",
    "        scales.append(scale)\n",
    "\n",
    "        scene.add(Mesh.from_trimesh(mesh), pose=np.eye(4))\n",
    "    \n",
    "    plt.figure(figsize=(12, 9))\n",
    "    plt.imshow(img)\n",
    "    plt.axis('off')\n",
    "    plt.show()\n",
    "    \n",
    "    rendered, _ = r.render(scene, RenderFlags.SKIP_CULL_FACES)\n",
    "    r.delete()\n",
    "\n",
    "    plt.figure(figsize=(12, 9))\n",
    "    plt.imshow(rendered)\n",
    "    plt.axis('off')\n",
    "    plt.show()\n",
    "    \n",
    "    plt.figure(figsize=(12, 9))\n",
    "    depth_vis = (np.clip((depth - 0.5) / 2., 0, 1) * 255).astype(np.uint8)\n",
    "    plt.imshow(depth_vis, cmap='gray')\n",
    "    plt.axis('off')\n",
    "    plt.show()\n",
    "    \n",
    "    plt.figure(figsize=(12, 9))\n",
    "    plt.imshow(nocs)\n",
    "    plt.axis('off')\n",
    "    plt.show()\n",
    "\n",
    "    pose_vis = np.copy(img)\n",
    "    for i in range(len(rots)):\n",
    "        # since articulated objects are not centered, we shift bounding box for a better visualization\n",
    "        draw_pose(pose_vis, intrinsic, rots[i], centers[i] + rots[i] @ center_offsets[i], scales[i], color=(255, 255, 0))\n",
    "    \n",
    "    plt.figure(figsize=(12, 9))\n",
    "    plt.imshow(pose_vis)\n",
    "    plt.axis('off')\n",
    "    plt.show()\n",
    "    \n",
    "    seg = np.zeros_like(rendered[..., 0])\n",
    "    for i in range(len(rots)):\n",
    "        mask = cv2.imread(f'{root}/{split}/{scene_id:06d}/mask_visib/{img_id:06d}_{i:06d}.png', 0)\n",
    "        seg[mask > 0] = i + 1\n",
    "    \n",
    "    seg_vis = np.copy(img)\n",
    "    for i in range(1, seg.max() + 1):\n",
    "        seg_vis[seg == i] = np.array(palette[i]) * 255\n",
    "    \n",
    "    plt.figure(figsize=(12, 9))\n",
    "    plt.imshow(seg_vis)\n",
    "    plt.axis('off')\n",
    "    plt.show()\n",
    "    \n",
    "    # visualize coco annotations\n",
    "    for suffix in ['', '_inst', '_partcat']:\n",
    "        coco_json_path = f'{root}/{split}/{scene_id:06d}/scene_gt_coco_det_modal{suffix}.json'\n",
    "        # Visualize detections\n",
    "        visualize_coco_detections(coco_json_path, f'{root}/{split}/{scene_id:06d}', tgt_img_id=int(img_id))\n",
    "\n",
    "    break"
   ]
  }
 ],
 "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}