convert_to_bop.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# README"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This conversion script is incomplete. Namely, symmetry information is not yet converted. Please see TODOs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Documentation of Format"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The IPD native format and BOP format uses similar terminology in different ways. Here are some of the differences.\n",
    "\n",
    "- IPD parts are considered BOP objects\n",
    "    - See generated `ipd_part_to_bop_obj_id.json` for mapping\n",
    "- IPD datasets are considered BOP scenes\n",
    "    - See generated `ipd_dataset_to_bop_scene.json` for mapping\n",
    "- IPD objects (referenced by part and instance) are considered BOP ground truth instances.\n",
    "    - See generated `test[camera]/[bop_scene/ipd_obj_to_bop_gt_id.json` for mapping\n",
    "\n",
    "Conversion details:\n",
    "- IPD cameras are considered different BOP `split_type`s\n",
    "- IPD lighting conditions are specified by different `test_targets`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```\n",
    "/ipd\n",
    "######## BASE ZIP\n",
    "├─ camera_photoneo.json\n",
    "├─ camera_basler_hr.json\n",
    "├─ camera_basler_lr.json\n",
    "├─ camera_flir_polar.json\n",
    "├─ ipd_part_to_bop_obj_id.json\n",
    "\t- mapping from part name to BOP OBJ_ID\n",
    "├─ ipd_dataset_to_bop_scene.json\n",
    "\t- mapping from dataset_id / background to BOP SCENE_ID\n",
    "├─ test_targets_bop19_[all, room, day, spot].json\n",
    "\t- instances for each object in each scene, in each dataset, for different subsets of lighting conditions\n",
    "########\n",
    "\n",
    "######## MODELS ZIP\n",
    "├─ models\n",
    "│  ├─ models_info.json\n",
    "│  ├─ obj_OBJ_ID.ply\n",
    "├─ models_stl\n",
    "│  ├─ models_info.json\n",
    "│  ├─ obj_OBJ_ID.stl\n",
    "########\n",
    "\n",
    "######## PHOTONEO ZIP\n",
    "├─ test_photoneo\n",
    "│  ├─ BOP SCENE_ID\n",
    "│  │  ├─ scene_camera.json\n",
    "\t\t- camera info for each IMG_ID\n",
    "│  │  ├─ scene_gt.json\n",
    "\t\t- List[6D pose and OBJ_ID in GT_ID order] for each IMG_ID\n",
    "│  │  ├─ scene_gt_info.json\n",
    "\t\t- List[bounding boxes in GT_ID order] for each IMG_ID\n",
    "│  │  ├─ depth\n",
    "│  │  │  ├─ IMGID.png\n",
    "│  │  ├─ mask\n",
    "│  │  │  ├─ IMGID_GTID.png\n",
    "│  │  ├─ mask_visib\n",
    "│  │  │  ├─ IMGID_GTID.png\n",
    "│  │  ├─ rgb #if multiple resolutions, combined as HDR image\n",
    "│  │  │  ├─ IMGID.png\n",
    "######## \n",
    "\n",
    "######## [OTHER CAMERAS] ZIP\n",
    "...\n",
    "########\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## To Install"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Option 1: Install Dev Deps via PDM\n",
    "1. Install pdm \n",
    "2. Clone `ipd` repo \n",
    "3. Sync `bop_toolkit` submodule: `git submodule update --init --recursive`\n",
    "4. `pdm install`\n",
    "\n",
    "Note: `bop_toolkit` should be an editable install!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Option 2: Manual Install via pip"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### IPD Toolkit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !python3 -m pip install -e ."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### BOP Toolkit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !git clone git@github.com:thodan/bop_toolkit.git\n",
    "# !python3 -m pip install -r bop_toolkit/requirements.txt -e bop_toolkit/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Other"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !python3 -m pip install open3d, pymeshlab"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Begin Conversion"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Imports & Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, shutil\n",
    "import numpy as np\n",
    "import open3d as o3d\n",
    "import pymeshlab\n",
    "import cv2\n",
    "import json\n",
    "from collections import defaultdict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "from intrinsic_ipd import IPDReader\n",
    "import intrinsic_ipd.constants"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bop_toolkit.bop_toolkit_lib import misc, inout"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "logging.basicConfig(\n",
    "    level=logging.INFO,  # Set the logging level (INFO, DEBUG, WARNING, etc.)\n",
    "    format=\"%(asctime)s - %(name)s - %(levelname)s - %(filename)s - %(funcName)s - %(message)s\",\n",
    "    datefmt=\"%Y-%m-%d %H:%M:%S\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def json_load_if_exists(path_to_json, default):\n",
    "    if os.path.exists(path_to_json):\n",
    "        with open(path_to_json, 'r') as fp:\n",
    "            return json.load(fp, object_hook = lambda d: {int(k) \n",
    "                         if k.lstrip('-').isdigit() else k: v for k, v in d.items()})\n",
    "    else:\n",
    "        return default"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Read Dataset & Setup Destinations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "camera = intrinsic_ipd.constants.IPDCamera.PHOTONEO #TODO: CHANGE ME!\n",
    "\n",
    "#Choose dataset\n",
    "bop_scene = 0 # TODO: CHANGE ME!\n",
    "datasets = intrinsic_ipd.constants.DATASET_IDS\n",
    "dataset = datasets[bop_scene]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "corner_bracket5.stl: 0.00B [00:00, ?B/s]\n",
      "ERROR:root:url: https://storage.googleapis.com/akasha-public/industrial_plenoptic_dataset/cad_models/corner_bracket5.stl failed to download with error: HTTP Error 404: Not Found\n"
     ]
    }
   ],
   "source": [
    "# Read/Download dataset\n",
    "reader = IPDReader(\"./datasets\", dataset, camera, lighting=intrinsic_ipd.constants.IPDLightCondition.ALL, download=True) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make ipd dest\n",
    "bop_dest = \"./bop_datasets\"\n",
    "ipd_dest = os.path.join(bop_dest, \"ipd\")\n",
    "os.makedirs(ipd_dest, exist_ok=True)\n",
    "\n",
    "# Make camera dest\n",
    "camera_dest = os.path.join(ipd_dest, f\"test_{camera.name.lower()}\")\n",
    "os.makedirs(camera_dest, exist_ok=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make bop_scene map\n",
    "bop_scene_map_file = os.path.join(ipd_dest, \"ipd_dataset_to_bop_scene.json\")\n",
    "bop_scene_map = json_load_if_exists(bop_scene_map_file, {did : i for i, did in enumerate(intrinsic_ipd.constants.DATASET_IDS)})\n",
    "inout.save_json(bop_scene_map_file, bop_scene_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make dataset (bop_scene) dest\n",
    "scene_dest = os.path.join(camera_dest, f\"{bop_scene:06}\")\n",
    "os.makedirs(scene_dest, exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Convert Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make models dest\n",
    "models_stl_dest = os.path.join(ipd_dest, \"models_stl\")\n",
    "models_dest = os.path.join(ipd_dest, \"models\")\n",
    "# models_eval_dest = os.path.join(ipd_dest, \"models_eval\")\n",
    "os.makedirs(models_stl_dest, exist_ok=True)\n",
    "os.makedirs(models_dest, exist_ok=True)\n",
    "# os.makedirs(models_eval_dest, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "# CONVERT USING PYMESHLAB\n",
    "def pymeshlab_stl_to_ply(in_stl_path, out_ply_path):\n",
    "    ms = pymeshlab.MeshSet()\n",
    "    ms.load_new_mesh(in_stl_path)\n",
    "    ms.save_current_mesh(out_ply_path,\n",
    "                         binary = False,\n",
    "                         save_vertex_normal = True\n",
    "                         )\n",
    "\n",
    "# Alternative: CONVERT USING OPEN3D\n",
    "def o3d_stl_to_ply(in_stl_path, out_ply_path, sample=False, num_points = 10000):\n",
    "    mesh = o3d.io.read_triangle_mesh(in_stl_path)\n",
    "    if sample:\n",
    "        cloud = mesh.sample_points_uniformly(num_points, use_triangle_normal=True)\n",
    "    else:\n",
    "        mesh.compute_vertex_normals()\n",
    "        mesh.paint_uniform_color((1, 0.75, 0))\n",
    "        mesh.compute_vertex_normals()\n",
    "        cloud = o3d.geometry.PointCloud()\n",
    "        cloud.points = mesh.vertices\n",
    "        cloud.normals = mesh.vertex_normals \n",
    "    o3d.io.write_point_cloud(out_ply_path, cloud, write_ascii=True)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "TEST = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NO STL FILE FOUND FOR: corner_bracket5 at ./datasets/models/corner_bracket5.stl\n"
     ]
    }
   ],
   "source": [
    "models_info = {}\n",
    "sample = False\n",
    "parts = intrinsic_ipd.constants.PART_NAMES\n",
    "if TEST:\n",
    "    parts = ['gear2']\n",
    "\n",
    "part2obid = {}\n",
    "for obj_id, part in enumerate(parts):\n",
    "    obj_id += 1 # index starts with 1\n",
    "    part2obid[part] = obj_id\n",
    "\n",
    "    ######## SAVE MODEL\n",
    "    # copy stl model to models_stl\n",
    "    model_stl_path = os.path.join(reader.root, 'models', f'{part}.stl')\n",
    "    if not os.path.exists(model_stl_path): \n",
    "        print(f\"NO STL FILE FOUND FOR: {part} at {model_stl_path}\")\n",
    "        continue\n",
    "    dst = os.path.join(models_stl_dest, f'obj_{obj_id:06}.ply')\n",
    "    shutil.copyfile(model_stl_path, dst)\n",
    "\n",
    "    # create ply model in models\n",
    "    model_ply_path = os.path.join(models_dest, f'obj_{obj_id:06}.ply')\n",
    "    pymeshlab_stl_to_ply(model_stl_path, model_ply_path)\n",
    "    \n",
    "    # TODO: create eval ply model in models_eval\n",
    "    # - 'Uniformly' resamples and decimates 3D object models for evaluation. \n",
    "    # - See bop_toolkit/scripts/remesh_models_for_eval.py\n",
    "    # !!!! DOES NOT WORK !!!! some error with pymeshlab. Try bop_toolkit/scripts/remesh_models_for_eval.py\n",
    "    if False:\n",
    "        ms = pymeshlab.MeshSet()\n",
    "        ms.load_new_mesh(model_ply_path)\n",
    "        ms.load_filter_script('./bop_toolkit/scripts/meshlab_scripts/remesh_for_eval_cell=0.25.mlx')\n",
    "        ms.apply_filter_script()\n",
    "        ms.save_current_mesh(os.path.join(models_eval_dest, f'obj_{obj_id:06}.ply'),\n",
    "                            binary = False,\n",
    "                            save_vertex_normal = True\n",
    "                            )\n",
    "\n",
    "    ######## SAVE MODEL INFO: see bop_toolkit/scripts/calc_model_info.py\n",
    "    model = inout.load_ply(model_ply_path)\n",
    "    ref_pt = list(map(float, model[\"pts\"].min(axis=0).flatten()))\n",
    "    size = list(map(float, (model[\"pts\"].max(axis=0) - ref_pt).flatten()))\n",
    "    diameter = misc.calc_pts_diameter(model[\"pts\"])\n",
    "\n",
    "    model_info = {\n",
    "        \"min_x\": ref_pt[0],\n",
    "        \"min_y\": ref_pt[1],\n",
    "        \"min_z\": ref_pt[2],\n",
    "        \"size_x\": size[0],\n",
    "        \"size_y\": size[1],\n",
    "        \"size_z\": size[2],\n",
    "        \"diameter\": diameter,\n",
    "    }\n",
    "\n",
    "    # TODO: process symmetries\n",
    "    # see: https://github.com/thodan/bop_toolkit/blob/97badc48dae87d03fa86c0f4ccce94ffdaaae4c5/bop_toolkit_lib/misc.py#L47\n",
    "    \n",
    "    symm = reader._get_symm_params(part)\n",
    "    # print(symm)\n",
    "\n",
    "    # # TODO: list of continuous symmetries arrays\n",
    "    # model_info['symmetries_discrete'] = [\n",
    "    #     np.eye(4).flatten() # 4x4 matrix flattened!\n",
    "    # ]\n",
    "\n",
    "    # # TODO: list of continuous symmetries dictionaries\n",
    "    # model_info['symmetries_continuous'] = [\n",
    "    #     {\n",
    "    #         \"axis\": [ 0, 0, 1 ],\n",
    "    #         \"offset\": [ 0, 0, 0 ]\n",
    "    #     }\n",
    "    # ]\n",
    "    models_info[obj_id] = model_info\n",
    "\n",
    "inout.save_json(os.path.join(models_dest, 'models_info.json'), models_info)\n",
    "inout.save_json(os.path.join(models_stl_dest, 'models_info.json'), models_info)\n",
    "# inout.save_json(os.path.join(models_eval_dest, 'models_info.json'), models_info) # TODO: create eval ply\n",
    "\n",
    "inout.save_json(os.path.join(ipd_dest, 'ipd_part_to_bop_obj_id.json'), part2obid)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Convert IPD object into BOP ground truth instances"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'gear1': {0: 0, 1: 1, 2: 2, 3: 3},\n",
       " 'pegboard_basket': {0: 4},\n",
       " 'u_bolt': {0: 5, 1: 6, 2: 7}}"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def get_bop_gt_id_map(reader, scene_dest):\n",
    "    map_json = os.path.join(scene_dest, 'ipd_obj_to_bop_gt_id.json')\n",
    "\n",
    "    bop_gt_id_map = json_load_if_exists(map_json, defaultdict(dict))\n",
    "\n",
    "    for i, obj in enumerate(reader.objects):\n",
    "          bop_gt_id_map[obj[0]][obj[1]] = i\n",
    "    with open(map_json, 'w') as fp:\n",
    "            json.dump(bop_gt_id_map, fp, sort_keys=True, indent=4)\n",
    "    return bop_gt_id_map\n",
    "\n",
    "get_bop_gt_id_map(reader, scene_dest)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Move/Convert Images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make image dests\n",
    "depth_dest = os.path.join(scene_dest, 'depth')\n",
    "mask_dest = os.path.join(scene_dest, 'mask_ipd')\n",
    "rgb_dest = os.path.join(scene_dest, 'rgb')\n",
    "\n",
    "os.makedirs(depth_dest, exist_ok=True)\n",
    "os.makedirs(mask_dest, exist_ok=True)\n",
    "os.makedirs(rgb_dest, exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Move RBGD images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "def merge_exposures(img_paths):\n",
    "    img_list = [cv2.imread(path) for path in img_paths]\n",
    "    exposure_times = np.array([1, 30, 80, 200], dtype=np.float32)\n",
    "    merge_debevec = cv2.createMergeDebevec()\n",
    "    hdr_debevec = merge_debevec.process(img_list, times=exposure_times.copy())\n",
    "    return hdr_debevec"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For photoneo, will move & rename rgb and depth files.\n",
    "\n",
    "For other cameras, will merge into an hdr photo and save"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "for bop_image_id in reader.scenes.keys():\n",
    "    # move or merge rgb & depth photos\n",
    "    if reader.camera == intrinsic_ipd.IPDCamera.PHOTONEO:\n",
    "        from_path = reader._get_img_file(bop_image_id, intrinsic_ipd.IPDImage.PHOTONEO_DEPTH)\n",
    "        to_path = os.path.join(depth_dest, f'{bop_image_id:06}.png')\n",
    "        shutil.copy(from_path, to_path)\n",
    "\n",
    "        from_path = reader._get_img_file(bop_image_id, intrinsic_ipd.IPDImage.PHOTONEO_HDR)\n",
    "        to_path = os.path.join(rgb_dest, f'{bop_image_id:06}.png')\n",
    "        shutil.copy(from_path, to_path)\n",
    "    else:\n",
    "        img_paths = [reader._get_img_file(bop_image_id, image_type) for image_type in reader.camera.images]\n",
    "        hdr_photo = merge_exposures(img_paths)\n",
    "        to_path = os.path.join(rgb_dest, f'{bop_image_id:06}.png')\n",
    "        cv2.imwrite(to_path, hdr_photo)\n",
    "    \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Move masks based on ground truth id"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Move and rename masks based on ground truth id."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "# move and rename mask photos\n",
    "bop_gt_id_map = None\n",
    "for bop_image_id in reader.scenes.keys():\n",
    "    for object in reader.objects:\n",
    "        part, instance = object\n",
    "        bop_gt_id_map = get_bop_gt_id_map(reader, scene_dest)\n",
    "        bop_gt_id = bop_gt_id_map[part][instance]\n",
    "        _, ipd_mask_path = reader.get_mask(bop_image_id, part, instance, return_path=True)\n",
    "        bop_mask_path = os.path.join(mask_dest, f'{bop_image_id:06}_{bop_gt_id:06}.png')\n",
    "        shutil.copy(ipd_mask_path, bop_mask_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Process Labels"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### scene_camera.json\n",
    "cam_K, cam_R_w2c, cam_t_w2c, depth_scale, elev, mode \n",
    "for each bop image (ipd scene)\n",
    "\n",
    "TODO: Read and add to this file when loading other lighting conditions??? (or will that be a new type?)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "scene_camera_path = os.path.join(scene_dest, 'scene_camera.json')\n",
    "camera_path = os.path.join(ipd_dest, f\"camera_{reader.camera.name.lower()}.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [],
   "source": [
    "w2c = np.linalg.inv(reader.cam_c2w)\n",
    "camera_info = {\n",
    "    'cam_K': reader.cam_K.flatten().tolist(),\n",
    "    'cam_R_w2c': w2c[:3,:3].flatten().tolist(),\n",
    "    'cam_t_w2c': w2c[:3,3].flatten().tolist(),\n",
    "}\n",
    "\n",
    "if reader.camera is intrinsic_ipd.constants.IPDCamera.PHOTONEO:\n",
    "    camera_info['depth_scale']= 1.0\n",
    "\n",
    "image = reader.get_img(list(reader.scenes.keys())[0])\n",
    "height, width = image.shape[:2]\n",
    "\n",
    "scene_camera = {bop_image_id : camera_info  for bop_image_id in reader.scenes.keys()}\n",
    "inout.save_json(scene_camera_path, scene_camera)\n",
    "inout.save_json(camera_path, {\n",
    "  \"cx\": reader.cam_K[0, 2],\n",
    "  \"cy\": reader.cam_K[1, 2],\n",
    "  \"depth_scale\": 1.0,\n",
    "  \"fx\": reader.cam_K[0, 0],\n",
    "  \"fy\": reader.cam_K[1, 1],\n",
    "  \"height\": height,\n",
    "  \"width\": width\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### scene_gt.json\n",
    "map IMG_ID to List[6D pose and OBJ_ID in GT_ID order]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "scene_gt_path = os.path.join(scene_dest, 'scene_gt.json')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_gt_info(reader, ipd_scene_id, ipd_dest, scene_dest):\n",
    "    with open(os.path.join(ipd_dest, \"ipd_part_to_bop_obj_id.json\"), 'r') as fp:\n",
    "        bop_obj_id_map = json.load(fp)\n",
    "    ipd_objects = reader.objects\n",
    "    o2c = reader.o2c.sel(scene=ipd_scene_id)\n",
    "    gt_info = {}\n",
    "    bop_gt_id_map = get_bop_gt_id_map(reader, scene_dest)\n",
    "    for part, instance in ipd_objects:\n",
    "        gt_id = bop_gt_id_map[part][instance]\n",
    "        gt_o2c = o2c.sel(part=part, instance=instance).data\n",
    "        obj_id = bop_obj_id_map[part]\n",
    "        gt_info[gt_id] = {\n",
    "            'obj_id': obj_id, \n",
    "            'cam_R_m2c': gt_o2c[:3,:3].flatten().tolist(),\n",
    "            'cam_t_m2c': gt_o2c[:3, 3].flatten().tolist(),\n",
    "            'gt_id': gt_id,\n",
    "            'ipd_object': (part, instance)\n",
    "        }\n",
    "    gt_keys = [int(k) for k in gt_info.keys()]\n",
    "    max_gt_id = max(gt_keys)\n",
    "    return [gt_info.get(gt_id, {}) for gt_id in range(max_gt_id)]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "scene_gt = {bop_image_id: get_gt_info(reader, bop_image_id, ipd_dest, scene_dest) for bop_image_id in reader.scenes.keys()}\n",
    "inout.save_json(scene_gt_path, scene_gt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### test_targets.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "import itertools, operator\n",
    "\n",
    "bop_obj_id_map = json_load_if_exists(os.path.join(ipd_dest, 'ipd_part_to_bop_obj_id.json'), {})\n",
    "\n",
    "\n",
    "conditions = [intrinsic_ipd.constants.IPDLightCondition.ALL,\n",
    "              intrinsic_ipd.constants.IPDLightCondition.DAY, \n",
    "              intrinsic_ipd.constants.IPDLightCondition.ROOM,\n",
    "              intrinsic_ipd.constants.IPDLightCondition.SPOT]\n",
    "\n",
    "bop_scene = json_load_if_exists(os.path.join(ipd_dest, 'ipd_dataset_to_bop_scene.json'), {})[reader.dataset_id]\n",
    "for condition in conditions:\n",
    "    target_file = os.path.join(ipd_dest, f'test_targets_bop19_{condition.name.lower()}.json')\n",
    "    targets = json_load_if_exists(target_file, [])    \n",
    "    for part, objects in itertools.groupby(reader.objects, operator.itemgetter(0)):\n",
    "        obj_id = bop_obj_id_map[part]\n",
    "        inst_count = len(list(objects))\n",
    "        for bop_image_id in condition.scenes:\n",
    "            target = {\n",
    "                \"im_id\": bop_image_id,\n",
    "                \"obj_id\": obj_id,\n",
    "                \"inst_count\": inst_count,\n",
    "                \"scene_id\": bop_scene,\n",
    "            }\n",
    "            if target not in targets:\n",
    "                targets.append(target)\n",
    "    \n",
    "    inout.save_json(target_file, targets)\n",
    "\n",
    "    targets_24_file = os.path.join(ipd_dest, f'test_targets_bop24_{condition.name.lower()}.json')\n",
    "    targets_24 = json_load_if_exists(targets_24_file, [])\n",
    "    for bop_image_id in condition.scenes:\n",
    "        target_24 = {\n",
    "                \"im_id\": bop_image_id,\n",
    "                \"scene_id\": bop_scene,\n",
    "            }\n",
    "        if target_24 not in targets_24:\n",
    "            targets_24.append(target_24)\n",
    "    inout.save_json(targets_24_file, targets_24)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Repeat above steps for all dataset ids and camera"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6.Run BOP scripts to generate rest of dataset info "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To run the bop_toolkit scripts, need to make some edits:\n",
    "\n",
    "Changelog to `bop_toolkit` as reflected in @carynbear's fork:\n",
    "- in `bop_toolkit/bop_toolkit_lib/dataset_params.py`\n",
    "    - CHANGED: added `ipd` params throughout\n",
    "    - TODO: indicate which objects (parts) have symmetry.\n",
    "    - TODO: get sizes of images for other cameras (photoneo done)\n",
    "    - TODO: calculate depth_range, azimuth_range, elev_range \n",
    "- in `bop_toolkit/bop_toolkit_lib/config.py`\n",
    "    - CHANGED: `output_path`\n",
    "- in `bop_toolkit/scripts/calc_gt_info.py`\n",
    "    - CHANGED: run with `ipd` and `vis`\n",
    "    - CHANGED: try catch to skip missing cad models\n",
    "- in `bop_toolkit/scripts/calc_gt_masks.py`\n",
    "    - CHANGED: run with `ipd`\n",
    "    - CHANGED: try catch to skip missing cad models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"BOP_PATH\"] = bop_dest"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calculate scene_gt_info.json\n",
    "map IMG_ID to List[bounding boxes in GT_ID order]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "9/20|22:18:08: Initializing renderer...\n",
      "vispy uses app: egl, gl: gl2\n",
      "9/20|22:18:10: Calculating GT info - dataset: ipd (test, None), scene: 0, im: 0\n"
     ]
    }
   ],
   "source": [
    "!python3 ./bop_toolkit/scripts/calc_gt_info.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate masks from gt labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "9/20|22:27:31: Initializing renderer...\n",
      "vispy uses app: egl, gl: gl2\n",
      "9/20|22:27:32: Calculating masks - dataset: ipd (test, None), scene: 0, im: 0\n"
     ]
    }
   ],
   "source": [
    "!python3 ./bop_toolkit/scripts/calc_gt_masks.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate models_evel with pymeshlab\n",
    "- TODO: Specify the meshlab_server_path to the executable in `bop_toolkit/bop_toolkit_lib/config.py`\n",
    "- TODO: Edit the script to run with `ipd` dataset!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !python3 ./bop_toolkit/scripts/remesh_models_for_eval.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate training images\n",
    "- TODO: INCOMPLETE!!! Edit and try running!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !python3 ./bop_toolkit/scripts/render_train_imgs.py"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}