text: import json
  import logging
  import os.path as osp
  from Queue import Empty, Queue
  from threading import Thread, current_thread
  
  import numpy as np
  
  from config import SHAPENET_IM
  from loader import read_camera, read_depth, read_im, read_quat, read_vol
  from custom_db import *
  
  
  def get_split(split_js='data/splits.json'):
      dir_path = osp.dirname(osp.realpath(__file__))
      with open(osp.join(dir_path, split_js), 'r') as f:
          js = json.load(f)
  
      return js
  
  
  class ShapeNet(object):
      def __init__(self,
                   im_dir=SHAPENET_IM,
                   split_file='data/splits.json',
                   vox_dir=None,
                   shape_ids=None,
                   num_renders=20,
                   rng_seed=0,
                   custom_db=False):
          self.vox_dir = vox_dir
          self.im_dir = im_dir
          self.split_file = split_file
          self.splits_all = get_split(split_file)
          self.shape_ids = (self.splits_all.keys()
                            if shape_ids is None else shape_ids)
          self.splits = {k: self.splits_all[k] for k in self.shape_ids}
  
          self.shape_cls = [
              self.splits[x]['name'].split(',')[0] for x in self.shape_ids
          ]
          self.rng = rng_seed
          self.num_renders = num_renders
          if custom_db:
              self.load_func = {
                  'im': self.get_im_custom,
                  'depth': self.get_depth_custom,
                  'K': self.get_K_custom,
                  'R': self.get_R_custom,
                  'quat': self.get_quat,
                  'vol': self.get_vol,
                  'shape_id': self.get_sid,
                  'model_id': self.get_mid,
                  'view_idx': self.get_view_idx
              }
          else:
              self.load_func = {
                  'im': self.get_im,
                  'depth': self.get_depth,
                  'K': self.get_K,
                  'R': self.get_R,
                  'quat': self.get_quat,
                  'vol': self.get_vol,
                  'shape_id': self.get_sid,
                  'model_id': self.get_mid,
                  'view_idx': self.get_view_idx
              }
  
          self.all_items = self.load_func.keys()
  
          self.logger = logging.getLogger('mview3d.' + __name__)
          np.random.seed(self.rng)
  
      # custom training functions
      def read_camera_custom(self, f):
          cam = json.load(open(f))
          R = np.array(cam['R'])
          t = np.array(cam['t']).reshape(3,1)
          Rt = np.hstack((R,t))
          K = np.array(cam['K'])
          return K, Rt
  
      def get_K_custom(self, sid, mid, idx):
          rand_idx = idx
          cams = []
          for ix in rand_idx:
              f = osp.join(self.im_dir, sid, mid, '{:d}.json'.format(ix))
              cams.append(self.read_camera_custom(f))
          camK = np.stack([c[0] for c in cams], axis=0)
          return camK
  
      def get_R_custom(self, sid, mid, idx):
          rand_idx = idx
          cams = []
          for ix in rand_idx:
              f = osp.join(self.im_dir, sid, mid, '{:d}.json'.format(ix))
              cams.append(self.read_camera_custom(f))
          camR = np.stack([c[1] for c in cams], axis=0)
          return camR
  
      def get_depth_custom(self, sid, mid, idx):
          rand_idx = idx
          depths = []
          for ix in rand_idx:
              f = osp.join(self.im_dir, sid, mid, '{:d}_depthmap.png'.format(ix))
              depths.append(read_depth(f))
          return np.stack(depths, axis=0)
  
      def get_im_custom(self, sid, mid, idx):
          rand_idx = idx
          ims = []
          for ix in rand_idx:
              f = osp.join(self.im_dir, sid, mid, '{:d}.png'.format(ix))
              ims.append(read_im(f))
          return np.stack(ims, axis=0)
  
      ################################################
  
      def get_mids(self, sid):
          return self.splits[sid]
  
      def get_smids(self, split):
          smids = []
          for k, v in self.splits.iteritems():
              smids.extend([(k, m) for m in v[split]])
          smids = np.random.permutation(smids)
          return smids
  
      def get_sid(self, sid, mid, idx=None):
          return np.array([sid])
  
      def get_view_idx(self, sid, mid, idx):
          return idx
  
      def get_mid(self, sid, mid, idx=None):
          return np.array([mid])
  
      def get_K(self, sid, mid, idx):
          rand_idx = idx
          cams = []
          for ix in rand_idx:
              f = osp.join(self.im_dir, sid, mid, 'camera_{:d}.mat'.format(ix))
              cams.append(read_camera(f))
          camK = np.stack([c[0] for c in cams], axis=0)
          return camK
  
      def get_R(self, sid, mid, idx):
          rand_idx = idx
          cams = []
          for ix in rand_idx:
              f = osp.join(self.im_dir, sid, mid, 'camera_{:d}.mat'.format(ix))
              cams.append(read_camera(f))
          camR = np.stack([c[1] for c in cams], axis=0)
          return camR
  
      def get_quat(self, sid, mid, idx):
          rand_idx = idx
          cams = []
          for ix in rand_idx:
              f = osp.join(self.im_dir, sid, mid, 'camera_{:d}.mat'.format(ix))
              cams.append(read_quat(f))
          camq = np.stack(cams, axis=0)
          return camq
  
      def get_depth(self, sid, mid, idx):
          rand_idx = idx
          depths = []
          for ix in rand_idx:
              f = osp.join(self.im_dir, sid, mid, 'depth_{:d}.png'.format(ix))
              depths.append(read_depth(f))
          return np.stack(depths, axis=0)
  
      def get_im(self, sid, mid, idx):
          rand_idx = idx
          ims = []
          for ix in rand_idx:
              f = osp.join(self.im_dir, sid, mid, 'render_{:d}.png'.format(ix))
              ims.append(read_im(f))
          return np.stack(ims, axis=0)
  
      def get_vol(self, sid, mid, idx=None, tsdf=False):
          if self.vox_dir is None:
              self.logger.error('Voxel dir not defined')
          f = osp.join(self.vox_dir, sid, mid)
          return read_vol(f, tsdf)
  
      def fetch_data(self, smids, items, im_batch):
          with self.coord.stop_on_exception():
              while not self.coord.should_stop():
                  data = {}
                  try:
                      data_idx = self.queue_idx.get(timeout=0.5)
                  except Empty:
                      self.logger.debug('Index queue empty - {:s}'.format(
                          current_thread().name))
                      continue
  
                  view_idx = np.random.choice(
                      self.num_renders, size=(im_batch, ), replace=False)
                  sid, mid = smids[data_idx]
                  for i in items:
                      data[i] = self.load_func[i](sid, mid, view_idx)
  
                  self.queue_data.put(data)
                  if self.loop_data:
                      self.queue_idx.put(data_idx)
  
      def fetch_data_single_thread(self, smids, items, im_batch):
          data = {}
          try:
              data_idx = self.queue_idx.get(timeout=0.5)
          except Empty:
              self.logger.debug('Index queue empty - {:s}'.format(
                  current_thread().name))
              return
  
          view_idx = np.random.choice(
              self.num_renders, size=(im_batch, ), replace=False)
          sid, mid = smids[data_idx]
          for i in items:
              data[i] = self.load_func[i](sid, mid, view_idx)
  
          self.queue_data.put(data)
          if self.loop_data:
              self.queue_idx.put(data_idx)
  
      def init_queue(self,
                     smids,
                     im_batch,
                     items,
                     coord,
                     nepochs=None,
                     qsize=32,
                     nthreads=4):
          self.coord = coord
          self.queue_data = Queue(maxsize=qsize)
          if nepochs is None:
              nepochs = 1
              self.loop_data = True
          else:
              self.loop_data = False
          self.total_items = nepochs * len(smids)
          self.queue_idx = Queue(maxsize=self.total_items)
  
          for nx in range(nepochs):
              for rx in range(len(smids)):
                  self.queue_idx.put(rx)
  
          self.qthreads = []
          self.logger.info('Starting {:d} prefetch threads'.format(nthreads))
          for qx in range(nthreads):
              worker = Thread(
                  target=self.fetch_data, args=(smids, items, im_batch))
              worker.start()
              self.coord.register_thread(worker)
              self.qthreads.append(worker)
  
          # self.fetch_data_single_thread(smids, items, im_batch)
  
  
      def close_queue(self, e=None):
          self.logger.debug('Closing queue')
          self.coord.request_stop(e)
          try:
              while True:
                  self.queue_idx.get(block=False)
          except Empty:
              self.logger.debug('Emptied idx queue')
  
          try:
              while True:
                  self.queue_data.get(block=False)
          except Empty:
              self.logger.debug("Emptied data queue")
  
      def next_batch(self, items, batch_size, timeout=0.5):
          data = []
          cnt = 0
          while cnt < batch_size:
              try:
                  dt = self.queue_data.get(timeout=timeout)
                  self.total_items -= 1
                  data.append(dt)
              except Empty:
                  self.logger.debug('Example queue empty')
                  if self.total_items <= 0 and not self.loop_data:
                      # Exhausted all data
                      self.close_queue()
                      break
                  else:
                      continue
              cnt += 1
  
          if len(data) == 0:
              return
  
          batch_data = {}
          for k in items:
              batch_data[k] = []
              for dt in data:
                  batch_data[k].append(dt[k])
              batched = np.stack(batch_data[k])
              batch_data[k] = batched
  
          return batch_data
  
      def reset(self):
          np.random.seed(self.rng)
  