File size: 1,977 Bytes
404d2af |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 |
# Copyright 2019-present NAVER Corp.
# CC BY-NC-SA 3.0
# Available only for non-commercial use
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] + ')'
|