Ejemplo n.º 1
0
 def get_landmark(self, img, label, dataset, use_essh):
     if use_essh:
         ret = self.detector.detect(img, threshold=0.4)
         if ret is None or ret.shape[0] == 0:
             return None, None
         bindex = self.get_maxpos(img, ret)
         face = ret[bindex]
         bbox = face[0:4]
         points = face[5:15].reshape(5, 2)
         # b = bbox
         # cv2.rectangle(img, (int(b[0]), int(b[1])), (int(b[2]), int(b[3])), (255, 255, 255))
         # for p in landmark:
         #   cv2.circle(img, (int(p[0]), int(p[1])), 1, (0, 0, 255), 2)
         # cv2.imshow("detection result", img)
         # cv2.waitKey(0)
         # for i in range(bbox.shape[0]):
         rimg, label2, trans1 = img_helper.preprocess(
             img, points, img.shape[0])
         ret2 = self.detector.detect(rimg, threshold=0.4)
         if ret2 is None or ret2.shape[0] == 0:
             return None, None
         bindex2 = self.get_maxpos(rimg, ret2)
         rimg, trans2 = img_helper.transform2(rimg, None,
                                              self.image_size[0],
                                              ret2[bindex2, 0:4], dataset)
     else:
         rimg, label2, trans1 = img_helper.preprocess(
             img, label, img.shape[0])
         rimg, trans2 = img_helper.transform2(rimg, label2,
                                              self.image_size[0], None,
                                              dataset)
     trans = self.trans_dot(trans1, trans2)
     # cv2.imshow("rimg", rimg)
     # cv2.waitKey(0)
     # img2 = cv2.cvtColor(rimg, cv2.COLOR_BGR2RGB)
     img2 = np.transpose(rimg, (2, 0, 1))  #3*128*128, RGB
     input_blob = np.zeros((1, 3, self.image_size[1], self.image_size[0]),
                           dtype=np.uint8)
     input_blob[0] = img2
     data = mx.nd.array(input_blob)
     db = mx.io.DataBatch(data=(data, ))
     self.model.forward(db, is_train=False)
     alabel = self.model.get_outputs()[-1].asnumpy()[0]
     IM = cv2.invertAffineTransform(trans)
     landmark = np.zeros((68, 2), dtype=np.float32)
     for i in xrange(alabel.shape[0]):
         a = cv2.resize(alabel[i], (self.image_size[1], self.image_size[0]))
         ind = np.unravel_index(np.argmax(a, axis=None), a.shape)
         point = (ind[1], ind[0], 1.0)  #w, h
         point = np.dot(IM, point)
         landmark[i] = point[0:2]
         npt = img_helper.transform_pt(label[i], trans)
         if config.landmark_type == '2d':
             npt = np.floor(npt)
         else:
             npt = np.round(npt)
         point = (npt[0], npt[1], 1.0)
         point = np.dot(IM, point)
         label[i] = point[0:2]
     return landmark, label
Ejemplo n.º 2
0
    def get_data(self, data, label, annot):
        if self.vis:
            self.img_num += 1
            #if self.img_num<=self.vis:
            #  filename = './vis/raw_%d.jpg' % (self.img_num)
            #  print('save', filename)
            #  draw = data.copy()
            #  for i in xrange(label.shape[0]):
            #    cv2.circle(draw, (label[i][1], label[i][0]), 1, (0, 0, 255), 2)
            #  scipy.misc.imsave(filename, draw)

        rotate = 0
        #scale = 1.0
        if 'scale' in annot:
            scale = annot['scale']
        else:
            scale = max(data.shape[0], data.shape[1])
        if 'center' in annot:
            center = annot['center']
        else:
            center = np.array((data.shape[1] / 2, data.shape[0] / 2))
        max_retry = 3
        if self.aug_level == 0:  #validation mode
            max_retry = 6
        retry = 0
        found = False
        base_scale = scale
        while retry < max_retry:
            retry += 1
            succ = True
            _scale = base_scale
            if self.aug_level > 0:
                rotate = np.random.randint(-40, 40)
                scale_config = 0.2
                #rotate = 0
                #scale_config = 0.0
                scale_ratio = min(
                    1 + scale_config,
                    max(1 - scale_config,
                        (np.random.randn() * scale_config) + 1))
                _scale = int(base_scale * scale_ratio)
                #translate = np.random.randint(-5, 5, size=(2,))
                #center += translate
                data_out, trans = img_helper.transform(data, center,
                                                       self.input_img_size,
                                                       _scale, rotate)
            else:
                data, label, trans = img_helper.preprocess(
                    data, label, data.shape[0])
                data_out, trans = img_helper.transform2(
                    data, label, self.input_img_size)
            #data_out = img_helper.crop2(data, center, _scale, (self.input_img_size, self.input_img_size), rot=rotate)
            label_out = np.zeros(self.label_shape, dtype=np.float32)
            #print('out shapes', data_out.shape, label_out.shape)
            for i in xrange(label.shape[0]):
                pt = label[i].copy()
                #pt = pt[::-1]
                npt = img_helper.transform_pt(pt, trans)
                if npt[0] >= data_out.shape[1] or npt[1] >= data_out.shape[
                        0] or npt[0] < 0 or npt[1] < 0:
                    succ = False
                    #print('err npt', npt)
                    break
                if config.losstype == 'heatmap':
                    pt_scale = float(
                        self.output_label_size) / self.input_img_size
                    npt *= pt_scale
                    npt = npt.astype(np.int32)
                    img_helper.gaussian(label_out[i], npt, config.gaussian)
                else:
                    label_out[i] = (npt / self.input_img_size)
                #print('before gaussian', label_out[i].shape, pt.shape)
                #trans = img_helper.transform(pt, center, _scale, (self.output_label_size, self.output_label_size), rot=rotate)
                #print(trans.shape)
                #if not img_helper.gaussian(label_out[i], trans, _g):
                #    succ = False
                #    break
            if not succ:
                if self.aug_level == 0:
                    base_scale += 20
                continue

            flip_data_out = None
            flip_label_out = None
            if config.net_coherent:
                flip_data_out, flip_label_out = self.get_flip(
                    data_out, label_out)
            elif ((self.aug_level > 0 and np.random.rand() < 0.5)
                  or self.force_mirror):  #flip aug
                flip_data_out, flip_label_out = self.get_flip(
                    data_out, label_out)
                data_out, label_out = flip_data_out, flip_label_out

            found = True
            break

        #self.stats[0]+=1
        if not found:
            #self.stats[1]+=1
            #print('find aug error', retry)
            #print(self.stats)
            #print('!!!ERR')
            return None
        #print('found with scale', _scale, rotate)

        if self.vis > 0 and self.img_num <= self.vis:
            print('crop', data.shape, center, _scale, rotate, data_out.shape)
            filename = './vis/cropped_%d.jpg' % (self.img_num)
            print('save', filename)
            draw = data_out.copy()
            alabel = label_out.copy()
            for i in xrange(label.shape[0]):
                a = cv2.resize(alabel[i],
                               (self.input_img_size, self.input_img_size))
                ind = np.unravel_index(np.argmax(a, axis=None), a.shape)
                cv2.circle(draw, (ind[1], ind[0]), 1, (0, 0, 255), 2)
            scipy.misc.imsave(filename, draw)
            filename = './vis/raw_%d.jpg' % (self.img_num)
            scipy.misc.imsave(filename, data)

        return data_out, label_out, flip_data_out, flip_label_out