dfnet.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DenseFusion Submission"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from learning.densefusion import DenseFuseNet\n",
    "from learning.loss_batch import quat_to_rot\n",
    "from learning.utils import OBJ_NAMES, OBJ_NAMES_TO_IDX, IDX_TO_OBJ_NAMES\n",
    "from icp import run_icp, rigid_transform, R_and_T, to_T\n",
    "\n",
    "from pathlib import Path\n",
    "from learning.load import PoseDataset\n",
    "import pickle\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "\n",
    "def load_model(model: torch.nn.Module, optimizer: torch.optim.Optimizer, load_path, device=torch.device('cpu')):\n",
    "    checkpoint = torch.load(load_path, map_location=device)\n",
    "    model.load_state_dict(checkpoint['model'])\n",
    "    optimizer.load_state_dict(checkpoint['optimizer'])\n",
    "    return model, optimizer\n",
    "\n",
    "def free(x: torch.Tensor) -> np.ndarray:\n",
    "    return x.detach().cpu().numpy()\n",
    "\n",
    "def pred_over_raw_data(\n",
    "        output_json_name='dfnet_pred.json',\n",
    "        processed_data_dir='processed_data_unet',\n",
    "        raw_data_dir='raw_data_all',\n",
    "        checkpoint='select-checkpoints/dfnet.pt',\n",
    "        refine=False,\n",
    "        icp_thresh=1e-30,\n",
    "    ):\n",
    "    def get_stripped_lines(fp, levels=[1, 2, 3]):\n",
    "        return [x.strip() for x in open(fp, 'r').readlines() if int(x[0]) in levels]\n",
    "    \n",
    "    device = torch.device('cuda')\n",
    "\n",
    "    num_objects = len(OBJ_NAMES)\n",
    "    dfnet = torch.nn.DataParallel(DenseFuseNet(num_objects)).to(device)\n",
    "    optimizer = torch.optim.Adam(dfnet.parameters(), lr=1e-4)\n",
    "\n",
    "    dfnet, optimizer = load_model(dfnet, optimizer, checkpoint, device=torch.device('cuda'))\n",
    "    dfnet = dfnet.eval()\n",
    "\n",
    "    raw_data_dir = Path(raw_data_dir)\n",
    "    raw_test_dir = raw_data_dir / 'testing_data'\n",
    "    raw_test_obj_dir = raw_test_dir / 'v2.2'\n",
    "    processed_data_dir = Path(processed_data_dir)\n",
    "    processed_test_dir = processed_data_dir / 'test'\n",
    "\n",
    "    test_scene_names = get_stripped_lines(raw_test_dir / 'test.txt')\n",
    "    test_ds = PoseDataset(data_dir=processed_test_dir, train=False, cloud=True, rgb=True, model=True, choose=True)\n",
    "\n",
    "    data_point_num = 0\n",
    "    all_data = dict()\n",
    "    pbar = tqdm(test_scene_names)\n",
    "    for scene_name in pbar:\n",
    "        meta_path = raw_test_obj_dir / f'{scene_name}_meta.pkl'\n",
    "        meta = pickle.load(open(meta_path, 'rb'))\n",
    "\n",
    "        scene_data = dict(poses_world=[None] * 79)\n",
    "\n",
    "        for obj_id, obj_name in zip(meta['object_ids'], meta['object_names']):\n",
    "\n",
    "            pbar.set_description(f'dp_num={data_point_num}')\n",
    "\n",
    "            cloud, rgb, model, choose, obj_idxs = test_ds[data_point_num]\n",
    "            if len(cloud) < 1:\n",
    "                T = np.eye(4)\n",
    "            else:\n",
    "                cloud, rgb, model, choose, obj_idxs = cloud.to(device), rgb.to(device), model.to(device), choose.to(device), obj_idxs.to(device)\n",
    "                cloud, rgb, model, choose, obj_idxs = cloud.unsqueeze(0), rgb.unsqueeze(0), model.unsqueeze(0), choose.unsqueeze(0), obj_idxs.unsqueeze(0)\n",
    "                cloud, rgb, model, choose, obj_idxs = cloud.float(), rgb.float(), model.float(), choose.float(), obj_idxs\n",
    "\n",
    "                cloud_new = cloud.transpose(2, 1)\n",
    "                rgb_new = torch.moveaxis(rgb, -1, 1)\n",
    "                choose_new = choose.view(choose.size(0), -1)\n",
    "                R_quat_pred, t_pred, c_pred = dfnet(cloud_new, rgb_new, choose_new, obj_idxs)\n",
    "\n",
    "                R_pred = quat_to_rot(R_quat_pred)\n",
    "\n",
    "                R_pred_opt = R_pred[0][torch.argmax(c_pred[0])]\n",
    "                t_pred_opt = t_pred[0][torch.argmax(c_pred[0])]\n",
    "\n",
    "                R_pred, t_pred = free(R_pred_opt), free(t_pred_opt)\n",
    "\n",
    "                if refine:\n",
    "                    # icp refinement\n",
    "                    cloud = free(cloud.squeeze(0))\n",
    "                    model = free(model.squeeze(0))\n",
    "\n",
    "                    pred_model = model @ R_pred.T + t_pred\n",
    "                    R_pred_cloud, t_pred_cloud = R_and_T(run_icp(cloud, pred_model, max_attempts=1, max_iters=1000, finish_loop_thresh=icp_thresh, acceptable_thresh=icp_thresh))\n",
    "                    T = rigid_transform(model, (pred_model - t_pred_cloud) @ R_pred_cloud)\n",
    "                else:\n",
    "                    T = to_T(R_pred, t_pred)\n",
    "\n",
    "            scene_data['poses_world'][obj_id] = T.tolist()\n",
    "\n",
    "            data_point_num += 1\n",
    "\n",
    "        all_data[scene_name] = scene_data\n",
    "\n",
    "    import json\n",
    "    with open(output_json_name, 'w') as fp:\n",
    "        json.dump(all_data, fp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## HW2: Levels 1-2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "DenseFusion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "dp_num=1499: 100%|██████████| 200/200 [00:43<00:00,  4.64it/s]\n"
     ]
    }
   ],
   "source": [
    "pred_over_raw_data(output_json_name='dfnet_pred.json', processed_data_dir='processed_data', raw_data_dir='raw_data', checkpoint='select-checkpoints/dfnet.pt', refine=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "DenseFusion + ICP Refinement"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "dp_num=1499: 100%|██████████| 200/200 [00:54<00:00,  3.64it/s]\n"
     ]
    }
   ],
   "source": [
    "pred_over_raw_data(output_json_name='dfnet_pred_refine.json', processed_data_dir='processed_data', raw_data_dir='raw_data', checkpoint='select-checkpoints/dfnet.pt', refine=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## HW3: Levels 1-3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "DenseFusion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "dp_num=4999: 100%|██████████| 600/600 [02:17<00:00,  4.35it/s]\n"
     ]
    }
   ],
   "source": [
    "pred_over_raw_data(output_json_name='dfnet_pred_unet.json', processed_data_dir='processed_data_unet', raw_data_dir='raw_data_all', checkpoint='select-checkpoints/dfnet.pt', refine=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "DenseFusion + ICP Refinement"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "dp_num=4999: 100%|██████████| 600/600 [03:41<00:00,  2.71it/s]\n"
     ]
    }
   ],
   "source": [
    "pred_over_raw_data(output_json_name='dfnet_pred_unet_refine.json', processed_data_dir='processed_data_unet', raw_data_dir='raw_data_all', checkpoint='select-checkpoints/dfnet.pt', refine=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cse275",
   "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.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}