third_party/r2d2/datasets/web_images.py

# Copyright 2019-present NAVER Corp.
# CC BY-NC-SA 3.0
# Available only for non-commercial use

import os, pdb
from tqdm import trange

from .dataset import Dataset


class RandomWebImages(Dataset):
    """1 million distractors from Oxford and Paris Revisited
    see http://ptak.felk.cvut.cz/revisitop/revisitop1m/
    """

    def __init__(self, start=0, end=1024, root="data/revisitop1m"):
        Dataset.__init__(self)
        self.root = root

        bar = None
        self.imgs = []
        for i in range(start, end):
            try:
                # read cached list
                img_list_path = os.path.join(self.root, "image_list_%d.txt" % i)
                cached_imgs = [e.strip() for e in open(img_list_path)]
                assert cached_imgs, f"Cache '{img_list_path}' is empty!"
                self.imgs += cached_imgs

            except IOError:
                if bar is None:
                    bar = trange(start, 4 * end, desc="Caching")
                    bar.update(4 * i)

                # create it
                imgs = []
                for d in range(
                    i * 4, (i + 1) * 4
                ):  # 4096 folders in total, on average 256 each
                    key = hex(d)[2:].zfill(3)
                    folder = os.path.join(self.root, key)
                    if not os.path.isdir(folder):
                        continue
                    imgs += [f for f in os.listdir(folder) if verify_img(folder, f)]
                    bar.update(1)
                assert imgs, f"No images found in {folder}/"
                open(img_list_path, "w").write("\n".join(imgs))
                self.imgs += imgs

        if bar:
            bar.update(bar.total - bar.n)
        self.nimg = len(self.imgs)

    def get_key(self, i):
        key = self.imgs[i]
        return os.path.join(key[:3], key)


def verify_img(folder, f):
    path = os.path.join(folder, f)
    if not f.endswith(".jpg"):
        return False
    try:
        from PIL import Image

        Image.open(path).convert("RGB")  # try to open it
        return True
    except:
        return False