def gen_landmark_data(srcTxt, net, augment=False):
    '''
    srcTxt: each line is: 0=path, 1-4=bbox, 5-14=landmark 5points
    net: PNet or RNet or ONet
    augment: if enable data augmentation
    '''
    print(">>>>>> Start landmark data create...Stage: %s"%(net))
    srcTxt = os.path.join(rootPath, srcTxt)
    saveFolder = os.path.join(rootPath, "tmp/data/%s/"%(net))
    saveImagesFolder = os.path.join(saveFolder, "landmark")
    sizeOfNet = {"pnet": 12, "rnet": 24, "onet": 48}
    if net not in sizeOfNet:
        raise Exception("The net type error!")
    if not os.path.isdir(saveImagesFolder):
        os.makedirs(saveImagesFolder)
    saveF = open(join(saveFolder, "landmark.txt"), 'w')
    imageCnt = 0
    # image_path bbox landmark(5*2)
    for (imgPath, bbox, landmarkGt) in getBboxLandmarkFromTxt(srcTxt):
        F_imgs = []
        F_landmarks = []        
        img = cv2.imread(imgPath)
        assert(img is not None)
        img_h, img_w, img_c = img.shape
        gt_box = np.array([bbox.left, bbox.top, bbox.right, bbox.bottom])
        f_face = img[bbox.top: bbox.bottom+1, bbox.left: bbox.right+1]
        #print("-----{}".format(sizeOfNet[net]))
        try:
            f_face = cv2.resize(f_face, (sizeOfNet[net], sizeOfNet[net]))
        except:
            print("An exception occurred")
            continue        
        landmark = np.zeros((landmark_number, 2))
        #normalize
        for index, one in enumerate(landmarkGt):
            rv = ((one[0]-gt_box[0])/(gt_box[2]-gt_box[0]), (one[1]-gt_box[1])/(gt_box[3]-gt_box[1]))
            landmark[index] = rv
        F_imgs.append(f_face)
        F_landmarks.append(landmark.reshape(landmark_number*2))
        landmark = np.zeros((landmark_number, 2))
        if augment:
            x1, y1, x2, y2 = gt_box
            #gt's width
            gt_w = x2 - x1 + 1
            #gt's height
            gt_h = y2 - y1 + 1        
            if max(gt_w, gt_h) < 40 or x1 < 0 or y1 < 0:
                continue
            #random shift
            for i in range(landmark_number*2):
                bbox_size = np.random.randint(int(min(gt_w, gt_h) * 0.8), np.ceil(1.25 * max(gt_w, gt_h)))
                delta_x = np.random.randint(-gt_w * 0.2, gt_w * 0.2)
                delta_y = np.random.randint(-gt_h * 0.2, gt_h * 0.2)
                nx1 = max(x1+gt_w/2-bbox_size/2+delta_x,0)
                ny1 = max(y1+gt_h/2-bbox_size/2+delta_y,0)
                
                nx2 = nx1 + bbox_size
                ny2 = ny1 + bbox_size
                if nx2 > img_w or ny2 > img_h:
                    continue
                crop_box = np.array([nx1,ny1,nx2,ny2])
                cropped_im = img[ny1:ny2+1,nx1:nx2+1,:]
                resized_im = cv2.resize(cropped_im, (sizeOfNet[net], sizeOfNet[net]))
                #cal iou
                iou = IoU(crop_box, np.expand_dims(gt_box,0))
                if iou <= 0.65:
                    continue
                F_imgs.append(resized_im)
                #normalize
                for index, one in enumerate(landmarkGt):
                    rv = ((one[0]-nx1)/bbox_size, (one[1]-ny1)/bbox_size)
                    landmark[index] = rv
                F_landmarks.append(landmark.reshape(landmark_number*2))
                landmark = np.zeros((landmark_number, 2))
                landmark_ = F_landmarks[-1].reshape(-1,2)
                bbox = BBox([nx1,ny1,nx2,ny2])                    

                #mirror                    
                if random.choice([0,1]) > 0:
                    face_flipped, landmark_flipped = flip(resized_im, landmark_)
                    face_flipped = cv2.resize(face_flipped, (sizeOfNet[net], sizeOfNet[net]))
                    #c*h*w
                    F_imgs.append(face_flipped)
                    F_landmarks.append(landmark_flipped.reshape(landmark_number*2))
                #rotate
                if random.choice([0,1]) > 0:
                    face_rotated_by_alpha, landmark_rotated = rotate(img, bbox, \
                                                                     bbox.reprojectLandmark(landmark_), 5)#逆时针旋转
                    #landmark_offset
                    landmark_rotated = bbox.projectLandmark(landmark_rotated)
                    face_rotated_by_alpha = cv2.resize(face_rotated_by_alpha, (sizeOfNet[net], sizeOfNet[net]))
                    F_imgs.append(face_rotated_by_alpha)
                    F_landmarks.append(landmark_rotated.reshape(landmark_number*2))
                
                    #flip
                    face_flipped, landmark_flipped = flip(face_rotated_by_alpha, landmark_rotated)
                    face_flipped = cv2.resize(face_flipped, (sizeOfNet[net], sizeOfNet[net]))
                    F_imgs.append(face_flipped)
                    F_landmarks.append(landmark_flipped.reshape(landmark_number*2))
                
                #inverse clockwise rotation
                if random.choice([0,1]) > 0: 
                    face_rotated_by_alpha, landmark_rotated = rotate(img, bbox, \
                                                                     bbox.reprojectLandmark(landmark_), -5)#clockwise rotation
                    landmark_rotated = bbox.projectLandmark(landmark_rotated)
                    face_rotated_by_alpha = cv2.resize(face_rotated_by_alpha, (sizeOfNet[net], sizeOfNet[net]))
                    F_imgs.append(face_rotated_by_alpha)
                    F_landmarks.append(landmark_rotated.reshape(landmark_number*2))
                
                    face_flipped, landmark_flipped = flip(face_rotated_by_alpha, landmark_rotated)
                    face_flipped = cv2.resize(face_flipped, (sizeOfNet[net], sizeOfNet[net]))
                    F_imgs.append(face_flipped)
                    F_landmarks.append(landmark_flipped.reshape(landmark_number*2))
        F_imgs, F_landmarks = np.asarray(F_imgs), np.asarray(F_landmarks)
        for i in range(len(F_imgs)):
            path = os.path.join(saveImagesFolder, "%d.jpg"%(imageCnt))
            cv2.imwrite(path, F_imgs[i])
            landmarks = map(str, list(F_landmarks[i]))
            saveF.write(path + " -2 " + " ".join(landmarks)+"\n")
            imageCnt += 1
        printStr = "\rCount: {}".format(imageCnt)
        sys.stdout.write(printStr)
        sys.stdout.flush()
    saveF.close()
    print "\nLandmark create done!"
