gen_gmm-mine.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pylab import *\n",
    "import os\n",
    "import sys\n",
    "from mixture import GaussianMixture\n",
    "import pickle\n",
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_dir = 'rgbd_dataset_freiburg3_long_office_household'\n",
    "depth_dir = 'depth'\n",
    "gmm_dir = 'gmm_qqvga_mine_pixel_fixed'\n",
    "\n",
    "# og \n",
    "fx = 525.0  # focal length x\n",
    "fy = 525.0  # focal length y\n",
    "cx = 319.5  # optical center x\n",
    "cy = 239.5  # optical center y\n",
    "\n",
    "# fri3 \n",
    "fx = 535.4  # focal length x\n",
    "fy = 539.2  # focal length y\n",
    "cx = 320.1  # optical center x\n",
    "cy = 247.6  # optical center y\n",
    "factor = 5000 # for the 16-bit PNG files\n",
    "# OR: factor = 1 # for the 32-bit float images in the ROS bag files\n",
    "baseline = 0.075 \n",
    "\n",
    "full_depth_dir = os.path.join(dataset_dir,depth_dir)\n",
    "full_gmm_dir = os.path.join(dataset_dir,gmm_dir)\n",
    "if not os.path.exists(full_gmm_dir):\n",
    "    os.mkdir(full_gmm_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_disp = 325 #fx*baseline/0.125\n",
    "min_disp = 2#f\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_grid = np.repeat(np.arange(480)[:,None],640,1)\n",
    "x_grid = np.repeat(np.arange(640)[None,:],480,0)\n",
    "\n",
    "size_of_disp_error = 1.0 #0.125\n",
    "\n",
    "def img_to_pts(img):\n",
    "    Z = img/factor\n",
    "    X = (x_grid-cx) * Z /fx\n",
    "    Y = (y_grid-cy) * Z /fy\n",
    "    \n",
    "    d =  (fx*baseline)/Z\n",
    "    zf = (fx*baseline)/np.clip(d+size_of_disp_error,min_disp,max_disp)\n",
    "    zb = (fx*baseline)/np.clip(d-size_of_disp_error,min_disp,max_disp)\n",
    "    \n",
    "    subsample = 4\n",
    "    step_frac = 0.5 * subsample\n",
    "    Xp = x_grid-cx+step_frac\n",
    "    Yp = y_grid-cy+step_frac\n",
    "    Xm = x_grid-cx-step_frac\n",
    "    Ym = y_grid-cy-step_frac\n",
    "    frtptT = np.array(((Xp)/fx,(Yp)/fy,np.ones_like(d)))*zf\n",
    "    lftptT = np.array(((Xm)/fx,(Yp)/fy,np.ones_like(d)))*Z\n",
    "    rgtptT = np.array(((Xp)/fx,(Yp)/fy,np.ones_like(d)))*Z\n",
    "    bckptT = np.array(((Xm)/fx,(Yp)/fy,np.ones_like(d)))*zb\n",
    "    \n",
    "    frtptB = np.array(((Xp)/fx,(Ym)/fy,np.ones_like(d)))*zf\n",
    "    lftptB = np.array(((Xm)/fx,(Ym)/fy,np.ones_like(d)))*Z\n",
    "    rgtptB = np.array(((Xp)/fx,(Ym)/fy,np.ones_like(d)))*Z\n",
    "    bckptB = np.array(((Xm)/fx,(Ym)/fy,np.ones_like(d)))*zb\n",
    "    \n",
    "    box = np.stack([frtptT,lftptT,rgtptT,bckptT,frtptB,lftptB,rgtptB,bckptB])\n",
    "    print(box.shape)\n",
    "    \n",
    "    m = box.mean(0)\n",
    "    print(m.shape)\n",
    "\n",
    "    m1 = (m.reshape((3,-1)).T)[:,None]\n",
    "    m2 = m1.reshape((-1,3,1))\n",
    "    print(m1.shape)\n",
    "    print(m2.shape)\n",
    "\n",
    "    \n",
    "    M = m2 @ m1\n",
    "    print(M.shape)\n",
    "    print(box.shape)\n",
    "    \n",
    "    covar = np.zeros( (640*480,3,3) )\n",
    "    for pt in box:\n",
    "        pt = pt.reshape((3,-1)).T\n",
    "        pt1 = pt.reshape((-1,1,3))\n",
    "        pt2 = pt.reshape((-1,3,1))\n",
    "\n",
    "        covar += (pt2 @ pt1 - M)\n",
    "    covar = covar/12\n",
    "    covar = covar.reshape((480,640,3,3))[::subsample,::subsample,:,:]\n",
    "    X = X[::subsample,::subsample]\n",
    "    Y = Y[::subsample,::subsample]\n",
    "    Z = Z[::subsample,::subsample]\n",
    "    print(X.shape)\n",
    "\n",
    "    xl = X[Z > 0]\n",
    "    yl = Y[Z > 0]\n",
    "    zl = Z[Z > 0]\n",
    "    c1 = (covar.reshape((120,160,-1))[Z > 0]).reshape( (-1,3,3) )\n",
    "    n = zl.shape[0]\n",
    "    print(n,c1.shape)\n",
    "    \n",
    "    idx = np.random.randint(n, size=n//16)\n",
    "\n",
    "    return np.vstack([xl,yl,zl]).T, c1\n",
    "\n",
    "def load_data(file):\n",
    "    image = Image.open(file)   \n",
    "    pixel = np.array(image)\n",
    "    return img_to_pts(pixel)\n",
    "\n",
    "for fl in sorted(os.listdir(full_depth_dir)):\n",
    "    name,ext = os.path.splitext(fl)\n",
    "    new_name = name + '.pkl'\n",
    "    \n",
    "    clf = GaussianMixture(100)\n",
    "    data = load_data(os.path.join(full_depth_dir,fl))\n",
    "    print([_.shape for _ in data])\n",
    "    t1 = time.time()\n",
    "    \n",
    "    n = data[0].shape[0]\n",
    "    clf.set_areas(np.ones(n))\n",
    "    clf.set_covars(data[1])\n",
    "    clf.fit(data[0])\n",
    "    clf.set_areas(None)\n",
    "    clf.set_covars(None)\n",
    "    \n",
    "    print(time.time()-t1)\n",
    "    with open(os.path.join(full_gmm_dir,new_name),'wb') as fp:\n",
    "        pickle.dump(clf,fp)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "frtptB.shape"
   ]
  },
  {
   "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}