def retrain(net, epoch):
    print('\nEpoch: %d' % epoch)
    global best_acc
    net.train()
    train_loss = 0
    total = 0
    correct = 0

    for batch_idx, (inputs, targets) in enumerate(train_loader):
        if use_cuda:
            inputs, targets = inputs.cuda(), targets.cuda()
        optimizer.zero_grad()
        inputs, targets = Variable(inputs), Variable(targets)
        outputs = net(inputs)
        loss = criterion(outputs, targets)
        loss.backward()

        if args.fixed:
            net = util.quantize(net, args.pprec)

        optimizer.step()

        train_loss += loss.data.item()
        _, predicted = torch.max(outputs.data, 1)
        total += targets.size(0)
        correct += float(predicted.eq(targets.data).cpu().sum())

        progress_bar(
            batch_idx, len(train_loader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)' %
            (train_loss /
             (batch_idx + 1), 100. * correct / total, correct, total))
        acc = 100. * correct / total
    if args.fixed:
        net = util.quantize(net, args.pprec)
Ejemplo n.º 2
0
def quantization_train(x_train, y_train, network, learning_rate, epochs, bit):

    loss = np.zeros([epochs, 1])
    loss_fn = torch.nn.MSELoss()

    for param in network.parameters():
        param.data = util.quantize(param.data, bit)

    network = util.normalize(network)

    for ix in range(epochs):
        y_hat = network(x_train)  #Forward pass
        loss_var = loss_fn(y_hat, y_train)
        loss[ix] = loss_var.data.cpu()
        network.zero_grad()
        loss_var.backward()

        for param in network.parameters():
            param.data = param.data - util.quantize(
                learning_rate * param.grad.data, bit)

        network = util.normalize(network)

        for param in network.parameters():
            param.data = util.quantize(param.data, bit)

        if ix % 1000 == 0:
            print('Current epochs : ' + str(ix) + ' th epochs')

    return network, loss
Ejemplo n.º 3
0
def draw_keypoints(img, keypoints, draw_prob):
    """Draws for each keypoint a circle (roughly matching the sigma of the scale)
    with a line for the orientation.
    Args:
        img    The image to which to add the keypoints (gets copied)
        keypoints The keypoints to draw
        draw_prob Probability of drawing a keypoint (the lower the less keypoints are drawn)
    Returns:
        Image with keypoints"""
    height, width = img.shape
    img = np.copy(img)

    # convert from grayscale image to RGB so that keypoints can be drawn in red
    img = img[:, :, np.newaxis]
    img = np.repeat(img, 3, axis=2)

    for (y, x), orientation, scale_idx, scale_size, kp_type in keypoints:
        if draw_prob < 1.0 and random.random() <= draw_prob:
            # draw the circle
            radius = int(scale_size)
            rr, cc = draw.circle_perimeter(y, x, radius, shape=img.shape)
            img[rr, cc, 0] = 1.0
            img[rr, cc, 1:] = 0

            # draw orientation
            orientation = util.quantize(orientation, [-135, -90, -45, 0, 45, 90, 135, 180])
            x_start = x
            y_start = y
            if orientation == 0:
                x_end = x + radius
                y_end = y
            elif orientation == 45:
                x_end = x + radius*(1/math.sqrt(2))
                y_end = y + radius*(1/math.sqrt(2))
            elif orientation == 90:
                x_end = x
                y_end = y + radius
            elif orientation == 135:
                x_end = x - radius*(1/math.sqrt(2))
                y_end = y - radius*(1/math.sqrt(2))
            elif orientation == 180:
                x_end = x - radius
                y_end = y
            elif orientation == -135:
                x_end = x - radius*(1/math.sqrt(2))
                y_end = y - radius*(1/math.sqrt(2))
            elif orientation == -90:
                x_end = x
                y_end = y - radius
            elif orientation == -45:
                x_end = x + radius*(1/math.sqrt(2))
                y_end = y - radius*(1/math.sqrt(2))
            x_end = np.clip(x_end, 0, width-1)
            y_end = np.clip(y_end, 0, height-1)
            rr, cc = draw.line(int(y_start), int(x_start), int(y_end), int(x_end))
            img[rr, cc, 0] = 1.0
            img[rr, cc, 1:] = 0
    img = np.clip(img, 0, 1.0)

    return img
