def extract_wallybody(dirpath, anns_path):
    img_prc = ImageProcessing()
    anns = open(anns_path, 'r')

    paths = glob.glob('{}/*.jpg'.format(dirpath))
    names = os.listdir(dirpath)

    # Body 가 들어있는 이미지의 정보를 가져옵니다
    lines = anns.readlines()
    body_dict = {}
    for line in lines[1:]:
        fpath, x1, x2, y1, y2 = line.split(',')
        x1, x2, y1, y2 = map(lambda ele: int(ele.strip()), [x1, x2, y1, y2])
        name = utils.get_name(fpath)
        # first
        if not name in body_dict.keys():
            body_dict[name] = [(x1, x2, y1, y2)]
        else:
            body_dict[name].append((x1, x2, y1, y2))

    fg_imgs_list = []
    bg_imgs_list = []
    fgs = []
    # get wally face list
    for p, path in enumerate(paths):
        name = utils.get_name(path)
        img = np.asarray(Image.open(path).convert('RGB'))
        # extract wally
        if name in body_dict.keys():
            for i, coord in enumerate(body_dict[name]):
                x1, x2, y1, y2 = coord
                fg = img[y1:y2, x1:x2, :]

                fgs.append(fg)
                fg_imgs, fg_coords = img_prc.guarantee_stride_cropping(
                    img, (400, 400), [x1, y1, x2, y2], (25, 25))
                if len(fg_imgs) == 0:
                    print path, x2 - x1, y2 - y1
                else:
                    fg_imgs = img_prc.resize_npImages(fg_imgs, (80, 80))
                    fg_imgs_list.append(fg_imgs)

                img = copy.deepcopy(img)
                # fill rectangle for extract back grounds images
                cv2.rectangle(img, (x1, y1), (x2, y2), (0, 0, 0), -1)

        bg_imgs, bg_coords = img_prc.stride_cropping(img, 200, 200, 400, 400)
        bg_imgs = img_prc.resize_npImages(bg_imgs, (80, 80))
        bg_imgs_list.append(bg_imgs)

    fgs = np.vstack(fg_imgs_list)
    bgs = np.vstack(bg_imgs_list)

    return fgs, bgs
Beispiel #2
0
class WallyDataset_ver4(Wally_dataset):
    """
    Class Usage :

    1. self.infos 는
        {이미지이름.jpg :
        {'whole_img' : 이름 원본 이미지와(Numpy)} ,
        {'face_img'  , 원본 이미지 (Numpy)}
        {'coords': face coord }
        형태로 저장되어 있습니다.

    2. 월리 얼굴 형태를 포함해서 cropping 합니다
        def cropping_with_face:
        src_img : 원본 이미지입니다  , self.infos['whole_img] 에서 가져와 사용하면 됩니다.
        crop_size : 얼마만큼 crop 할지 결정합니다
        coord : 얼굴 좌표를 말합니다  ,self.infos['coords'] 에서 가져와 사용하면 됩니다.
        stride_size
    """
    def __init__(self, wholeImg_dir, faceImg_dir, anns):
        self.img_prc = ImageProcessing()
        self.wholeImg_dir = wholeImg_dir
        self.faceImg_dir = faceImg_dir
        self.anns = anns
        #
        self.face_paths = glob.glob(os.path.join(faceImg_dir, '*.jpg'))
        self.whole_paths = glob.glob(os.path.join(wholeImg_dir, '*.jpg'))
        self.names = get_names(self.face_paths)
        #
        f = open(self.anns, 'r')
        lines = f.readlines()
        # not include first line
        self.infos = {}
        #
        print '# line {}'.format(len(lines))

        for line in lines[1:][:]:
            fname, x1, x2, y1, y2, w, h = line.split(',')
            x1, x2, y1, y2, w, h = map(lambda x: int(x.strip()),
                                       [x1, x2, y1, y2, w, h])
            coord = [x1, y1, x2, y2, w, h]

            whole_img = self.img_prc.path2img(os.path.join(
                self.wholeImg_dir, fname),
                                              resize=None)
            face_img = self.img_prc.path2img(os.path.join(
                self.faceImg_dir, fname),
                                             resize=None)
            self.infos[fname] = {
                'coord': coord,
                'whole_img': whole_img,
                'face_img': face_img
            }

    def cropping_with_face(self, src_img, crop_size, coord, stride_size):
        # src_img = 원본 이미지
        # crop = (crop_h , crop_w)
        # coord = [x1,y1,x2,y2]
        # stride_size = stride_h , stride_w
        cropped_imgs, coords = self.img_prc.guarantee_stride_cropping(
            src_img, crop_size, coord, stride_size)
        return cropped_imgs

    def generate_tfrecord(self, tfrecord_path, n_fg, fg_imgs, n_bg, bg_imgs):
        self.img_prc.make_tfrecord(tfrecord_path, None, (n_fg, fg_imgs),
                                   (n_bg, bg_imgs))

    def get_wallyface(self):
        fg_train_savepath = os.path.join('Wally_ver3', 'numpy_imgs',
                                         'fg_train.npy')
        fg_test_savepath = os.path.join('Wally_ver3', 'numpy_imgs',
                                        'fg_test.npy')
        fg_val_savepath = os.path.join('Wally_ver3', 'numpy_imgs',
                                       'fg_val.npy')

        if os.path.exists(fg_train_savepath) and os.path.exists(
                fg_test_savepath) and os.path.exists(fg_val_savepath):
            self.fg_train_imgs = np.load(fg_train_savepath)
            self.fg_test_imgs = np.load(fg_test_savepath)
            self.fg_val_imgs = np.load(fg_val_savepath)
        else:
            print 'Generating WallyFace Data....'
            fg_list = []
            for key in wally_dp.infos.keys():
                target_coords = self.infos[key]['coord']
                x1, y1, x2, y2 = target_coords[:4]
                whole_img = self.infos[key]['whole_img']
                # Wally 얼굴을 포함하는걸 보장하며 crop 합니다
                cropped_imgs = self.cropping_with_face(whole_img, (64, 64),
                                                       [x1, y1, x2, y2],
                                                       (1, 1))
                # Save crop
                np.save('./Wally_ver3/cropped_img_with_face/{}'.format(
                    os.path.splitext(key)[0]),
                        arr=cropped_imgs)
                fg_imgs = np.load(
                    './Wally_ver3/cropped_img_with_face/{}.npy'.format(
                        os.path.splitext(key)[0]))
                fg_list.append(fg_imgs)
            fg_imgs = np.vstack(fg_list)
            print 'foreground shape : {}'.format(np.shape(fg_imgs))

            # divide Train , Val , Test
            self.fg_test_imgs = fg_imgs[:5000]
            self.fg_val_imgs = fg_imgs[5000:5000 * 2]
            self.fg_train_imgs = fg_imgs[5000 * 2:]

            # save imgs to numpy
            np.save(fg_train_savepath, self.fg_train_imgs)
            np.save(fg_test_savepath, self.fg_test_imgs)
            np.save(fg_val_savepath, self.fg_val_imgs)

    def get_train(self):
        imgs = np.vstack([self.fg_train_imgs, self.bg_train_imgs])
        labs = np.asarray([0] * len(self.fg_train_imgs) +
                          [1] * len(self.bg_train_imgs))
        trainImgs_savepath = os.path.join('Wally_ver3', 'numpy_imgs',
                                          'train_imgs.npy')
        trainLabs_savepath = os.path.join('Wally_ver3', 'numpy_imgs',
                                          'train_labs.npy')
        np.save(trainImgs_savepath, imgs)
        np.save(trainLabs_savepath, labs)