text: import torch.utils.data as data
  from PIL import Image
  import os
  import os.path
  import errno
  import torch
  import json
  import codecs
  import numpy as np
  import sys
  import torchvision.transforms as transforms
  import argparse
  import json
  import time
  import random
  import numpy.ma as ma
  import copy
  import scipy.misc
  import scipy.io as scio
  import yaml
  import cv2
  
  
  class PoseDataset(data.Dataset):
      def __init__(self, mode, num, add_noise, root, noise_trans, refine):
          self.objlist = [1, 2, 4, 5, 6, 8, 9, 10, 11, 12, 13, 14, 15]
          self.mode = mode
  
          self.list_rgb = []
          self.list_depth = []
          self.list_label = []
          self.list_obj = []
          self.list_rank = []
          self.meta = {}
          self.pt = {}
          self.root = root
          self.noise_trans = noise_trans
          self.refine = refine
  
          item_count = 0
          for item in self.objlist:
              if self.mode == 'train':
                  input_file = open('{0}/data/{1}/train.txt'.format(self.root, '%02d' % item))
              else:
                  input_file = open('{0}/data/{1}/test.txt'.format(self.root, '%02d' % item))
              while 1:
                  item_count += 1
                  input_line = input_file.readline()
                  if self.mode == 'test' and item_count % 10 != 0:
                      continue
                  if not input_line:
                      break
                  if input_line[-1:] == '\n':
                      input_line = input_line[:-1]
                  self.list_rgb.append('{0}/data/{1}/rgb/{2}.png'.format(self.root, '%02d' % item, input_line))
                  self.list_depth.append('{0}/data/{1}/depth/{2}.png'.format(self.root, '%02d' % item, input_line))
                  if self.mode == 'eval':
                      self.list_label.append('{0}/segnet_results/{1}_label/{2}_label.png'.format(self.root, '%02d' % item, input_line))
                  else:
                      self.list_label.append('{0}/data/{1}/mask/{2}.png'.format(self.root, '%02d' % item, input_line))
                  
                  self.list_obj.append(item)
                  self.list_rank.append(int(input_line))
  
              meta_file = open('{0}/data/{1}/gt.yml'.format(self.root, '%02d' % item), 'r')
              self.meta[item] = yaml.load(meta_file)
              self.pt[item] = ply_vtx('{0}/models/obj_{1}.ply'.format(self.root, '%02d' % item))
              
              print("Object {0} buffer loaded".format(item))
  
          self.length = len(self.list_rgb)
  
          self.cam_cx = 325.26110
          self.cam_cy = 242.04899
          self.cam_fx = 572.41140
          self.cam_fy = 573.57043
  
          self.xmap = np.array([[j for i in range(640)] for j in range(480)])
          self.ymap = np.array([[i for i in range(640)] for j in range(480)])
          
          self.num = num
          self.add_noise = add_noise
          self.trancolor = transforms.ColorJitter(0.2, 0.2, 0.2, 0.05)
          self.norm = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
          self.border_list = [-1, 40, 80, 120, 160, 200, 240, 280, 320, 360, 400, 440, 480, 520, 560, 600, 640, 680]
          self.num_pt_mesh_large = 500
          self.num_pt_mesh_small = 500
          self.symmetry_obj_idx = [7, 8]
          self.transforms=transforms.Compose(transforms.Resize([3,160,120]),transforms.ToTensor)
      def __getitem__(self, index):
          img = Image.open(self.list_rgb[index])
          ori_img = np.array(img)
          img=self.transforms(img)
          depth = np.array(Image.open(self.list_depth[index]))
          label = np.array(Image.open(self.list_label[index]))
          obj = self.list_obj[index]
          rank = self.list_rank[index]        
  
          if obj == 2:
              for i in range(0, len(self.meta[obj][rank])):
                  if self.meta[obj][rank][i]['obj_id'] == 2:
                      meta = self.meta[obj][rank][i]
                      break
          else:
              meta = self.meta[obj][rank][0]
  
          mask_depth = ma.getmaskarray(ma.masked_not_equal(depth, 0))
          if self.mode == 'eval':
              mask_label = ma.getmaskarray(ma.masked_equal(label, np.array(255)))
          else:
              mask_label = ma.getmaskarray(ma.masked_equal(label, np.array([255, 255, 255])))[:, :, 0]
          
          mask = mask_label * mask_depth
  
          if self.add_noise:
              img = self.trancolor(img)
  
          img = np.array(img)[:, :, :3]
          img = np.transpose(img, (2, 0, 1))
          img_masked = img
  
          if self.mode == 'eval':
              rmin, rmax, cmin, cmax = get_bbox(mask_to_bbox(mask_label))
          else:
              rmin, rmax, cmin, cmax = get_bbox(meta['obj_bb'])
  
          img_masked = img_masked[:, rmin:rmax, cmin:cmax]
          #p_img = np.transpose(img_masked, (1, 2, 0))
          #scipy.misc.imsave('evaluation_result/{0}_input.png'.format(index), p_img)
  
          target_r = np.resize(np.array(meta['cam_R_m2c']), (3, 3))
          target_t = np.array(meta['cam_t_m2c'])
          add_t = np.array([random.uniform(-self.noise_trans, self.noise_trans) for i in range(3)])
  
          choose = mask[rmin:rmax, cmin:cmax].flatten().nonzero()[0]
          if len(choose) == 0:
              cc = torch.LongTensor([0])
              return(cc, cc, cc, cc, cc, cc)
  
          if len(choose) > self.num:
              c_mask = np.zeros(len(choose), dtype=int)
              c_mask[:self.num] = 1
              np.random.shuffle(c_mask)
              choose = choose[c_mask.nonzero()]
          else:
              choose = np.pad(choose, (0, self.num - len(choose)), 'wrap')
          #add a dim to make pt012 concatenate different dims
          depth_masked = depth[rmin:rmax, cmin:cmax].flatten()[choose][:, np.newaxis].astype(np.float32)
          xmap_masked = self.xmap[rmin:rmax, cmin:cmax].flatten()[choose][:, np.newaxis].astype(np.float32)
          ymap_masked = self.ymap[rmin:rmax, cmin:cmax].flatten()[choose][:, np.newaxis].astype(np.float32)
          choose = np.array([choose])
  
  #pixel to point cloud
          cam_scale = 1.0
          pt2 = depth_masked / cam_scale
          pt0 = (ymap_masked - self.cam_cx) * pt2 / self.cam_fx
          pt1 = (xmap_masked - self.cam_cy) * pt2 / self.cam_fy
          cloud = np.concatenate((pt0, pt1, pt2), axis=1)
          cloud = cloud / 1000.0
  
          if self.add_noise:
              cloud = np.add(cloud, add_t)
  
          #fw = open('evaluation_result/{0}_cld.xyz'.format(index), 'w')
          #for it in cloud:
          #    fw.write('{0} {1} {2}\n'.format(it[0], it[1], it[2]))
          #fw.close()
  
          model_points = self.pt[obj] / 1000.0
          dellist = [j for j in range(0, len(model_points))]
          #select points from dellist,the two progress only slect
          dellist = random.sample(dellist, len(model_points) - self.num_pt_mesh_small)
          model_points = np.delete(model_points, dellist, axis=0)
  
          #fw = open('evaluation_result/{0}_model_points.xyz'.format(index), 'w')
          #for it in model_points:
          #    fw.write('{0} {1} {2}\n'.format(it[0], it[1], it[2]))
          #fw.close()
  
          target = np.dot(model_points, target_r.T)
          if self.add_noise:
              target = np.add(target, target_t / 1000.0 + add_t)
              out_t = target_t / 1000.0 + add_t
          else:
              target = np.add(target, target_t / 1000.0)
              out_t = target_t / 1000.0
  
          #fw = open('evaluation_result/{0}_tar.xyz'.format(index), 'w')
          #for it in target:
          #    fw.write('{0} {1} {2}\n'.format(it[0], it[1], it[2]))
          #fw.close()
  
          return torch.from_numpy(cloud.astype(np.float32)), \
                 torch.LongTensor(choose.astype(np.int32)), \
                 self.norm(torch.from_numpy(img_masked.astype(np.float32))), \
                 torch.from_numpy(target.astype(np.float32)), \
                 torch.from_numpy(model_points.astype(np.float32)), \
                 torch.LongTensor([self.objlist.index(obj)])
  
      def __len__(self):
          return self.length
  
      def get_sym_list(self):
          return self.symmetry_obj_idx
  
      def get_num_points_mesh(self):
          if self.refine:
              return self.num_pt_mesh_large
          else:
              return self.num_pt_mesh_small
  
  
  
  border_list = [-1, 40, 80, 120, 160, 200, 240, 280, 320, 360, 400, 440, 480, 520, 560, 600, 640, 680]
  img_width = 480
  img_length = 640
  
  
  def mask_to_bbox(mask):
      mask = mask.astype(np.uint8)
      _, contours, _ = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
      x = 0
      y = 0
      w = 0
      h = 0
      for contour in contours:
          tmp_x, tmp_y, tmp_w, tmp_h = cv2.boundingRect(contour)
          if tmp_w * tmp_h > w * h:
              x = tmp_x
              y = tmp_y
              w = tmp_w
              h = tmp_h
      return [x, y, w, h]
  
  
  def get_bbox(bbox):
      bbx = [bbox[1], bbox[1] + bbox[3], bbox[0], bbox[0] + bbox[2]]
      if bbx[0] < 0:
          bbx[0] = 0
      if bbx[1] >= 480:
          bbx[1] = 479
      if bbx[2] < 0:
          bbx[2] = 0
      if bbx[3] >= 640:
          bbx[3] = 639                
      rmin, rmax, cmin, cmax = bbx[0], bbx[1], bbx[2], bbx[3]
      r_b = rmax - rmin
      for tt in range(len(border_list)):
          if r_b > border_list[tt] and r_b < border_list[tt + 1]:
              r_b = border_list[tt + 1]
              break
      c_b = cmax - cmin
      for tt in range(len(border_list)):
          if c_b > border_list[tt] and c_b < border_list[tt + 1]:
              c_b = border_list[tt + 1]
              break
      center = [int((rmin + rmax) / 2), int((cmin + cmax) / 2)]
      rmin = center[0] - int(r_b / 2)
      rmax = center[0] + int(r_b / 2)
      cmin = center[1] - int(c_b / 2)
      cmax = center[1] + int(c_b / 2)
      if rmin < 0:
          delt = -rmin
          rmin = 0
          rmax += delt
      if cmin < 0:
          delt = -cmin
          cmin = 0
          cmax += delt
      if rmax > 480:
          delt = rmax - 480
          rmax = 480
          rmin -= delt
      if cmax > 640:
          delt = cmax - 640
          cmax = 640
          cmin -= delt
      return rmin, rmax, cmin, cmax
  
  
  def ply_vtx(path):
      f = open(path)
      assert f.readline().strip() == "ply"
      f.readline()
      f.readline()
      N = int(f.readline().split()[-1])
      while f.readline().strip() != "end_header":
          continue
      pts = []
      for _ in range(N):
          pts.append(np.float32(f.readline().split()[:3]))
      return np.array(pts)
  