Ejemplo n.º 4
0
def readfile(source_type, path):

    # Only point cloud (poc) is acceptable.
    assert source_type == "poc", "Only point cloud is acceptable"

    # xyz has the position x, y, z with normals nx, ny, nz in this order
    xyz = np.loadtxt(path)

    # q_xyx is a matrix of x, y, z, nx, ny, nz over all pixels
    q_xyz = util.quantize(xyz)

    return q_xyz
Ejemplo n.º 5
0
def non_maximum_suppression(g_magnitudes, g_orientations):
    """Apply Non Maximum Suppression to the gradient of an image.
    Args:
        g_magnitudes    Magnitude of the gradient of an image (in 2D).
        g_orientations  Orientations of the gradient of an image (in 2D).
    Returns:
        Modified gradient magnitudes
    """
    gm, go = g_magnitudes, g_orientations
    gm_out = np.copy(gm)
    height, width = gm.shape
    for y in range(height):
        for x in range(width):
            theta = np.degrees(go[y, x])
            theta = 180 + theta if theta < 0 else theta
            theta = util.quantize(theta, [0, 45, 90, 135, 180])

            north = gm[y - 1, x] if y > 0 else 0
            south = gm[y + 1, x] if y < height - 1 else 0
            west = gm[y, x - 1] if x > 0 else 0
            east = gm[y, x + 1] if x < width - 1 else 0
            northwest = gm[y - 1, x - 1] if y > 0 and x > 0 else 0
            northeast = gm[y - 1, x + 1] if y > 0 and x < width - 1 else 0
            southwest = gm[y + 1, x - 1] if y < height - 1 and x > 0 else 0
            southeast = gm[y + 1,
                           x + 1] if y < height - 1 and x < width - 1 else 0

            if theta == 0 or theta == 180:
                gm_out[y, x] = gm[
                    y, x] if gm[y, x] >= north and gm[y, x] >= south else 0
            elif theta == 45:
                gm_out[y, x] = gm[y, x] if gm[y, x] >= northwest and gm[
                    y, x] >= southeast else 0
            elif theta == 90:
                gm_out[y,
                       x] = gm[y,
                               x] if gm[y, x] >= west and gm[y,
                                                             x] >= east else 0
            else:  # theta == 135
                gm_out[y, x] = gm[y, x] if gm[y, x] >= northeast and gm[
                    y, x] >= southwest else 0

    return gm_out
Ejemplo n.º 6
0
def non_maximum_suppression(g_magnitudes, g_orientations):
    """Apply Non Maximum Suppression to the gradient of an image.
    Args:
        g_magnitudes    Magnitude of the gradient of an image (in 2D).
        g_orientations  Orientations of the gradient of an image (in 2D).
    Returns:
        Modified gradient magnitudes
    """
    gm, go = g_magnitudes, g_orientations
    gm_out = np.copy(gm)
    height, width = gm.shape
    for y in range(height):
        for x in range(width):
            theta = np.degrees(go[y, x])
            theta = 180 + theta if theta < 0 else theta
            theta = util.quantize(theta, [0, 45, 90, 135, 180])

            north = gm[y - 1, x] if y > 0 else 0
            south = gm[y + 1, x] if y < height - 1 else 0
            west = gm[y, x - 1] if x > 0 else 0
            east = gm[y, x + 1] if x < width - 1 else 0
            northwest = gm[y - 1, x - 1] if y > 0 and x > 0 else 0
            northeast = gm[y - 1, x + 1] if y > 0 and x < width - 1 else 0
            southwest = gm[y + 1, x - 1] if y < height - 1 and x > 0 else 0
            southeast = gm[y + 1, x + 1] if y < height - 1 and x < width - 1 else 0

            if theta == 0 or theta == 180:
                gm_out[y, x] = gm[y, x] if gm[y, x] >= north and gm[y, x] >= south else 0
            elif theta == 45:
                gm_out[y, x] = gm[y, x] if gm[y, x] >= northwest and gm[y, x] >= southeast else 0
            elif theta == 90:
                gm_out[y, x] = gm[y, x] if gm[y, x] >= west and gm[y, x] >= east else 0
            else:  # theta == 135
                gm_out[y, x] = gm[y, x] if gm[y, x] >= northeast and gm[y, x] >= southwest else 0

    return gm_out
