scene_util.py

import os
import sys
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.append(BASE_DIR)

import numpy as np
from sklearn.neighbors import NearestNeighbors
from numpy import linalg as la
import scipy.io as sio

def cart2sph(xyz):
  xy = xyz[:,0]**2+xyz[:,1]**2
  aer = np.zeros(xyz.shape)
  aer[:,2] = np.sqrt(xy+xyz[:,2]**2)
  aer[:,1] = np.arctan2(xyz[:,2],np.sqrt(xy))
  aer[:,0] = np.arctan2(xyz[:,1],xyz[:,0])
  return aer

# generate virtual scan of a scene by subsampling the point cloud
def virtual_scan(xyz, mode=-1):
  camloc = np.mean(xyz,axis=0)
  camloc[2] = 1.5 # human height
  if mode==-1:
    view_dr = np.array([2*np.pi*np.random.random(), np.pi/10*(np.random.random()-0.75)])
    camloc[:2] -= (0.8+0.7*np.random.random())*np.array([np.cos(view_dr[0]),np.sin(view_dr[0])])
  else:
    view_dr = np.array([np.pi/4*mode, 0])
    camloc[:2] -= np.array([np.cos(view_dr[0]),np.sin(view_dr[0])])
  ct_ray_dr = np.array([np.cos(view_dr[1])*np.cos(view_dr[0]), np.cos(view_dr[1])*np.sin(view_dr[0]), np.sin(view_dr[1])])
  hr_dr = np.cross(ct_ray_dr, np.array([0,0,1]))
  hr_dr /= la.norm(hr_dr)
  vt_dr = np.cross(hr_dr, ct_ray_dr)
  vt_dr /= la.norm(vt_dr)
  xx = np.linspace(-0.6,0.6,200) #200
  yy = np.linspace(-0.45,0.45,150) #150
  xx, yy = np.meshgrid(xx,yy)
  xx = xx.reshape(-1,1)
  yy = yy.reshape(-1,1)
  rays = xx*hr_dr.reshape(1,-1)+yy*vt_dr.reshape(1,-1)+ct_ray_dr.reshape(1,-1)
  rays_aer = cart2sph(rays)
  local_xyz = xyz-camloc.reshape(1,-1)
  local_aer = cart2sph(local_xyz)
  nbrs = NearestNeighbors(n_neighbors=1, algorithm='kd_tree').fit(rays_aer[:,:2])
  mindd, minidx = nbrs.kneighbors(local_aer[:,:2])
  mindd = mindd.reshape(-1)
  minidx = minidx.reshape(-1)

  sub_idx = mindd<0.01
  if sum(sub_idx)<100:
    return np.ones(0)
  sub_r = local_aer[sub_idx,2]
  sub_minidx = minidx[sub_idx]
  min_r = float('inf')*np.ones(np.max(sub_minidx)+1)
  for i in xrange(len(sub_r)):
    if sub_r[i]<min_r[sub_minidx[i]]:
      min_r[sub_minidx[i]] = sub_r[i]
  sub_smpidx = np.ones(len(sub_r))
  for i in xrange(len(sub_r)):
    if sub_r[i]>min_r[sub_minidx[i]]:
      sub_smpidx[i] = 0
  smpidx = np.where(sub_idx)[0]
  smpidx = smpidx[sub_smpidx==1]
  return smpidx

if __name__=='__main__':
  pc = np.load('scannet_dataset/scannet_scenes/scene0015_00.npy')
  print pc.shape
  xyz = pc[:,:3]
  seg = pc[:,7]
  smpidx = virtual_scan(xyz,mode=2)
  xyz = xyz[smpidx,:]
  seg = seg[smpidx]
  sio.savemat('tmp.mat',{'pc':xyz,'seg':seg})