test_seg.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "from PIL import Image\n",
    "import pickle\n",
    "# import utils\n",
    "\n",
    "def load_pickle(filename):\n",
    "    with open(filename, 'rb') as f:\n",
    "        return pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !unzip ./testing_data_final_filtered.zip\n",
    "# !ls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Change these \n",
    "testing_data_dir = \"testing_data_final_filtered/testing_data/v2.2\"\n",
    "split_dir_test = \"testing_data_final_filtered/testing_data/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_split_files(split_name,split_dir,data_dir):\n",
    "    with open(os.path.join(split_dir, f\"{split_name}.txt\"), 'r') as f:\n",
    "        prefix = [os.path.join(data_dir, line.strip()) for line in f if line.strip()]\n",
    "        rgb = [p + \"_color_kinect.png\" for p in prefix]\n",
    "        depth = [p + \"_depth_kinect.png\" for p in prefix]\n",
    "#         label = [p + \"_label_kinect.png\" for p in prefix]\n",
    "        meta = [p + \"_meta.pkl\" for p in prefix]\n",
    "    return rgb, depth,meta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rgb_files, depth_files, meta_files = get_split_files('test',split_dir_test,testing_data_dir)\n",
    "with open(os.path.join(split_dir_test, f\"{'test'}.txt\"), 'r') as f:\n",
    "    scenes = [line.strip()for line in f if line.strip()]\n",
    "# print((scenes))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scenes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib.cm import get_cmap\n",
    "NUM_OBJECTS = 79\n",
    "cmap = get_cmap('rainbow', NUM_OBJECTS)\n",
    "COLOR_PALETTE = np.array([cmap(i)[:3] for i in range(NUM_OBJECTS + 3)])\n",
    "COLOR_PALETTE = np.array(COLOR_PALETTE * 255, dtype=np.uint8)\n",
    "COLOR_PALETTE[-3] = [119, 135, 150]\n",
    "COLOR_PALETTE[-2] = [176, 194, 216]\n",
    "COLOR_PALETTE[-1] = [255, 255, 225]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get familiar with the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rgb = np.array(Image.open(rgb_files[0])) / 255   # convert 0-255 to 0-1\n",
    "depth = np.array(Image.open(depth_files[0])) / 1000   # convert from mm to m\n",
    "# label = np.array(Image.open(label_files[0]))\n",
    "print(rgb.shape)\n",
    "plt.figure(figsize=(15, 10))\n",
    "plt.subplot(1, 3, 1)\n",
    "plt.imshow(rgb)\n",
    "plt.subplot(1, 3, 2)\n",
    "plt.imshow(depth)\n",
    "# plt.subplot(1, 3, 3)\n",
    "# plt.imshow(COLOR_PALETTE[label])  # draw colorful segmentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import Dataset\n",
    "from torch.utils.data import DataLoader\n",
    "import pandas as pd\n",
    "\n",
    "objects_csv = 'testing_data_final_filtered/testing_data/objects_v1.csv'\n",
    "\n",
    "def read_image(img_path):\n",
    "    '''\n",
    "    inputs:\n",
    "    img_path : the location of the image to be read\n",
    "    outputs:\n",
    "    image converted to torch.tensor\n",
    "    '''\n",
    "    image = np.array(Image.open(img_path))\n",
    "#     print(image)\n",
    "    image = torch.from_numpy(image)\n",
    "    return image\n",
    "\n",
    "class mydataset(Dataset):\n",
    "    # define the init method\n",
    "    def __init__(self, img_files, img_dir, scene_names, object_files, transform=None, target_transform = None) -> None:\n",
    "        super().__init__()\n",
    "#         self.target_labels = annotations_files\n",
    "        self.img_dir = img_dir\n",
    "        self.img_files = img_files\n",
    "        self.scenes = scene_names\n",
    "        self.objects = pd.read_csv(object_files)\n",
    "        self.transform = transform\n",
    "        self.target_transform = target_transform\n",
    "\n",
    "    # define the len method\n",
    "    def __len__(self):\n",
    "        return len(self.img_files)\n",
    "\n",
    "    # define the getitem() method\n",
    "    def __getitem__(self,idx):\n",
    "        img_path = self.img_files[idx]\n",
    "        scene_idx = self.scenes[idx]#os.path.join(self.img_dir, self.img_files[idx])\n",
    "#         target_path = self.target_labels[idx]#os.path.join(self.img_dir, self.target_labels[idx])\n",
    "        image = read_image(img_path)/255.0 # divide by 255 or do some normalization using transforms\n",
    "#         label = read_image(target_path)\n",
    "        if self.transform:\n",
    "            image  = self.transform(image)\n",
    "        return image,scene_idx\n",
    "\n",
    "test_data = mydataset(rgb_files,testing_data_dir,scenes,objects_csv)\n",
    "\n",
    "test_dataloader = DataLoader(test_data, batch_size=1, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "\n",
    "class Segmentation(nn.Module): \n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.c1 = nn.Sequential(\n",
    "            nn.Conv2d(3, 64, 3, padding=1, padding_mode=\"reflect\"),\n",
    "            nn.BatchNorm2d(64),\n",
    "            nn.ReLU(),\n",
    "            # nn.Conv2d(64, 64, 3, padding=1, padding_mode=\"reflect\"),\n",
    "            # nn.BatchNorm2d(64),\n",
    "            # nn.ReLU(),\n",
    "        )\n",
    "        self.c2 = nn.Sequential(\n",
    "            nn.Conv2d(64, 128, 3, padding=1, padding_mode=\"reflect\"),\n",
    "            nn.BatchNorm2d(128),\n",
    "            nn.ReLU(),\n",
    "            # nn.Conv2d(128, 128, 3, padding=1, padding_mode=\"reflect\"),\n",
    "            # nn.BatchNorm2d(128),\n",
    "            # nn.ReLU(),\n",
    "        )\n",
    "        self.c3 = nn.Sequential(\n",
    "            nn.Conv2d(128, 256, 3, padding=1, padding_mode=\"reflect\"),\n",
    "            nn.BatchNorm2d(256),\n",
    "            nn.ReLU(),\n",
    "            # nn.Conv2d(256, 256, 3, padding=1, padding_mode=\"reflect\"),\n",
    "            # nn.BatchNorm2d(256),\n",
    "            # nn.ReLU(),\n",
    "        )\n",
    "        self.c4 = nn.Sequential(\n",
    "            nn.Conv2d(256, 512, 3, padding=1, padding_mode=\"reflect\"),\n",
    "            nn.BatchNorm2d(512),\n",
    "            nn.ReLU(),\n",
    "            # nn.Conv2d(512, 512, 3, padding=1, padding_mode=\"reflect\"),\n",
    "            # nn.BatchNorm2d(512),\n",
    "            # nn.ReLU(),\n",
    ")\n",
    "        self.p1 = nn.MaxPool2d(2)\n",
    "        self.p2 = nn.MaxPool2d(2)\n",
    "        self.p3 = nn.MaxPool2d(2)\n",
    "        self.d1 = nn.ConvTranspose2d(128, 64, 2, 2)\n",
    "        self.d2 = nn.ConvTranspose2d(256, 128, 2, 2)\n",
    "        self.d3 = nn.ConvTranspose2d(512, 256, 2, 2)\n",
    "        self.dc1 = nn.Sequential(\n",
    "            nn.Conv2d(128, 64, 3, padding=1, padding_mode=\"reflect\"),\n",
    "            nn.BatchNorm2d(64),\n",
    "            nn.ReLU(),\n",
    "            # nn.Conv2d(64, 64, 3, padding=1, padding_mode=\"reflect\"),\n",
    "            # nn.BatchNorm2d(64),\n",
    "            # nn.ReLU(),\n",
    "            nn.Conv2d(64, 82, 1),\n",
    "#             nn.BatchNorm2d(82),\n",
    "#             nn.ReLU(),\n",
    "        )\n",
    "        self.dc2 = nn.Sequential(\n",
    "            nn.Conv2d(256, 128, 3, padding=1, padding_mode=\"reflect\"),\n",
    "            nn.BatchNorm2d(128),\n",
    "            nn.ReLU(),\n",
    "            # nn.Conv2d(128, 128, 3, padding=1, padding_mode=\"reflect\"),\n",
    "            # nn.BatchNorm2d(128),\n",
    "            # nn.ReLU(),\n",
    ")\n",
    "        self.dc3 = nn.Sequential(\n",
    "            nn.Conv2d(512, 256, 3, padding=1, padding_mode=\"reflect\"),\n",
    "            nn.BatchNorm2d(256),\n",
    "            nn.ReLU(),\n",
    "            # nn.Conv2d(256, 256, 3, padding=1, padding_mode=\"reflect\"),\n",
    "            # nn.BatchNorm2d(256),\n",
    "            # nn.ReLU(),\n",
    ")\n",
    "    def forward(self, x): \n",
    "        x1 = self.c1(x)\n",
    "\n",
    "        x2 = self.c2(self.p1(x1))\n",
    "\n",
    "        x3 = self.c3(self.p2(x2))\n",
    "\n",
    "        x4 = self.c4(self.p3(x3))\n",
    "\n",
    "        y3 = torch.cat([x3, self.d3(x4)], dim=1)\n",
    "\n",
    "        y2 = torch.cat([x2, self.d2(self.dc3(y3))], dim=1)\n",
    "\n",
    "        y1 = torch.cat([x1, self.d1(self.dc2(y2))], dim=1)\n",
    "        output = self.dc1(y1).squeeze(1) \n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy\n",
    "\n",
    "device = \"cuda\"\n",
    "model = torch.load('./model_new_5.pth')\n",
    "model.to(device)\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    for data in test_dataloader:\n",
    "        torch.cuda.empty_cache()\n",
    "        temp_out = model(torch.permute(data[0],(0,3,1,2)).to(device))\n",
    "        test_labels = torch.argmax(temp_out,dim=1)\n",
    "        test_labels = test_labels.squeeze(0).detach().cpu().numpy()\n",
    "        im = Image.fromarray(test_labels.astype(np.uint8))\n",
    "        print(f\"./testing_data_final_filtered/testing_data/v2.2/{data[1][0]}_label_kinect.png is saved\")\n",
    "        im.save(f\"testing_data_final_filtered/testing_data/v2.2/{data[1][0]}_label_kinect.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(test_dataloader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Lift depth to point cloud"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You can use other visualization from previous homeworks, like Open3D\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "\n",
    "\n",
    "def show_points(points):\n",
    "    fig = plt.figure()\n",
    "    ax = fig.add_subplot(projection='3d')\n",
    "    ax.set_xlim3d([-2, 2])\n",
    "    ax.set_ylim3d([-2, 2])\n",
    "    ax.set_zlim3d([0, 4])\n",
    "    ax.scatter(points[:, 0], points[:, 2], points[:, 1])\n",
    "    \n",
    "\n",
    "def compare_points(points1, points2):\n",
    "    fig = plt.figure()\n",
    "    ax = fig.add_subplot(projection='3d')\n",
    "    ax.set_xlim3d([-0.5, 0.5])\n",
    "    ax.set_ylim3d([-0.5, 0.5])\n",
    "    ax.set_zlim3d([0, 1])\n",
    "    ax.scatter(points1[:, 0], points1[:, 2], points1[:, 1])\n",
    "    ax.scatter(points2[:, 0], points2[:, 2], points2[:, 1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import copy\n",
    "import open3d\n",
    "from numpy.linalg import inv\n",
    "from tqdm.notebook import trange, tqdm\n",
    "\n",
    "def draw_registration_result(source, target, transformation):\n",
    "    source_temp = copy.deepcopy(source)\n",
    "    target_temp = copy.deepcopy(target)\n",
    "    source_temp.paint_uniform_color([1, 0.706, 0])\n",
    "    target_temp.paint_uniform_color([0, 0.651, 0.929])\n",
    "    source_temp.transform(transformation)\n",
    "    open3d.visualization.draw_geometries([source_temp, target_temp],\n",
    "                                      zoom=3,\n",
    "                                      front=[0.9288, -0.2951, -0.2242],\n",
    "                                      lookat=[1.6784, 2.0612, 1.4451],\n",
    "                                      up=[-0.3402, -0.9189, -0.1996])\n",
    "\n",
    "src = open3d.geometry.PointCloud()\n",
    "tgt = open3d.geometry.PointCloud()\n",
    "\n",
    "# od = 4\n",
    "test_dump = {}\n",
    "for i in trange(len(depth_files)):\n",
    "    trans = []\n",
    "    poses_test = [None]*79\n",
    "    meta_test = load_pickle(meta_files[i])\n",
    "    train_dump = load_pickle('train_dump')\n",
    "    intrinsic = meta_test['intrinsic']\n",
    "    depth = np.array(Image.open(depth_files[i])) / 1000\n",
    "    label = np.array(Image.open(label_files[i]))\n",
    "    z = depth\n",
    "    v, u = np.indices(z.shape)\n",
    "    uv1 = np.stack([u + 0.5, v + 0.5, np.ones_like(z)], axis=-1)\n",
    "    points_viewer = uv1 @ np.linalg.inv(intrinsic).T * z[..., None]  # [H, W, 3]\n",
    "    crops_pcd = np.array([points_viewer[label==idx] for idx in meta_test['object_ids']])\n",
    "    c2w_test = inv(meta_test['extrinsic'])\n",
    "    for od in range(len(meta_test['object_ids'])):\n",
    "        rmse = []\n",
    "        t_list = []\n",
    "        for k in range(len(train_dump[meta_test['object_names'][od]]['pcd'])):\n",
    "            src_pcd = train_dump[meta_test['object_names'][od]]['pcd'][k]\n",
    "            src_pcd = src_pcd.squeeze(0)\n",
    "            if(len(src_pcd)<600):\n",
    "                continue\n",
    "            src_pose = train_dump[meta_test['object_names'][od]]['gt_world'][k]\n",
    "            inv_gt = inv(src_pose)\n",
    "            src_pcd = src_pcd@inv_gt[:3,:3].T + inv_gt[:3,3]\n",
    "            src_pcd = src_pcd.reshape([-1,3])\n",
    "\n",
    "            tgt_pcd = crops_pcd[od]\n",
    "            tgt_pcd = tgt_pcd@c2w_test[:3,:3].T + c2w_test[:3,3]\n",
    "            \n",
    "            src_pcd = src_pcd[:600,:]\n",
    "            tgt_pcd = tgt_pcd[:600,:]\n",
    "\n",
    "            src.points = open3d.utility.Vector3dVector(src_pcd.reshape([-1, 3]))\n",
    "\n",
    "            tgt.points = open3d.utility.Vector3dVector(tgt_pcd.reshape([-1, 3]))\n",
    "\n",
    "\n",
    "            threshold = 0.02\n",
    "\n",
    "\n",
    "            #dump_5.json with inv_gt\n",
    "            trans_init = inv(inv_gt)\n",
    "            q_bar = np.mean(tgt_pcd,axis=0)\n",
    "            p_bar = np.mean(src_pcd,axis=0)\n",
    "            trans_init[:3,3] = q_bar - np.matmul(trans_init[:3,:3],p_bar)\n",
    "\n",
    "\n",
    "            reg_p2p = open3d.pipelines.registration.registration_icp(\n",
    "                src, tgt, threshold, trans_init,\n",
    "                open3d.pipelines.registration.TransformationEstimationPointToPoint(),\n",
    "                open3d.pipelines.registration.ICPConvergenceCriteria(max_iteration=200)\n",
    "            )\n",
    "            if(len(reg_p2p.correspondence_set)==0 or reg_p2p.inlier_rmse==0):\n",
    "                continue\n",
    "            rmse.append(reg_p2p.inlier_rmse)\n",
    "            t_list.append(reg_p2p.transformation)\n",
    "\n",
    "        # the below T is in transform from camera_src to camera_tgt\n",
    "        if(len(rmse)==0):\n",
    "            T = np.eye(4)\n",
    "        else:\n",
    "            T = t_list[np.argmin(rmse)]\n",
    "#         T = reg_p2p.transformation\n",
    "        trans.append(T)\n",
    "        poses_test[meta_test['object_ids'][od]] = T.tolist()\n",
    "\n",
    "    #write to json\n",
    "    test_dump[scenes[i]] = {\"poses_world\":poses_test}\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "print(type(test_dump))\n",
    "file = open('test_dump_last_final.json', 'w')\n",
    "\n",
    "# dump information to that file\n",
    "json.dump(test_dump, file)\n",
    "\n",
    "# close the file\n",
    "file.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Draw bounding boxes of poses on 2D image\n",
    "If you are curious, take a look at `utils.py`. It is very simple."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s =160\n",
    "scene = scenes[s]\n",
    "boxed_image = np.array(Image.open(rgb_files[s])) / 255 \n",
    "meta = load_pickle(meta_files[s])\n",
    "poses_world = np.array([test_dump[scene]['poses_world'][idx] for idx in meta['object_ids']])\n",
    "box_sizes = np.array([meta['extents'][idx] * meta['scales'][idx] for idx in meta['object_ids']])\n",
    "for i in range(len(poses_world)):\n",
    "    utils.draw_projected_box3d(\n",
    "        boxed_image, poses_world[i][:3,3], box_sizes[i], poses_world[i][:3, :3], meta['extrinsic'], meta['intrinsic'],\n",
    "        thickness=2)\n",
    "\n",
    "Image.fromarray((boxed_image * 255).astype(np.uint8))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test data\n",
    "Test data has everything but the poses. Testing data and training data are from the same distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "load_pickle(\"./testing_data_pose_filtered/testing_data/v2.2/1-1-1_meta.pkl\").keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Image.open(\"./testing_data/v2.2/1-1-1_color_kinect.png\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}