Ejemplo n.º 7
0
    def _aux_generator(self, batch_size=16, sample_set='train', datatype = None, depthres = 256, seg_joint_res = 64):
        """ Auxiliary Generator
        Args:
            See Args section in self._generator
        """
        generated_batch = {}
        random.seed(time.time())
        generated_batch['train_img'] = np.zeros((batch_size, 256, 256, 3), dtype=np.float32)
        generated_batch['train_gtseg'] = np.zeros([batch_size, seg_joint_res, seg_joint_res], dtype = np.int8)
        generated_batch['train_gt2dheat'] = np.zeros([batch_size, seg_joint_res, seg_joint_res, self.joints_num], dtype = np.float32)
        generated_batch['train_gtjoints'] = np.zeros((batch_size, 64, 64, self.joints_num * self.Zres_joint), dtype=np.float32)
        generated_batch['train_gtdepthre'] =  np.zeros((batch_size, depthres, depthres), dtype= np.float32)
        generated_batch['train_mask'] = np.zeros([batch_size, depthres, depthres],dtype = np.bool)
        generated_batch['train_2djoints'] = np.zeros([batch_size, 2, self.joints_num ],dtype= np.float32)
        generated_batch['train_3djoints'] = np.zeros([batch_size, 3, self.joints_num ],dtype= np.float32)

        i=0
        if datatype == 'normal_dataset':
            generated_batch['normal_train_img'] = np.zeros((batch_size, self.normalres[0], self.normalres[1], 3), dtype=np.float32)
            generated_batch['normal_train_gtnormal'] = np.zeros([batch_size, self.normalres[0], self.normalres[1], 3],
                                                                dtype=np.float32)
            generated_batch['normal_train_gtdepthre'] = np.zeros((batch_size, self.normalres[0], self.normalres[1]),
                                                                 dtype=np.float32)
            generated_batch['normal_train_mask'] = np.zeros([batch_size, self.normalres[0], self.normalres[1]],
                                                            dtype=np.bool)
            while i < batch_size:
                img_name = self.filelist[self.currentindex]
                type_dir = os.path.join(self.test_dir, img_name.split('/')[-4])#random.sample(getsubfolders(self.train_dir), 1)[0])
                depth_dir = type_dir + '/depth_maps'
                normal_dir = type_dir + '/normals'

                view_type = img_name.split('/')[-2]

                depth_dir = os.path.join(depth_dir, view_type)
                normal_dir = os.path.join(normal_dir, view_type)

                index = img_name[-9:-5]
                depth_name = depth_dir + '/depth_' + index + '.npz'
                normal_name = normal_dir + '/normals_' + index + '.npz'


                bg_name = os.path.join(self.bg_dir, random.sample(os.listdir(self.bg_dir), 1)[0])
                bg_name = os.path.join(bg_name, random.sample(os.listdir(bg_name), 1)[0])

                try:
                    bg_img = io.imread(bg_name)
                except:
                    self.currentindex +=1
                    continue
                bg_img = scipy.misc.imresize(bg_img, [self.normalres[0], self.normalres[1]], interp='bilinear')
                img = io.imread(img_name)
                nmap = np.load(normal_name)['normals']
                dmap = np.load(depth_name)['depth']
                mask = dmap > 1e-4

                generated_mask = np.zeros([self.normalres[0], self.normalres[1]], dtype=np.bool)
                generated_mask[15:239, 15:239] = mask
                generated_batch['normal_train_mask'][i] = generated_mask
                img_pad = np.zeros((self.normalres[0], self.normalres[1], 3), dtype=np.uint8)
                img_pad[15: 239, 15: 239, :] = img.astype(np.float32)
                bg_img[generated_mask] = img_pad[generated_mask]

                # plt.figure()
                # plt.imshow(bg_img, aspect='auto',
                #            cmap=plt.get_cmap('jet'))
                # plt.show()

                bg_img = bg_img.astype(np.float32)
                # color augmentation
                if sample_set == 'train':
                    for j in range(3):
                        bg_img[:, :, j] = np.clip(
                            bg_img[:, :, j].astype(np.float32) / 255 * np.random.uniform(0.6, 1.4), 0.0,
                            1.0)
                else:
                    for j in range(3):
                        bg_img[:, :, j] = np.clip(bg_img[:, :, j].astype(np.float32) / 255, 0.0, 1.0)
                # print('color augmentation done!')

                # whitening rgb image
                meanstd = load_lua(self.meanRgb_dir)
                for j in range(3):
                    bg_img[:, :, j] = bg_img[:, :, j] - meanstd['mean'][j]
                    bg_img[:, :, j] = bg_img[:, :, j] / meanstd['std'][j]
                generated_batch['normal_train_img'][i,:,:,:] = bg_img

                generated_batch['normal_train_gtnormal'][i, 15:239, 15:239, :] = nmap


                if self.show:
                    plt.figure()
                    plt.imshow(generated_batch['normal_train_gtnormal'][i, :, :, 0], aspect='auto', cmap=plt.get_cmap('jet'))
                    plt.show()
                #
                # plt.figure()
                # plt.imshow(generated_batch['normal_train_gtnormal'][i, :, :, 1], aspect='auto', cmap=plt.get_cmap('jet'))
                # plt.show()
                #
                # plt.figure()
                # plt.imshow(generated_batch['normal_train_gtnormal'][i, :, :, 2], aspect='auto', cmap=plt.get_cmap('jet'))
                # plt.show()
                # print(generated_batch['normal_train_mask'].shape)
                # plt.figure()
                # plt.imshow(generated_batch['normal_train_mask'][i, :, :, 0], aspect='auto', cmap=plt.get_cmap('jet'))
                # plt.show()

                generated_batch['normal_train_gtdepthre'][i, 15:239, 15:239] = dmap

                i = i + 1

                self.currentindex+=1
                if(self.currentindex == self.datanum-1):
                    self._reset_filelist(datatype,sample_set)
            return  generated_batch


        if datatype == 'realtest':
            while i < batch_size:
                #name = random.sample(glob.glob(self.test_dir + "/*.jpg"), 1)[0]
                name = self.filelist[self.currentindex]
                testimg = io.imread(name)
                testimg = scipy.misc.imresize(testimg, [self.insize[1], self.insize[1]], interp='bilinear').astype(np.float32)
                meanstd = load_lua(self.meanRgb_dir)
                for j in range(3):
                    testimg[:, :, j] = np.clip(testimg[:, :, j].astype(np.float32) / 255.0, 0.0, 1.0)
                    testimg[:, :, j] = testimg[:, :, j] - meanstd['mean'][j]
                    testimg[:, :, j] = testimg[:, :, j] / meanstd['std'][j]
                generated_batch['train_img'][i] = cv2.resize(testimg, (self.insize[0], self.insize[1]), interpolation=cv2.INTER_NEAREST)
                i += 1
                self.currentindex += 1

                if self.show:
                    plt.figure()
                    plt.imshow(generated_batch['train_img'][0], aspect='auto', cmap=plt.get_cmap('jet'))
                    plt.show()


                if(self.currentindex == self.datanum-1):
                    self._reset_filelist('realtest','test')
            return generated_batch

        while i < batch_size:
            if datatype != 'detail_data' and datatype != 'up-3d':
                name = self.filelist[self.currentindex]
                #name = '/home/sicong/surreal/data/SURREAL/data/cmu/train/run1/ung_91_33/ung_91_33_c0001.mp4'
                cap = cv2.VideoCapture(name)
                length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
                frameindex = random.randint(1, length)
                #frameindex = 82
                if(sample_set == 'test'):
                    print('test file: ',name, 'frameindex: ', frameindex)
                cap.set(1, frameindex - 1)
                _, img_full = cap.read()
                try:
                    img_full = cv2.cvtColor(img_full, cv2.COLOR_BGR2RGB)
                except:
                    continue
                bodyinfo = sio.loadmat(name[0:-4] + '_info.mat')

            elif datatype == 'detail_data':
                name = self.filelist[self.currentindex]
                frameindex = name[-12:-8]
                name = '/home/sicong/detail_data/data/3/0235_rgb.png'
                try:
                    img_full = io.imread(name)
                except:
                    self.currentindex += 1
                    continue

            elif datatype == 'up-3d':
                if sample_set == 'train':
                    info_dir = self.train_dir + '/pose_prepared/91/500/up-p91/'
                    seg_dir = self.train_dir + '/segment/up-s31/s31/'
                elif sample_set == 'valid':
                    info_dir = self.valid_dir + '/pose_prepared/91/500/up-p91/'
                    seg_dir = self.valid_dir + '/segment/up-s31/s31/'
                elif sample_set == 'test':
                    info_dir = self.test_dir + '/pose_prepared/91/500/up-p91/'
                    seg_dir = self.test_dir + '/segment/up-s31/s31/'

                name = self.filelist[self.currentindex]
                if(sample_set == 'test'):
                    print('test file: ',name)
                # name = '/media/sicong/a86d93af-1a2e-469b-972c-f819c47cd5ee/datasets/pose_prepared/91/500/up-p91/04877_image.png'
                frameindex = name[-15:-10]

                try:
                    img_full = io.imread(name)
                except:
                    self.currentindex +=1
                    continue
                try:
                    bodyinfo = sio.loadmat(info_dir + frameindex+ '_info.mat')
                except:
                    self.currentindex += 1
                    continue
            if self.show:
                img = Image.fromarray(img_full, 'RGB')
                img.show()

            if datatype != 'detail_data':
                # load 2d joints to determine the bounding box
                # [2 x njoints]
                if datatype != 'up-3d':
                    if bodyinfo is None:
                        self.currentindex += 1
                        continue
                    joints2dfull = bodyinfo['joints2D']
                    if joints2dfull is None:
                        self.currentindex += 1
                        continue
                    if len(joints2dfull.shape) < 3:
                        self.currentindex += 1
                        continue
                    if frameindex - 1 >= joints2dfull.shape[2]:
                        self.currentindex += 1
                        continue
                    joints2d = joints2dfull[:, self.joints_subset, frameindex - 1].astype(np.int64)

                    joints3dfull = bodyinfo['joints3D']
                    if joints3dfull is None:
                        self.currentindex += 1
                        continue
                    if frameindex - 1 >= joints2dfull.shape[2]:
                        self.currentindex += 1
                        continue
                    joints3d = joints3dfull[:, self.joints_subset, frameindex - 1]

                    generated_batch['train_2djoints'][i,:] = joints2d
                    generated_batch['train_3djoints'][i,:] = joints3d

                    depth_full = sio.loadmat(name[0:-4] + '_depth.mat')['depth_' + str(frameindex)]
                elif datatype == 'up-3d':
                    if bodyinfo is None:
                        self.currentindex += 1
                        continue
                    joints2dfull = bodyinfo['joints2D']
                    if joints2dfull is None:
                        self.currentindex += 1
                        continue
                    if len(joints2dfull.shape) < 2:
                        self.currentindex += 1
                        continue
                    joints2d = joints2dfull[:, self.joints_subset].astype(np.int64)
                    joints3dfull = np.transpose(bodyinfo['joints3D'])
                    if joints3dfull is None:
                        self.currentindex += 1
                        continue
                    joints3d = joints3dfull[:, self.joints_subset]

                    depth_full = sio.loadmat(info_dir + frameindex+ '_depth.mat')['depth']

                #set pelvis as the original point
                camLoc = bodyinfo['camLoc'][0]
                if datatype == 'up-3d':
                    # camlocation = camLoc[2]
                    # joints3d[2, :] = camlocation - joints3d[2, :]
                    dPelvis = joints3d[2, 0]
                else:
                    camlocation = camLoc
                    joints3d[0, :] = camlocation - joints3d[0, :]
                    dPelvis = joints3d[0, 0]

                if datatype != 'up-3d':
                    segm_raw = sio.loadmat(name[0:-4] + '_segm.mat')['segm_'+str(frameindex)]

                    segm_full = util.changeSegmIx(segm_raw,
                                                  [2, 12, 9, 2, 13, 10, 2, 14, 11, 2, 14, 11, 2, 2, 2, 1, 6, 3, 7, 4, 8,
                                                   5, 8,
                                                   5]).astype(np.int8)

                else:
                    segm_raw = cv2.imread(seg_dir+ frameindex + '_ann_vis.png')
                    segm_full = util.up3dtosurreal(segm_raw)

                if self.show:
                    plt.figure()
                    plt.imshow(segm_full, aspect='auto', cmap=plt.get_cmap('jet'))
                    plt.show()

                if datatype == 'up-3d':
                    quantized_joints3d, _ = util.quantize(joints3d[2, :], dPelvis, self.stp_joint, self.Zres_joint)
                    quantized_joints3d = quantized_joints3d * -1
                    relative_depth, _ = util.relative_up3d(depth_full, dPelvis, self.stp, self.Zres)  # self.halfrange
                elif datatype != 'detail_data':
                    quantized_joints3d, _ = util.quantize(joints3d[0, :], dPelvis, self.stp_joint, self.Zres_joint)
                    quantized_joints3d = quantized_joints3d * -1
                    relative_depth, _ = util.relative(depth_full,dPelvis, self.stp, self.Zres) #self.halfrange

                # TODO: 1. resize quantized_depth 2. output dense continuous relative depth in util.quantize
                if self.show:
                    plt.figure()
                    plt.imshow(depth_full, aspect='auto', cmap=plt.get_cmap('jet'))
                    plt.show()
                if self.show:
                    plt.figure()
                    plt.imshow(relative_depth, aspect='auto', cmap=plt.get_cmap('jet'))
                    plt.show()
            else:
                depth_full = io.imread(name[0:-8] + '_depth.png')
                depthcount = np.sum(depth_full > 100)
                if depthcount < 100 * 100:
                    self.currentindex += 1
                    continue

            if datatype != 'detail_data':
                # crop, scale
                rot = 0
                scale = util.getScale(joints2d)
                center = util.getCenter(joints2d)
            else:
                # crop, scale
                rot = 0
                scale = util.getScale_detail(depth_full)
                center = util.getCenter_detail(depth_full)

            if (center[0] < 1 or center[1] < 1 or center[1] > img_full.shape[0] or center[0] > img_full.shape[1]):
                self.currentindex +=1
                continue


