|
|
|
|
|
|
|
|
|
import os |
|
import json |
|
import pdb |
|
import numpy as np |
|
|
|
|
|
class Dataset(object): |
|
''' Base class for a dataset. To be overloaded. |
|
''' |
|
root = '' |
|
img_dir = '' |
|
nimg = 0 |
|
|
|
def __len__(self): |
|
return self.nimg |
|
|
|
def get_key(self, img_idx): |
|
raise NotImplementedError() |
|
|
|
def get_filename(self, img_idx, root=None): |
|
return os.path.join(root or self.root, self.img_dir, self.get_key(img_idx)) |
|
|
|
def get_image(self, img_idx): |
|
from PIL import Image |
|
fname = self.get_filename(img_idx) |
|
try: |
|
return Image.open(fname).convert('RGB') |
|
except Exception as e: |
|
raise IOError("Could not load image %s (reason: %s)" % (fname, str(e))) |
|
|
|
def __repr__(self): |
|
res = 'Dataset: %s\n' % self.__class__.__name__ |
|
res += ' %d images' % self.nimg |
|
res += '\n root: %s...\n' % self.root |
|
return res |
|
|
|
|
|
|
|
class CatDataset (Dataset): |
|
''' Concatenation of several datasets. |
|
''' |
|
def __init__(self, *datasets): |
|
assert len(datasets) >= 1 |
|
self.datasets = datasets |
|
offsets = [0] |
|
for db in datasets: |
|
offsets.append(db.nimg) |
|
self.offsets = np.cumsum(offsets) |
|
self.nimg = self.offsets[-1] |
|
self.root = None |
|
|
|
def which(self, i): |
|
pos = np.searchsorted(self.offsets, i, side='right')-1 |
|
assert pos < self.nimg, 'Bad image index %d >= %d' % (i, self.nimg) |
|
return pos, i - self.offsets[pos] |
|
|
|
def get_key(self, i): |
|
b, i = self.which(i) |
|
return self.datasets[b].get_key(i) |
|
|
|
def get_filename(self, i): |
|
b, i = self.which(i) |
|
return self.datasets[b].get_filename(i) |
|
|
|
def __repr__(self): |
|
fmt_str = "CatDataset(" |
|
for db in self.datasets: |
|
fmt_str += str(db).replace("\n"," ") + ', ' |
|
return fmt_str[:-2] + ')' |
|
|
|
|
|
|
|
|
|
|
