Gather_nd_tests.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gather_object_pointcloud(point_cloud, mask, m_points, device='cuda'):\n",
    "    def mask_to_indicies(mask, batch_size):\n",
    "        indices = torch.zeros((batch_size, m_points, 2), dtype=torch.int32)\n",
    "        for i in range(batch_size):\n",
    "            pos_indices = np.where(mask[i, :] > 0.5)[0]\n",
    "            # skip cases when pos_indices is empty\n",
    "            if len(pos_indices) > 0:\n",
    "                if len(pos_indices) > m_points:\n",
    "                    choice = np.random.choice(len(pos_indices),\n",
    "                                              m_points, replace=False)\n",
    "                else:\n",
    "                    choice = np.random.choice(len(pos_indices),\n",
    "                                              m_points - len(pos_indices), replace=True)\n",
    "                    choice = np.concatenate((np.arange(len(pos_indices)), choice))\n",
    "                np.random.shuffle(choice)\n",
    "                indices[i, :, 1] = torch.from_numpy(pos_indices[choice])\n",
    "            indices[i, :, 0] = i\n",
    "        return indices.to(device)\n",
    "\n",
    "    batch_size = mask.size()[0]\n",
    "    n_channels = point_cloud.size()[2]\n",
    "\n",
    "    #indices = torch.autograd.Variable(mask_to_indicies(mask, batch_size))\n",
    "    indices = mask_to_indicies(mask, batch_size)\n",
    "    \n",
    "    object_point_cloud = torch.zeros((mask.size()[0], m_points, n_channels))\n",
    "    #object_point_cloud = torch.autograd.Variable(object_point_cloud)\n",
    "\n",
    "    for i in range(batch_size):\n",
    "        count = 0\n",
    "        for j in indices[i, :, 1]:\n",
    "            object_point_cloud[i, count, :] = point_cloud[i, j, :]\n",
    "            count += 1\n",
    "\n",
    "    return object_point_cloud.to(device), indices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[ 0.13329216, -0.42825173, -0.66352704, -0.25151226],\n",
       "        [-0.28417719,  0.38093518,  0.48201524, -0.19048918],\n",
       "        [-2.02072253,  1.30598666, -0.14669808,  1.29994792]],\n",
       "\n",
       "       [[-0.64609169,  0.97508493,  0.33305814,  1.08070891],\n",
       "        [ 2.40611802, -0.47110607, -0.45400631,  0.08251184],\n",
       "        [-0.28001176, -2.27080497, -0.02160056, -0.07199822]]])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "x = np.random.randn(2,3,4)\n",
    "x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "Expected tensor [2 x 3 x 2], src [2 x 6 x 4] and index [2 x 3 x 2] to have the same size apart from dimension 0 at /opt/conda/conda-bld/pytorch_1535491974311/work/aten/src/TH/generic/THTensorMath.cpp:620",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-20-bd9fc2148696>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[0mindics\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mLongTensor\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0midx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0mtorch_pc\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mFloatTensor\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpc\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 4\u001b[0;31m \u001b[0moutput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgather\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtorch_pc\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mindics\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      5\u001b[0m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutput\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msize\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      6\u001b[0m \u001b[0moutput\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: Expected tensor [2 x 3 x 2], src [2 x 6 x 4] and index [2 x 3 x 2] to have the same size apart from dimension 0 at /opt/conda/conda-bld/pytorch_1535491974311/work/aten/src/TH/generic/THTensorMath.cpp:620"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "indics = torch.LongTensor(idx)\n",
    "torch_pc = torch.FloatTensor(pc)\n",
    "output = torch.gather(torch_pc, 0, indics)\n",
    "print(output.size())\n",
    "output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "idx = np.ones((2,3,2))\n",
    "x = np.random.randn(2,3,4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[-1.56831218, -0.55415761,  0.51438201,  1.35513448],\n",
       "        [ 1.56367027, -0.73864353, -0.71441799,  0.78760628],\n",
       "        [ 0.08812818, -0.36317287,  0.60575778,  1.3710668 ]],\n",
       "\n",
       "       [[ 0.33651575,  0.66759067,  0.98694863, -1.39022862],\n",
       "        [-0.44095884, -1.5241819 , -0.67308512, -0.15392798],\n",
       "        [ 0.72538415, -0.739054  , -0.34964515,  1.31625855]]])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor 'GatherNd_1:0' shape=(2, 3, 4) dtype=float64>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "indics_tf = tf.constant(idx, dtype=np.int32)\n",
    "\n",
    "pc_tf = tf.convert_to_tensor(pc)\n",
    "\n",
    "output_tf = tf.gather_nd(pc_tf, indics_tf)\n",
    "output_tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(2, 3, 4)\n"
     ]
    }
   ],
   "source": [
    "sess = tf.InteractiveSession()\n",
    "output_tf_eval = output_tf.eval()\n",
    "print(output_tf_eval.shape)\n",
    "sess.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(2, 3, 4)\n",
      "[[[-0.2224 -0.5145  2.2887  0.    ]\n",
      "  [-0.1772  0.1358 -3.0995  0.    ]\n",
      "  [ 0.7007  0.5813  2.545   0.    ]]\n",
      "\n",
      " [[ 0.      0.      0.      0.    ]\n",
      "  [ 0.      0.      0.      0.    ]\n",
      "  [ 0.      0.      0.      0.    ]]]\n"
     ]
    }
   ],
   "source": [
    "o = gather_x(pc, idx)\n",
    "print(o.shape)\n",
    "print(o)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "point_cloud = x\n",
    "indices = np.ones((2,3,2), dtype=np.int32)\n",
    "batch_size = 2\n",
    "\n",
    "\n",
    "def gather_x(point_cloud, indices):\n",
    "    object_point_cloud = np.zeros((batch_size, 3, 4))\n",
    "    for i in range(batch_size):\n",
    "        count = 0\n",
    "        for j in indices[i, :, 1]:\n",
    "            object_point_cloud[i, count, :] = point_cloud[i, j, :]\n",
    "            count += 1\n",
    "    return object_point_cloud"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "idx = np.asarray([[[0, 1],[0, 2],[0, 4]], [[1, 2],[1, 2], [1, 2]]])\n",
    "idx_torch = np.asarray([[[1, 1,1,1],[2,2,2, 2],[4,4,4, 4]], [[2,2,2, 2],[2,2,2, 2], [2,2,2,2]]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "pc = np.asarray([[[ 0.3298, -0.3560,  1.5519,  0.0000],\n",
    "         [-0.2224, -0.5145,  2.2887,  0.0000],\n",
    "         [-0.1772,  0.1358, -3.0995,  0.0000],\n",
    "         [-0.9271, -0.4670, -3.5719,  0.0000],\n",
    "         [ 0.7007,  0.5813,  2.5450,  0.0000],\n",
    "         [ 0.2962,  0.6205,  0.2858,  0.1900]],\n",
    "\n",
    "        [[ 0.4187, -0.2572,  1.7973,  0.0000],\n",
    "         [ 0.1225,  1.0131,  0.6915,  0.2000],\n",
    "         [ 0.0000,  0.0000,  0.0000,  0.0000],\n",
    "         [-0.4852,  0.9913,  0.6028,  0.2900],\n",
    "         [-1.5547, -0.0345, 18.2407,  0.3100],\n",
    "         [ 0.4657,  0.6077, -0.4934,  0.0000]]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(2, 3, 4)\n",
      "[[[-0.2224 -0.5145  2.2887  0.    ]\n",
      "  [-0.1772  0.1358 -3.0995  0.    ]\n",
      "  [ 0.7007  0.5813  2.545   0.    ]]\n",
      "\n",
      " [[ 0.      0.      0.      0.    ]\n",
      "  [ 0.      0.      0.      0.    ]\n",
      "  [ 0.      0.      0.      0.    ]]]\n"
     ]
    }
   ],
   "source": [
    "o = gather_x(pc, idx)\n",
    "print(o.shape)\n",
    "print(o)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(2, 3, 4)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[[ 0.3298, -0.3560,  1.5519,  0.0000],\n",
       "         [-0.2224, -0.5145,  2.2887,  0.0000],\n",
       "         [-0.1772,  0.1358, -3.0995,  0.0000]],\n",
       "\n",
       "        [[-0.4852,  0.9913,  0.6028,  0.2900],\n",
       "         [-1.5547, -0.0345, 18.2407,  0.3100],\n",
       "         [ 0.4657,  0.6077, -0.4934,  0.0000]]])"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "idx_test = np.asarray([[[0,0,0,0],\n",
    "  [1,1,1,1],\n",
    "  [2,2,2,2]],\n",
    " [[3,3,3,3],\n",
    "  [4,4,4,4],\n",
    "  [5,5,5,5]]])\n",
    "print(idx_test.shape)\n",
    "torch.gather(torch_pc, 1, torch.LongTensor(idx_test))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 0.3298, -0.3560,  1.5519,  0.0000],\n",
       "         [-0.2224, -0.5145,  2.2887,  0.0000],\n",
       "         [-0.1772,  0.1358, -3.0995,  0.0000],\n",
       "         [-0.9271, -0.4670, -3.5719,  0.0000],\n",
       "         [ 0.7007,  0.5813,  2.5450,  0.0000],\n",
       "         [ 0.2962,  0.6205,  0.2858,  0.1900]],\n",
       "\n",
       "        [[ 0.4187, -0.2572,  1.7973,  0.0000],\n",
       "         [ 0.1225,  1.0131,  0.6915,  0.2000],\n",
       "         [ 0.0000,  0.0000,  0.0000,  0.0000],\n",
       "         [-0.4852,  0.9913,  0.6028,  0.2900],\n",
       "         [-1.5547, -0.0345, 18.2407,  0.3100],\n",
       "         [ 0.4657,  0.6077, -0.4934,  0.0000]]])"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch_pc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[-0.2224, -0.5145,  2.2887,  0.    ],\n",
       "        [-0.1772,  0.1358, -3.0995,  0.    ],\n",
       "        [ 0.7007,  0.5813,  2.545 ,  0.    ]],\n",
       "\n",
       "       [[ 0.    ,  0.    ,  0.    ,  0.    ],\n",
       "        [ 0.    ,  0.    ,  0.    ,  0.    ],\n",
       "        [ 0.    ,  0.    ,  0.    ,  0.    ]]])"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output_tf_eval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch_pc = torch.FloatTensor(pc)\n",
    "torch_idx_tensor = torch.LongTensor(idx_torch)\n",
    "\n",
    "output_torch = torch.gather(torch_pc, 1, torch_idx_tensor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[-0.2224, -0.5145,  2.2887,  0.0000],\n",
       "         [-0.1772,  0.1358, -3.0995,  0.0000],\n",
       "         [ 0.7007,  0.5813,  2.5450,  0.0000]],\n",
       "\n",
       "        [[ 0.0000,  0.0000,  0.0000,  0.0000],\n",
       "         [ 0.0000,  0.0000,  0.0000,  0.0000],\n",
       "         [ 0.0000,  0.0000,  0.0000,  0.0000]]])"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output_torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[-0.2224, -0.5145,  2.2887,  0.    ],\n",
       "        [-0.1772,  0.1358, -3.0995,  0.    ],\n",
       "        [ 0.7007,  0.5813,  2.545 ,  0.    ]],\n",
       "\n",
       "       [[ 0.    ,  0.    ,  0.    ,  0.    ],\n",
       "        [ 0.    ,  0.    ,  0.    ,  0.    ],\n",
       "        [ 0.    ,  0.    ,  0.    ,  0.    ]]])"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output_tf_eval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-1.3644, -1.3644, -1.3644, -1.3644],\n",
       "        [-0.1275, -0.1275, -0.1275, -0.1275],\n",
       "        [-2.3967, -2.3967, -2.3967, -2.3967],\n",
       "        [-0.4475, -0.4475, -0.4475, -0.4475],\n",
       "        [-0.6955, -0.6955, -0.6955, -0.6955],\n",
       "        [ 0.4388,  0.4388,  0.4388,  0.4388]], dtype=torch.float64)"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.from_numpy(np.random.randn((6))).unsqueeze(1).repeat(1,4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([6, 4])"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}