## for rgb image
            if datatype != 'up-3d' and datatype!= 'detail_data':
                img = util.cropfor3d(img_full, center, scale, rot, self.insize[1], 'bilinear')
            elif datatype == 'detail_data':
                img = util_detail.cropfor3d(img_full, center, scale, rot, self.insize[1], 'bilinear')
            elif datatype == 'up-3d':
                norm_factor = np.array([self.insize[1]/img_full.shape[1], self.insize[1]/img_full.shape[0]], dtype=np.float32)
                img = scipy.misc.imresize(img_full, [self.insize[1], self.insize[1]], interp= 'bilinear')
                badexample = False
                for j in range(joints2d.shape[1]):
                    joints2d_rescaled = np.multiply(joints2d[:,j],norm_factor).astype(np.int64)
                    if joints2d_rescaled[0] < 0 or joints2d_rescaled[0] > 256 or joints2d_rescaled[1] < 0 or joints2d_rescaled[1] > 256:
                        badexample = True
                if badexample:
                    self.currentindex += 1
                    continue

            if img is None:
                self.currentindex+=1
                continue
            if (img.shape[0] == 0 or img.shape[1] == 0):
                self.currentindex+=1
                continue

            if self.show:
                imgnew = Image.fromarray(img, 'RGB')
                imgnew.show()

            # color augmentation
            img_bak = img
            img = img.astype(np.float32)
            if sample_set == 'train':
                for j in range(3):
                    img[:, :, j] = np.clip(img[:, :, j].astype(np.float32) / 255 * np.random.uniform(0.6, 1.4), 0.0,
                                           1.0)
            else:
                for j in range(3):
                    img[:, :, j] = np.clip(img[:, :, j].astype(np.float32) / 255, 0.0, 1.0)
            # print('color augmentation done!')

            # whitening rgb image
            meanstd = load_lua(self.meanRgb_dir)
            for j in range(3):
                img[:, :, j] = img[:, :, j] - meanstd['mean'][j]
                img[:, :, j] = img[:, :, j] / meanstd['std'][j]

            generated_batch['train_img'][i] = img

