Spaces:
Running
Running
| # 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] + ')' | |