File size: 9,118 Bytes
8c212a5 |
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 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 |
# python3.7
"""Contains the class of dataset."""
import os
import pickle
import string
import zipfile
import numpy as np
import cv2
import lmdb
import torch
from torch.utils.data import Dataset
from .transforms import progressive_resize_image
from .transforms import crop_resize_image
from .transforms import resize_image
from .transforms import normalize_image
try:
import turbojpeg
BASE_DIR = os.path.dirname(os.path.relpath(__file__))
LIBRARY_NAME = 'libturbojpeg.so.0'
LIBRARY_PATH = os.path.join(BASE_DIR, LIBRARY_NAME)
jpeg = turbojpeg.TurboJPEG(LIBRARY_PATH)
except ImportError:
jpeg = None
__all__ = ['BaseDataset']
_FORMATS_ALLOWED = ['dir', 'lmdb', 'list', 'zip']
class ZipLoader(object):
"""Defines a class to load zip file.
This is a static class, which is used to solve the problem that different
data workers can not share the same memory.
"""
files = dict()
@staticmethod
def get_zipfile(file_path):
"""Fetches a zip file."""
zip_files = ZipLoader.files
if file_path not in zip_files:
zip_files[file_path] = zipfile.ZipFile(file_path, 'r')
return zip_files[file_path]
@staticmethod
def get_image(file_path, image_path):
"""Decodes an image from a particular zip file."""
zip_file = ZipLoader.get_zipfile(file_path)
image_str = zip_file.read(image_path)
image_np = np.frombuffer(image_str, np.uint8)
image = cv2.imdecode(image_np, cv2.IMREAD_COLOR)
return image
class LmdbLoader(object):
"""Defines a class to load lmdb file.
This is a static class, which is used to solve lmdb loading error
when num_workers > 0
"""
files = dict()
@staticmethod
def get_lmdbfile(file_path):
"""Fetches a lmdb file"""
lmdb_files = LmdbLoader.files
if 'env' not in lmdb_files:
env = lmdb.open(file_path,
max_readers=1,
readonly=True,
lock=False,
readahead=False,
meminit=False)
with env.begin(write=False) as txn:
num_samples = txn.stat()['entries']
cache_file = '_cache_' + ''.join(
c for c in file_path if c in string.ascii_letters)
if os.path.isfile(cache_file):
keys = pickle.load(open(cache_file, "rb"))
else:
with env.begin(write=False) as txn:
keys = [key for key, _ in txn.cursor()]
pickle.dump(keys, open(cache_file, "wb"))
lmdb_files['env'] = env
lmdb_files['num_samples'] = num_samples
lmdb_files['keys'] = keys
return lmdb_files
@staticmethod
def get_image(file_path, idx):
"""Decodes an image from a particular lmdb file"""
lmdb_files = LmdbLoader.get_lmdbfile(file_path)
env = lmdb_files['env']
keys = lmdb_files['keys']
with env.begin(write=False) as txn:
imagebuf = txn.get(keys[idx])
image_np = np.frombuffer(imagebuf, np.uint8)
image = cv2.imdecode(image_np, cv2.IMREAD_COLOR)
return image
class BaseDataset(Dataset):
"""Defines the base dataset class.
This class supports loading data from a full-of-image folder, a lmdb
database, or an image list. Images will be pre-processed based on the given
`transform` function before fed into the data loader.
NOTE: The loaded data will be returned as a directory, where there must be
a key `image`.
"""
def __init__(self,
root_dir,
resolution,
data_format='dir',
image_list_path=None,
mirror=0.0,
progressive_resize=True,
crop_resize_resolution=-1,
transform=normalize_image,
transform_kwargs=None,
**_unused_kwargs):
"""Initializes the dataset.
Args:
root_dir: Root directory containing the dataset.
resolution: The resolution of the returned image.
data_format: Format the dataset is stored. Supports `dir`, `lmdb`,
and `list`. (default: `dir`)
image_list_path: Path to the image list. This field is required if
`data_format` is `list`. (default: None)
mirror: The probability to do mirror augmentation. (default: 0.0)
progressive_resize: Whether to resize images progressively.
(default: True)
crop_resize_resolution: The resolution of the output after crop
and resize. (default: -1)
transform: The transform function for pre-processing.
(default: `datasets.transforms.normalize_image()`)
transform_kwargs: The additional arguments for the `transform`
function. (default: None)
Raises:
ValueError: If the input `data_format` is not supported.
NotImplementedError: If the input `data_format` is not implemented.
"""
if data_format.lower() not in _FORMATS_ALLOWED:
raise ValueError(f'Invalid data format `{data_format}`!\n'
f'Supported formats: {_FORMATS_ALLOWED}.')
self.root_dir = root_dir
self.resolution = resolution
self.data_format = data_format.lower()
self.image_list_path = image_list_path
self.mirror = np.clip(mirror, 0.0, 1.0)
self.progressive_resize = progressive_resize
self.crop_resize_resolution = crop_resize_resolution
self.transform = transform
self.transform_kwargs = transform_kwargs or dict()
if self.data_format == 'dir':
self.image_paths = sorted(os.listdir(self.root_dir))
self.num_samples = len(self.image_paths)
elif self.data_format == 'lmdb':
lmdb_file = LmdbLoader.get_lmdbfile(self.root_dir)
self.num_samples = lmdb_file['num_samples']
elif self.data_format == 'list':
self.metas = []
assert os.path.isfile(self.image_list_path)
with open(self.image_list_path) as f:
for line in f:
fields = line.rstrip().split(' ')
if len(fields) == 1:
self.metas.append((fields[0], None))
else:
assert len(fields) == 2
self.metas.append((fields[0], int(fields[1])))
self.num_samples = len(self.metas)
elif self.data_format == 'zip':
zip_file = ZipLoader.get_zipfile(self.root_dir)
image_paths = [f for f in zip_file.namelist()
if ('.jpg' in f or '.jpeg' in f or '.png' in f)]
self.image_paths = sorted(image_paths)
self.num_samples = len(self.image_paths)
else:
raise NotImplementedError(f'Not implemented data format '
f'`{self.data_format}`!')
def __len__(self):
return self.num_samples
def __getitem__(self, idx):
data = dict()
# Load data.
if self.data_format == 'dir':
image_path = self.image_paths[idx]
try:
in_file = open(os.path.join(self.root_dir, image_path), 'rb')
image = jpeg.decode(in_file.read())
except: # pylint: disable=bare-except
image = cv2.imread(os.path.join(self.root_dir, image_path))
elif self.data_format == 'lmdb':
image = LmdbLoader.get_image(self.root_dir, idx)
elif self.data_format == 'list':
image_path, label = self.metas[idx]
image = cv2.imread(os.path.join(self.root_dir, image_path))
label = None if label is None else torch.LongTensor(label)
# data.update({'label': label})
elif self.data_format == 'zip':
image_path = self.image_paths[idx]
image = ZipLoader.get_image(self.root_dir, image_path)
else:
raise NotImplementedError(f'Not implemented data format '
f'`{self.data_format}`!')
image = image[:, :, ::-1] # Converts BGR (cv2) to RGB.
# Transform image.
if self.crop_resize_resolution > 0:
image = crop_resize_image(image, self.crop_resize_resolution)
if self.progressive_resize:
image = progressive_resize_image(image, self.resolution)
image = image.transpose(2, 0, 1).astype(np.float32)
if np.random.uniform() < self.mirror:
image = image[:, :, ::-1] # CHW
image = torch.FloatTensor(image.copy())
if not self.progressive_resize:
image = resize_image(image, self.resolution)
if self.transform is not None:
image = self.transform(image, **self.transform_kwargs)
data.update({'image': image})
return data
|