## for depth
            if datatype == 'detail_data':
                depm_continue = util_detail.cropfor3d(depth_full,center,scale,rot,self.insize[1],'bilinear')

                if self.show:
                    plt.figure()
                    plt.imshow(depth_full, aspect='auto', cmap=plt.get_cmap('jet'))
                    plt.show()

                if self.show:
                    plt.figure()
                    plt.imshow(depm_continue, aspect='auto', cmap=plt.get_cmap('jet'))
                    plt.show()
                mask =depm_continue>100
                depm_continue[depm_continue < 100] = 15 * 1000.0
                final_depth = depm_continue/1000.0
                median_value =np.median(final_depth[final_depth<5])
                final_depth = final_depth - median_value + 0.10
                final_depth[final_depth>5] = 0.60
                generated_batch['train_gtdepthre'][i, :, :] = final_depth

                mask = ndimage.binary_erosion(mask).astype(mask.dtype)
                generated_batch['train_mask'][i,:,:] = mask

            elif datatype == 'up-3d':
                depm_continue = cv2.resize(relative_depth.astype(np.float32), (depthres, depthres), interpolation=cv2.INTER_NEAREST)
                generated_batch['train_gtdepthre'][i, :, :] = depm_continue
                mask = depm_continue<0.59

                mask = ndimage.binary_erosion(mask).astype(mask.dtype)
                generated_batch['train_mask'][i,:,:] = mask

            else:
                depm_continue = util.cropdepth(relative_depth,center,scale,rot,self.insize[1],0.60)
                generated_batch['train_gtdepthre'][i, :, :] = cv2.resize(depm_continue,(depthres, depthres),interpolation=cv2.INTER_NEAREST)
                mask = depm_continue<0.59

                mask = ndimage.binary_erosion(mask).astype(mask.dtype)
                generated_batch['train_mask'][i,:,:] = mask




            if self.show:
                plt.figure()
                plt.imshow(generated_batch['train_gtdepthre'][i, :, :], aspect='auto', cmap=plt.get_cmap('jet'))
                plt.show()

            # if self.show:
            #     plt.figure()
            #     plt.imshow(mask, aspect='auto', cmap=plt.get_cmap('jet'))
            #     plt.show()