Ejemplo n.º 2
0
def gen_hard_bbox_pnet(srcDataSet, srcAnnotations):
    srcDataSet = os.path.join(rootPath, srcDataSet)
    srcAnnotations = os.path.join(rootPath, srcAnnotations)
    saveFolder = os.path.join(rootPath, "tmp/data/pnet/")
    print(">>>>>> Gen hard samples for pnet...")
    typeName = ["pos", "neg", "part"]
    saveFiles = {}
    for tp in typeName:
        _saveFolder = os.path.join(saveFolder, tp)
        if not os.path.isdir(_saveFolder):
            os.makedirs(_saveFolder)
        saveFiles[tp] = open(os.path.join(saveFolder, "{}.txt".format(tp)),
                             'w')

    annotationsFile = open(srcAnnotations, "r")
    pIdx = 0  # positive
    nIdx = 0  # negative
    dIdx = 0  # dont care
    idx = 0
    for annotation in annotationsFile:
        annotation = annotation.strip().split(' ')
        # image path
        imPath = annotation[0]
        # boxed change to float type, bbox's type is array
        bbox = map(float, annotation[1:])
        bbox = map(int, bbox)
        # gt. each row mean bounding box, each image not only have one gt
        boxes = np.array(bbox, dtype=np.float32).reshape(-1, 4)
        #load image
        img = cv2.imread(os.path.join(srcDataSet, imPath + '.jpg'))
        idx += 1
        height, width, channel = img.shape
        # draw rectangle
        # cv2.rectangle(img, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (28, 28, 28), 4)

        # 1. NEG: random to crop negative sample image
        negNum = 0
        while negNum < 50:
            size = np.random.randint(12, min(width, height) / 2)
            # top_left
            nx = np.random.randint(0, width - size)
            ny = np.random.randint(0, height - size)
            # random crop
            crop_box = np.array([nx, ny, nx + size, ny + size])
            # cal iou and iou must below 0.3 for neg sample
            iou = IoU(crop_box, boxes)
            if np.max(iou) >= 0.3:
                continue
            # crop sample image
            cropped_im = img[ny:ny + size, nx:nx + size, :]
            # cv2.rectangle(img, (nx, ny), (nx+size, ny+size), (0, 0, 255), 4)
            resized_im = cv2.resize(cropped_im, (12, 12),
                                    interpolation=cv2.INTER_LINEAR)
            # now to save it
            save_file = os.path.join(saveFolder, "neg", "%s.jpg" % nIdx)
            saveFiles['neg'].write(save_file + ' 0\n')
            cv2.imwrite(save_file, resized_im)
            nIdx += 1
            negNum += 1
        for box in boxes:
            # box (x_left, y_top, x_right, y_bottom)
            x1, y1, x2, y2 = box
            #bbox's width and height
            w, h = x2 - x1 + 1, y2 - y1 + 1
            # ignore small faces
            # in case the ground truth boxes of small faces are not accurate
            if max(w, h) < 40 or x1 < 0 or y1 < 0:
                continue
            # 2. NEG: random to crop sample image in bbox inside
            for i in range(5):
                size = np.random.randint(12, min(width, height) / 2)
                # delta_x and delta_y are offsets of (x1, y1)
                delta_x = np.random.randint(max(-size, -x1), w)
                delta_y = np.random.randint(max(-size, -y1), h)
                nx1 = int(max(0, x1 + delta_x))
                ny1 = int(max(0, y1 + delta_y))
                if nx1 + size > width or ny1 + size > height:
                    continue
                crop_box = np.array([nx1, ny1, nx1 + size, ny1 + size])

                Iou = IoU(crop_box, boxes)
                if np.max(iou) >= 0.3:
                    continue
                cropped_im = img[ny1:ny1 + size, nx1:nx1 + size, :]
                # cv2.rectangle(img, (nx1, ny1), (nx1 + size, ny1 + size), (144, 20, 255), 4)
                resized_im = cv2.resize(cropped_im, (12, 12),
                                        interpolation=cv2.INTER_LINEAR)
                save_file = os.path.join(saveFolder, "neg", "%s.jpg" % nIdx)
                saveFiles['neg'].write(save_file + ' 0\n')
                cv2.imwrite(save_file, resized_im)
                nIdx += 1
            # 3. POS and PART
            for i in range(20):
                # pos and part face size [minsize*0.8,maxsize*1.25]
                size = np.random.randint(int(min(w, h) * 0.8),
                                         np.ceil(1.25 * max(w, h)))
                # delta here is the offset of box center
                delta_x = np.random.randint(-w * 0.2, w * 0.2)
                delta_y = np.random.randint(-h * 0.2, h * 0.2)
                #show this way: nx1 = max(x1+w/2-size/2+delta_x)
                nx1 = max(x1 + w / 2 + delta_x - size / 2, 0)
                #show this way: ny1 = max(y1+h/2-size/2+delta_y)
                ny1 = max(y1 + h / 2 + delta_y - size / 2, 0)
                nx2 = nx1 + size
                ny2 = ny1 + size

                if nx2 > width or ny2 > height:
                    continue
                crop_box = np.array([nx1, ny1, nx2, ny2])
                #yu gt de offset
                offset_x1 = (x1 - nx1) / float(size)
                offset_y1 = (y1 - ny1) / float(size)
                offset_x2 = (x2 - nx2) / float(size)
                offset_y2 = (y2 - ny2) / float(size)
                #crop
                cropped_im = img[int(ny1):int(ny2), int(nx1):int(nx2), :]
                #resize
                resized_im = cv2.resize(cropped_im, (12, 12),
                                        interpolation=cv2.INTER_LINEAR)

                box_ = box.reshape(1, -1)
                if IoU(crop_box, box_) >= 0.65:
                    save_file = os.path.join(saveFolder, "pos",
                                             "%s.jpg" % pIdx)
                    saveFiles['pos'].write(
                        save_file + ' 1 %.2f %.2f %.2f %.2f\n' %
                        (offset_x1, offset_y1, offset_x2, offset_y2))
                    cv2.imwrite(save_file, resized_im)
                    # cv2.rectangle(img, (int(nx1), int(ny1)), (int(nx2), int(ny2)), (78, 238, 148), 4)
                    pIdx += 1
                elif IoU(crop_box, box_) >= 0.4:
                    save_file = os.path.join(saveFolder, "part",
                                             "%s.jpg" % dIdx)
                    saveFiles['part'].write(
                        save_file + ' -1 %.2f %.2f %.2f %.2f\n' %
                        (offset_x1, offset_y1, offset_x2, offset_y2))
                    cv2.imwrite(save_file, resized_im)
                    # cv2.rectangle(img, (int(nx1), int(ny1)), (int(nx2), int(ny2)), (0, 238, 238), 4)
                    dIdx += 1
        printStr = "\r[{}] pos: {}  neg: {}  part:{}".format(
            idx, pIdx, nIdx, dIdx)
        # cv2.imwrite("001_new.jpg",img)
        sys.stdout.write(printStr)
        sys.stdout.flush()
    for f in saveFiles.values():
        f.close()
    print '\n'
