Add enhanced ParserImageInTar that can read images from tars within tars, folders with multiple tars, etc. Additional comment cleanup.
Ross Wightman
4 years ago
parent
55f7dfa9ea
commit
5d4c3d0af3
@ -1,107 +0,0 @@
|
||||
import os
|
||||
import tarfile
|
||||
import pickle
|
||||
from glob import glob
|
||||
import numpy as np
|
||||
|
||||
from timm.utils.misc import natural_key
|
||||
|
||||
from .parser import Parser
|
||||
from .class_map import load_class_map
|
||||
from .constants import IMG_EXTENSIONS
|
||||
|
||||
|
||||
def extract_tarinfos(root, class_name_to_idx=None, cache_filename=None, extensions=None):
|
||||
tar_filenames = glob(os.path.join(root, '*.tar'), recursive=True)
|
||||
assert len(tar_filenames)
|
||||
num_tars = len(tar_filenames)
|
||||
|
||||
cache_path = ''
|
||||
if cache_filename is not None:
|
||||
cache_path = os.path.join(root, cache_filename)
|
||||
if os.path.exists(cache_path):
|
||||
with open(cache_path, 'rb') as pf:
|
||||
tarinfo_map = pickle.load(pf)
|
||||
else:
|
||||
tarinfo_map = {}
|
||||
for fi, fn in enumerate(tar_filenames):
|
||||
if fi % 1000 == 0:
|
||||
print(f'DEBUG: tar {fi}/{num_tars}')
|
||||
# cannot keep this open across processes, reopen later
|
||||
name = os.path.splitext(os.path.basename(fn))[0]
|
||||
with tarfile.open(fn) as tf:
|
||||
if extensions is None:
|
||||
# assume all files are valid samples
|
||||
class_tarinfos = tf.getmembers()
|
||||
else:
|
||||
class_tarinfos = [m for m in tf.getmembers() if os.path.splitext(m.name)[1].lower() in extensions]
|
||||
tarinfo_map[name] = dict(tarinfos=class_tarinfos)
|
||||
print(f'DEBUG: {len(class_tarinfos)} images for class {name}')
|
||||
tarinfo_map = {k: v for k, v in sorted(tarinfo_map.items(), key=lambda k: natural_key(k[0]))}
|
||||
if cache_path:
|
||||
with open(cache_path, 'wb') as pf:
|
||||
pickle.dump(tarinfo_map, pf, protocol=pickle.HIGHEST_PROTOCOL)
|
||||
|
||||
tarinfos = []
|
||||
targets = []
|
||||
build_class_map = False
|
||||
if class_name_to_idx is None:
|
||||
class_name_to_idx = {}
|
||||
build_class_map = True
|
||||
for i, (name, metadata) in enumerate(tarinfo_map.items()):
|
||||
class_idx = i
|
||||
if build_class_map:
|
||||
class_name_to_idx[name] = i
|
||||
else:
|
||||
if name not in class_name_to_idx:
|
||||
# only samples with class in class mapping are added
|
||||
continue
|
||||
class_idx = class_name_to_idx[name]
|
||||
num_samples = len(metadata['tarinfos'])
|
||||
tarinfos.extend(metadata['tarinfos'])
|
||||
targets.extend([class_idx] * num_samples)
|
||||
|
||||
return tarinfos, np.array(targets), class_name_to_idx
|
||||
|
||||
|
||||
class ParserImageClassInTar(Parser):
|
||||
""" Multi-tarfile dataset parser where there is one .tar file per class
|
||||
"""
|
||||
|
||||
CACHE_FILENAME = '_tarinfos.pickle'
|
||||
|
||||
def __init__(self, root, class_map=''):
|
||||
super().__init__()
|
||||
|
||||
class_name_to_idx = None
|
||||
if class_map:
|
||||
class_name_to_idx = load_class_map(class_map, root)
|
||||
assert os.path.isdir(root)
|
||||
self.root = root
|
||||
self.tarinfos, self.targets, self.class_name_to_idx = extract_tarinfos(
|
||||
self.root, class_name_to_idx=class_name_to_idx,
|
||||
cache_filename=self.CACHE_FILENAME, extensions=IMG_EXTENSIONS)
|
||||
self.class_idx_to_name = {v: k for k, v in self.class_name_to_idx.items()}
|
||||
self.tarfiles = {} # to open lazily
|
||||
self.cache_tarfiles = False
|
||||
|
||||
def __len__(self):
|
||||
return len(self.tarinfos)
|
||||
|
||||
def __getitem__(self, index):
|
||||
tarinfo = self.tarinfos[index]
|
||||
target = self.targets[index]
|
||||
class_name = self.class_idx_to_name[target]
|
||||
if self.cache_tarfiles:
|
||||
tf = self.tarfiles.setdefault(
|
||||
class_name, tarfile.open(os.path.join(self.root, class_name + '.tar')))
|
||||
else:
|
||||
tf = tarfile.open(os.path.join(self.root, class_name + '.tar'))
|
||||
fileobj = tf.extractfile(tarinfo)
|
||||
return fileobj, target
|
||||
|
||||
def _filename(self, index, basename=False, absolute=False):
|
||||
filename = self.tarinfos[index].name
|
||||
if basename:
|
||||
filename = os.path.basename(filename)
|
||||
return filename
|
||||
@ -0,0 +1,219 @@
|
||||
""" A dataset parser that reads tarfile based datasets
|
||||
|
||||
This parser can read and extract image samples from:
|
||||
* a single tar of image files
|
||||
* a folder of multiple tarfiles containing imagefiles
|
||||
* a tar of tars containing image files
|
||||
|
||||
Labels are based on the combined folder and/or tar name structure.
|
||||
|
||||
Hacked together by / Copyright 2020 Ross Wightman
|
||||
"""
|
||||
import os
|
||||
import tarfile
|
||||
import pickle
|
||||
import logging
|
||||
import numpy as np
|
||||
from glob import glob
|
||||
from typing import List, Dict
|
||||
|
||||
from timm.utils.misc import natural_key
|
||||
|
||||
from .parser import Parser
|
||||
from .class_map import load_class_map
|
||||
from .constants import IMG_EXTENSIONS
|
||||
|
||||
|
||||
_logger = logging.getLogger(__name__)
|
||||
CACHE_FILENAME_SUFFIX = '_tarinfos.pickle'
|
||||
|
||||
|
||||
class TarState:
|
||||
|
||||
def __init__(self, tf: tarfile.TarFile = None, ti: tarfile.TarInfo = None):
|
||||
self.tf: tarfile.TarFile = tf
|
||||
self.ti: tarfile.TarInfo = ti
|
||||
self.children: Dict[str, TarState] = {} # child states (tars within tars)
|
||||
|
||||
def reset(self):
|
||||
self.tf = None
|
||||
|
||||
|
||||
def _extract_tarinfo(tf: tarfile.TarFile, parent_info: Dict, extensions=IMG_EXTENSIONS):
|
||||
sample_count = 0
|
||||
for i, ti in enumerate(tf):
|
||||
if not ti.isfile():
|
||||
continue
|
||||
dirname, basename = os.path.split(ti.path)
|
||||
name, ext = os.path.splitext(basename)
|
||||
ext = ext.lower()
|
||||
if ext == '.tar':
|
||||
with tarfile.open(fileobj=tf.extractfile(ti), mode='r|') as ctf:
|
||||
child_info = dict(
|
||||
name=ti.name, path=os.path.join(parent_info['path'], name), ti=ti, children=[], samples=[])
|
||||
sample_count += _extract_tarinfo(ctf, child_info, extensions=extensions)
|
||||
_logger.debug(f'{i}/?. Extracted child tarinfos from {ti.name}. {len(child_info["samples"])} images.')
|
||||
parent_info['children'].append(child_info)
|
||||
elif ext in extensions:
|
||||
parent_info['samples'].append(ti)
|
||||
sample_count += 1
|
||||
return sample_count
|
||||
|
||||
|
||||
def extract_tarinfos(root, class_name_to_idx=None, cache_tarinfo=None, extensions=IMG_EXTENSIONS, sort=True):
|
||||
root_is_tar = False
|
||||
if os.path.isfile(root):
|
||||
assert os.path.splitext(root)[-1].lower() == '.tar'
|
||||
tar_filenames = [root]
|
||||
root, root_name = os.path.split(root)
|
||||
root_name = os.path.splitext(root_name)[0]
|
||||
root_is_tar = True
|
||||
else:
|
||||
root_name = root.strip(os.path.sep).split(os.path.sep)[-1]
|
||||
tar_filenames = glob(os.path.join(root, '*.tar'), recursive=True)
|
||||
num_tars = len(tar_filenames)
|
||||
tar_bytes = sum([os.path.getsize(f) for f in tar_filenames])
|
||||
assert num_tars, f'No .tar files found at specified path ({root}).'
|
||||
|
||||
_logger.info(f'Scanning {tar_bytes/1024**2:.2f}MB of tar files...')
|
||||
info = dict(tartrees=[])
|
||||
cache_path = ''
|
||||
if cache_tarinfo is None:
|
||||
cache_tarinfo = True if tar_bytes > 10*1024**3 else False # FIXME magic number, 10GB
|
||||
if cache_tarinfo:
|
||||
cache_filename = '_' + root_name + CACHE_FILENAME_SUFFIX
|
||||
cache_path = os.path.join(root, cache_filename)
|
||||
if os.path.exists(cache_path):
|
||||
_logger.info(f'Reading tar info from cache file {cache_path}.')
|
||||
with open(cache_path, 'rb') as pf:
|
||||
info = pickle.load(pf)
|
||||
assert len(info['tartrees']) == num_tars, "Cached tartree len doesn't match number of tarfiles"
|
||||
else:
|
||||
for i, fn in enumerate(tar_filenames):
|
||||
path = '' if root_is_tar else os.path.splitext(os.path.basename(fn))[0]
|
||||
with tarfile.open(fn, mode='r|') as tf: # tarinfo scans done in streaming mode
|
||||
parent_info = dict(name=os.path.relpath(fn, root), path=path, ti=None, children=[], samples=[])
|
||||
num_samples = _extract_tarinfo(tf, parent_info, extensions=extensions)
|
||||
num_children = len(parent_info["children"])
|
||||
_logger.debug(
|
||||
f'{i}/{num_tars}. Extracted tarinfos from {fn}. {num_children} children, {num_samples} samples.')
|
||||
info['tartrees'].append(parent_info)
|
||||
if cache_path:
|
||||
_logger.info(f'Writing tar info to cache file {cache_path}.')
|
||||
with open(cache_path, 'wb') as pf:
|
||||
pickle.dump(info, pf)
|
||||
|
||||
samples = []
|
||||
labels = []
|
||||
build_class_map = False
|
||||
if class_name_to_idx is None:
|
||||
build_class_map = True
|
||||
|
||||
# Flatten tartree info into lists of samples and targets w/ targets based on label id via
|
||||
# class map arg or from unique paths.
|
||||
# NOTE: currently only flattening up to two-levels, filesystem .tars and then one level of sub-tar children
|
||||
# this covers my current use cases and keeps things a little easier to test for now.
|
||||
tarfiles = []
|
||||
|
||||
def _label_from_paths(*path, leaf_only=True):
|
||||
path = os.path.join(*path).strip(os.path.sep)
|
||||
return path.split(os.path.sep)[-1] if leaf_only else path.replace(os.path.sep, '_')
|
||||
|
||||
def _add_samples(info, fn):
|
||||
added = 0
|
||||
for s in info['samples']:
|
||||
label = _label_from_paths(info['path'], os.path.dirname(s.path))
|
||||
if not build_class_map and label not in class_name_to_idx:
|
||||
continue
|
||||
samples.append((s, fn, info['ti']))
|
||||
labels.append(label)
|
||||
added += 1
|
||||
return added
|
||||
|
||||
_logger.info(f'Collecting samples and building tar states.')
|
||||
for parent_info in info['tartrees']:
|
||||
# if tartree has children, we assume all samples are at the child level
|
||||
tar_name = None if root_is_tar else parent_info['name']
|
||||
tar_state = TarState()
|
||||
parent_added = 0
|
||||
for child_info in parent_info['children']:
|
||||
child_added = _add_samples(child_info, fn=tar_name)
|
||||
if child_added:
|
||||
tar_state.children[child_info['name']] = TarState(ti=child_info['ti'])
|
||||
parent_added += child_added
|
||||
parent_added += _add_samples(parent_info, fn=tar_name)
|
||||
if parent_added:
|
||||
tarfiles.append((tar_name, tar_state))
|
||||
del info
|
||||
|
||||
if build_class_map:
|
||||
# build class index
|
||||
sorted_labels = list(sorted(set(labels), key=natural_key))
|
||||
class_name_to_idx = {c: idx for idx, c in enumerate(sorted_labels)}
|
||||
|
||||
_logger.info(f'Mapping targets and sorting samples.')
|
||||
samples_and_targets = [(s, class_name_to_idx[l]) for s, l in zip(samples, labels) if l in class_name_to_idx]
|
||||
if sort:
|
||||
samples_and_targets = sorted(samples_and_targets, key=lambda k: natural_key(k[0][0].path))
|
||||
|
||||
_logger.info(f'Finished processing {len(samples_and_targets)} samples across {len(tarfiles)} tar files.')
|
||||
return samples_and_targets, class_name_to_idx, tarfiles
|
||||
|
||||
|
||||
class ParserImageInTar(Parser):
|
||||
""" Multi-tarfile dataset parser where there is one .tar file per class
|
||||
"""
|
||||
|
||||
def __init__(self, root, class_map='', cache_tarfiles=True, cache_tarinfo=None):
|
||||
super().__init__()
|
||||
|
||||
class_name_to_idx = None
|
||||
if class_map:
|
||||
class_name_to_idx = load_class_map(class_map, root)
|
||||
self.root = root
|
||||
self.samples_and_targets, self.class_name_to_idx, tarfiles = extract_tarinfos(
|
||||
self.root,
|
||||
class_name_to_idx=class_name_to_idx,
|
||||
cache_tarinfo=cache_tarinfo,
|
||||
extensions=IMG_EXTENSIONS)
|
||||
self.class_idx_to_name = {v: k for k, v in self.class_name_to_idx.items()}
|
||||
if len(tarfiles) == 1 and tarfiles[0][0] is None:
|
||||
self.root_is_tar = True
|
||||
self.tar_state = tarfiles[0][1]
|
||||
else:
|
||||
self.root_is_tar = False
|
||||
self.tar_state = dict(tarfiles)
|
||||
self.cache_tarfiles = cache_tarfiles
|
||||
|
||||
def __len__(self):
|
||||
return len(self.samples_and_targets)
|
||||
|
||||
def __getitem__(self, index):
|
||||
sample, target = self.samples_and_targets[index]
|
||||
sample_ti, parent_fn, child_ti = sample
|
||||
parent_abs = os.path.join(self.root, parent_fn) if parent_fn else self.root
|
||||
|
||||
tf = None
|
||||
cache_state = None
|
||||
if self.cache_tarfiles:
|
||||
cache_state = self.tar_state if self.root_is_tar else self.tar_state[parent_fn]
|
||||
tf = cache_state.tf
|
||||
if tf is None:
|
||||
tf = tarfile.open(parent_abs)
|
||||
if self.cache_tarfiles:
|
||||
cache_state.tf = tf
|
||||
if child_ti is not None:
|
||||
ctf = cache_state.children[child_ti.name].tf if self.cache_tarfiles else None
|
||||
if ctf is None:
|
||||
ctf = tarfile.open(fileobj=tf.extractfile(child_ti))
|
||||
if self.cache_tarfiles:
|
||||
cache_state.children[child_ti.name].tf = ctf
|
||||
tf = ctf
|
||||
|
||||
return tf.extractfile(sample_ti), target
|
||||
|
||||
def _filename(self, index, basename=False, absolute=False):
|
||||
filename = self.samples_and_targets[index][0][0].name
|
||||
if basename:
|
||||
filename = os.path.basename(filename)
|
||||
return filename
|
||||
Loading…
Reference in new issue