## for 2d segmentation

            if datatype == 'up-3d':
                segm = cv2.resize(segm_full, (seg_joint_res, seg_joint_res),
                                                            interpolation=cv2.INTER_NEAREST)
                generated_batch['train_gtseg'][i,:,:] = segm

            elif datatype != 'detail_data':
                segm = util.cropfor3d(segm_full, center, scale, rot, self.insize[1],'nearest')
                generated_batch['train_gtseg'][i,:,:] = cv2.resize(segm, (seg_joint_res, seg_joint_res),
                                                            interpolation=cv2.INTER_NEAREST)
                if self.show:
                    plt.figure()
                    plt.imshow(segm, aspect='auto', cmap=plt.get_cmap('jet'))
                    plt.show()

## for 2d joints

            if datatype != 'detail_data':
                # TODO: create 2d heatmaps
                sigma_2d_inscale = math.floor(2 * self.insize[0]/self.outsize[0])
                out_2d = np.zeros([self.insize[0], self.insize[1], self.joints_num])

                for j in range(self.joints_num):
                    if datatype == 'up-3d':
                        #pt = util.transform(joints2d[:, j], center, scale, 0, self.insize[0], False)
                        pt = np.multiply(joints2d[:, j], norm_factor).astype(np.int64)
                        # print('joints: ', joints2d[:, j], 'pt: ', pt)
                    else:
                        pt = util.transform(joints2d[:, j], center, scale, 0, self.insize[0], False)
                    heat_slice = util.Drawgaussian2D(img,pt,sigma_2d_inscale)
                    # if np.sum(heat_slice) > 1e-2:
                    #     heat_slice /= np.sum(heat_slice)
                    # else:
                    #     heat_slice *= 0
                    #print('heat_slice.shape',heat_slice.shape)

                    out_2d[:, :, j] = heat_slice
                    # if self.show:
                    #     plt.figure()
                    #     plt.imshow(heat_slice, aspect='auto', cmap=plt.get_cmap('jet'))
                    #     plt.show()

                out_2d = cv2.resize(out_2d,(seg_joint_res,seg_joint_res),interpolation=cv2.INTER_NEAREST)
                generated_batch['train_gt2dheat'][i] = out_2d
                if self.show:
                    # img4show = img
                    # for j in range(3):
                    #     img4show[:, :, j] = img4show[:, :, j] - meanstd['mean'][j]
                    #     img4show[:, :, j] = img4show[:, :, j] / meanstd['std'][j]
                    # img4show = img4show * 255.0
                    visualizer.draw2dskeleton(img_bak.astype(np.uint8), out_2d)

            