Ejemplo n.º 3
0
def __save_data(stage, data, save_path):
    im_idx_list = data['images']
    gt_boxes_list = data['bboxes']
    num_of_images = len(im_idx_list)
    # save files
    saveFolder = os.path.join(rootPath, "tmp/data/%s/" % (stage))
    print(">>>>>> Gen hard samples for %s..." % (stage))
    typeName = ["pos", "neg", "part"]
    saveFiles = {}
    for tp in typeName:
        _saveFolder = os.path.join(saveFolder, tp)
        if not os.path.isdir(_saveFolder):
            os.makedirs(_saveFolder)
        saveFiles[tp] = open(os.path.join(saveFolder, "{}.txt".format(tp)),
                             'w')
    #read detect result
    det_boxes = pickle.load(
        open(os.path.join(save_path, 'detections.pkl'), 'rb'))
    assert len(
        det_boxes) == num_of_images, "incorrect detections or ground truths"
    # index of neg, pos and part face, used as their image names
    n_idx, p_idx, d_idx = 0, 0, 0
    total_idx = 0
    for im_idx, dets, gts in zip(im_idx_list, det_boxes, gt_boxes_list):
        gts = np.array(gts, dtype=np.float32).reshape(-1, 4)
        if dets.shape[0] == 0:
            continue
        img = cv2.imread(im_idx)
        total_idx += 1
        #change to square
        dets = convert_to_square(dets)
        dets[:, 0:4] = np.round(dets[:, 0:4])
        neg_num = 0
        for box in dets:
            x_left, y_top, x_right, y_bottom, _ = box.astype(int)
            width = x_right - x_left + 1
            height = y_bottom - y_top + 1
            # ignore box that is too small or beyond image border
            if width < 20 or x_left < 0 or y_top < 0 or x_right > img.shape[
                    1] - 1 or y_bottom > img.shape[0] - 1:
                continue
            # compute intersection over union(IoU) between current box and all gt boxes
            Iou = IoU(box, gts)
            cropped_im = img[y_top:y_bottom + 1, x_left:x_right + 1, :]
            image_size = 24 if stage == "rnet" else 48
            resized_im = cv2.resize(cropped_im, (image_size, image_size),
                                    interpolation=cv2.INTER_LINEAR)
            # save negative images and write label
            # Iou with all gts must below 0.3
            if np.max(Iou) < 0.3 and neg_num < 60:
                # now to save it
                save_file = os.path.join(saveFolder, "neg", "%s.jpg" % n_idx)
                saveFiles['neg'].write(save_file + ' 0\n')
                cv2.imwrite(save_file, resized_im)
                n_idx += 1
                neg_num += 1
            else:
                # find gt_box with the highest iou
                idx = np.argmax(Iou)
                assigned_gt = gts[idx]
                x1, y1, x2, y2 = assigned_gt
                # compute bbox reg label
                offset_x1 = (x1 - x_left) / float(width)
                offset_y1 = (y1 - y_top) / float(height)
                offset_x2 = (x2 - x_right) / float(width)
                offset_y2 = (y2 - y_bottom) / float(height)
                # save positive and part-face images and write labels
                if np.max(Iou) >= 0.65:
                    save_file = os.path.join(saveFolder, "pos",
                                             "%s.jpg" % p_idx)
                    saveFiles['pos'].write(
                        save_file + ' 1 %.2f %.2f %.2f %.2f\n' %
                        (offset_x1, offset_y1, offset_x2, offset_y2))
                    cv2.imwrite(save_file, resized_im)
                    p_idx += 1
                elif np.max(Iou) >= 0.4:
                    save_file = os.path.join(saveFolder, "part",
                                             "%s.jpg" % d_idx)
                    saveFiles['part'].write(
                        save_file + ' -1 %.2f %.2f %.2f %.2f\n' %
                        (offset_x1, offset_y1, offset_x2, offset_y2))
                    cv2.imwrite(save_file, resized_im)
                    d_idx += 1
        printStr = "\r[{}] pos: {}  neg: {}  part:{}".format(
            total_idx, p_idx, n_idx, d_idx)
        sys.stdout.write(printStr)
        sys.stdout.flush()
    for f in saveFiles.values():
        f.close()
    print('\n')