import os import io import torch import tarfile from .parser import Parser from .class_map import load_class_map from .constants import IMG_EXTENSIONS from PIL import Image from timm.utils.misc import natural_key def extract_tar_info(tarfile, class_to_idx=None, sort=True): files = [] labels = [] for ti in tarfile.getmembers(): if not ti.isfile(): continue dirname, basename = os.path.split(ti.path) label = os.path.basename(dirname) ext = os.path.splitext(basename)[1] if ext.lower() in IMG_EXTENSIONS: files.append(ti) labels.append(label) if class_to_idx is None: unique_labels = set(labels) sorted_labels = list(sorted(unique_labels, key=natural_key)) class_to_idx = {c: idx for idx, c in enumerate(sorted_labels)} tarinfo_and_targets = [(f, class_to_idx[l]) for f, l in zip(files, labels) if l in class_to_idx] if sort: tarinfo_and_targets = sorted(tarinfo_and_targets, key=lambda k: natural_key(k[0].path)) return tarinfo_and_targets, class_to_idx class ParserImageTar(Parser): def __init__(self, root, load_bytes=False, class_map=''): super().__init__() class_to_idx = None if class_map: class_to_idx = load_class_map(class_map, root) assert os.path.isfile(root) self.root = root with tarfile.open(root) as tf: # cannot keep this open across processes, reopen later self.samples, self.class_to_idx = extract_tar_info(tf, class_to_idx) self.imgs = self.samples self.tarfile = None # lazy init in __getitem__ self.load_bytes = load_bytes def __getitem__(self, index): if self.tarfile is None: self.tarfile = tarfile.open(self.root) tarinfo, target = self.samples[index] iob = self.tarfile.extractfile(tarinfo) img = iob.read() if self.load_bytes else Image.open(iob).convert('RGB') return img, target def __len__(self): return len(self.samples) def _filename(self, index, basename=False, absolute=False): filename = self.samples[index][0].name if basename: filename = os.path.basename(filename) return filename