# for 3d joints

            #print('draw3d---------------------------------------------------')
            if datatype != 'detail_data':
                out = np.zeros([self.outsize[0], self.outsize[1], self.joints_num * self.Zres_joint])
                sigma_2d = 2
                size_z = 2 * math.floor((6* sigma_2d * self.Zres_joint / self.outsize[0] +1) / 2) + 1
                for j in range(self.joints_num):
                    #if joints2d[1,j] >= img_full.shape[0] or joints2d[0,j] >=img_full.shape[1] or joints2d[1,j]<0 or joints2d[0,j]<0:
                        #continue
                    z = quantized_joints3d[j]
                    if datatype == 'up-3d':
                        pt = np.multiply(joints2d[:, j], norm_factor/4).astype(np.int64)
                    else:
                        pt = util.transform(joints2d[:, j], center, scale, 0, self.outsize[0], False)
                    out[:,:,j * self.Zres_joint : (j+1) * self.Zres_joint] = util.Drawguassian3D(out[:,:,j * self.Zres_joint : (j+1) * self.Zres_joint], pt, z , sigma_2d, size_z)

                generated_batch['train_gtjoints'][i] = out
                if self.show:
                    visualizer.draw3dskeleton(self.joints_num,self.Zres_joint,out)
            i = i+1
            self.currentindex +=1
            if(self.currentindex==self.datanum-1):
                self._reset_filelist(datatype,sample_set)


        return generated_batch
                                       device_ids=range(
                                           torch.cuda.device_count()))
    cudnn.benchmark = True

