| | import os.path as op |
| | from zipfile import ZipFile, BadZipFile |
| | import torch.utils.data as data |
| | from PIL import Image |
| | from io import BytesIO |
| | import multiprocessing |
| |
|
| | _VALID_IMAGE_TYPES = ['.jpg', '.jpeg', '.tiff', '.bmp', '.png'] |
| |
|
| |
|
| | class ZipData(data.Dataset): |
| | _IGNORE_ATTRS = {'_zip_file'} |
| |
|
| | def __init__(self, path, map_file, |
| | transform=None, target_transform=None, |
| | extensions=None): |
| | self._path = path |
| | if not extensions: |
| | extensions = _VALID_IMAGE_TYPES |
| | self._zip_file = ZipFile(path) |
| | self.zip_dict = {} |
| | self.samples = [] |
| | self.transform = transform |
| | self.target_transform = target_transform |
| | self.class_to_idx = {} |
| | with open(map_file, 'r') as f: |
| | for line in iter(f.readline, ""): |
| | line = line.strip() |
| | if not line: |
| | continue |
| | cls_idx = [l for l in line.split('\t') if l] |
| | if not cls_idx: |
| | continue |
| | if (len(cls_idx) < 2): |
| | cls_idx = [l for l in line.split(' ') if l] |
| | if not cls_idx: |
| | continue |
| | assert len(cls_idx) >= 2, "invalid line: {}".format(line) |
| | idx = int(cls_idx[1]) |
| | cls = cls_idx[0] |
| | del cls_idx |
| | at_idx = cls.find('@') |
| | assert at_idx >= 0, "invalid class: {}".format(cls) |
| | cls = cls[at_idx + 1:] |
| | if cls.startswith('/'): |
| | |
| | cls = cls[1:] |
| | assert cls, "invalid class in line {}".format(line) |
| | prev_idx = self.class_to_idx.get(cls) |
| | assert prev_idx is None or prev_idx == idx, "class: {} idx: {} previously had idx: {}".format( |
| | cls, idx, prev_idx |
| | ) |
| | self.class_to_idx[cls] = idx |
| |
|
| | for fst in self._zip_file.infolist(): |
| | fname = fst.filename |
| | target = self.class_to_idx.get(fname) |
| | if target is None: |
| | continue |
| | if fname.endswith('/') or fname.startswith('.') or fst.file_size == 0: |
| | continue |
| | ext = op.splitext(fname)[1].lower() |
| | if ext in extensions: |
| | self.samples.append((fname, target)) |
| | assert len(self), "No images found in: {} with map: {}".format(self._path, map_file) |
| |
|
| | def __repr__(self): |
| | return 'ZipData({}, size={})'.format(self._path, len(self)) |
| |
|
| | def __getstate__(self): |
| | return { |
| | key: val if key not in self._IGNORE_ATTRS else None |
| | for key, val in self.__dict__.iteritems() |
| | } |
| |
|
| | def __getitem__(self, index): |
| | proc = multiprocessing.current_process() |
| | pid = proc.pid |
| | if pid not in self.zip_dict: |
| | self.zip_dict[pid] = ZipFile(self._path) |
| | zip_file = self.zip_dict[pid] |
| |
|
| | if index >= len(self) or index < 0: |
| | raise KeyError("{} is invalid".format(index)) |
| | path, target = self.samples[index] |
| | try: |
| | sample = Image.open(BytesIO(zip_file.read(path))).convert('RGB') |
| | except BadZipFile: |
| | print("bad zip file") |
| | return None, None |
| | if self.transform is not None: |
| | sample = self.transform(sample) |
| | if self.target_transform is not None: |
| | target = self.target_transform(target) |
| | return sample, target |
| |
|
| | def __len__(self): |
| | return len(self.samples) |
| |
|
