You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
67 lines
2.2 KiB
67 lines
2.2 KiB
4 years ago
|
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
|