Spaces:
Running
Running
File size: 1,964 Bytes
a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 |
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 |
# 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] + ")"
|