params = util.paramsGet(net)
tmp = (params.data != 0).sum()
print("Ratio of nonzero value : ", tmp.item() / params.size()[0])
print("Number of nonzero value : ", tmp.item())
print("Number of value", params.size()[0])

net2 = util.netMaskMul(net2, mask_null, isbias=1)
net2 = util.addNetwork(net2, net, isbias=1)
net2 = util.swapBatch(net2, net)

if args.fixed:
    net = util.quantize(net, args.pprec)
    net2 = util.quantize(net2, args.pprec)
    net_origin = util.quantize(net_origin, args.pprec)


def calcDiff(tensor):
    pass


def evalMetric(net, net_origin):
    params_net = utils.paramsGet(net)
    params_net_origin = utils.paramsGet(net_origin)
    print(params_net.size())
    print(params_net_origin.size())

def readfile(source_type, path, n_lights):
    # xyz has the position x, y, z with normals nx, ny, nz in this order
    if source_type == "poc":

        xyz = np.loadtxt(path)

        # q_xyx is a matrix of x, y, z, nx, ny, nz over all pixels
        q_xyz = util.quantize(xyz)

        # N as a matrix of Nx, Ny, Nz is extracted
        N = np.matrix(q_xyz[:, 3:6])

        # The light direction (uniformly distributed) is generated
        L = np.matrix(util.generateLight(n_lights))

        # Measurement
        M = np.asarray(N * L.T)

        # Noise (not Gaussian) makes the measurement (noise) times larger (no noise if zero)
        noise = 5
        noise_ratio = 0.1

        # Add noise on N
        n_noise = int(M.size * noise_ratio)
        noise_index = np.array([
            np.random.choice(np.arange(M.shape[0]), n_noise),
            np.random.choice(np.arange(M.shape[1]), n_noise)
        ])

        M[noise_index[0],
          noise_index[1]] = noise * M[noise_index[0], noise_index[1]]

    elif source_type == "image":

        dirname = path
        L = np.loadtxt(dirname + "lighting.txt").T
        '''
        ground.txt is ground truth of normal map that was extracted 
        from mat file (matlab). If the image size is m * n, the file has
        m * 3n matrix, which horizontally aligns m * n matrices of nx, ny and nz. 
        '''
        N = np.loadtxt(dirname + "ground.txt")
        pixels = int(N.shape[0]), int(N.shape[1] / 3)

        N = np.vstack([
            N[:, int(pixels[1] * i):int(pixels[1] * (i + 1))].reshape((-1, ))
            for i in range(3)
        ]).T

        imgfiles = glob.glob(dirname + "*.*.png")
        # Read img file as a gray-scale image
        M = np.vstack([
            cv2.imread(imgfiles[i], flags=0).reshape((-1, ))
            for i in range(len(imgfiles))
        ]).T

        # read mask file
        mask = cv2.imread(dirname + "mask.png", flags=0)

        mask_flat = mask.reshape((-1, ))
        mask_index = np.where(mask_flat == 255)

        # Extract only pixels that are defined in the mask (color is 255 in mask.png)
        N = N[mask_index]
        M = M[mask_index]

    return N, L, M, mask