def extract_train(tar_fname, target_dir, with_rec=False, num_thread=1):
mkdir(target_dir)
with tarfile.open(tar_fname) as tar:
print("Extracting " + tar_fname + "...")
# extract each class one-by-one
pbar = tqdm(total=len(tar.getnames()))
for class_tar in tar:
pbar.set_description('Extract ' + class_tar.name)
tar.extract(class_tar, target_dir)
class_fname = os.path.join(target_dir, class_tar.name)
class_dir = os.path.splitext(class_fname)[0]
os.mkdir(class_dir)
with tarfile.open(class_fname) as f:
f.extractall(class_dir)
os.remove(class_fname)
pbar.update(1)
pbar.close()
if with_rec:
build_rec_process(target_dir, True, num_thread)
def extract_val(tar_fname, target_dir, with_rec=False, num_thread=1):
mkdir(target_dir)
print('Extracting ' + tar_fname)
with tarfile.open(tar_fname) as tar:
tar.extractall(target_dir)
# build rec file before images are moved into subfolders
if with_rec:
build_rec_process(target_dir, False, num_thread)
# move images to proper subfolders
val_maps_file = os.path.join(os.path.dirname(__file__),
'imagenet_val_maps.pklz')
download(
'https://gluon-cv.mxnet.io/_downloads/f5c3f5262b5968d15a687bf7bd73db68/imagenet_val_maps.pklz',
val_maps_file)
with gzip.open(val_maps_file, 'rb') as f:
dirs, mappings = pickle.load(f)
for d in dirs:
os.makedirs(os.path.join(target_dir, d))
for m in mappings:
os.rename(os.path.join(target_dir, m[0]),
os.path.join(target_dir, m[1], m[0]))
def build_rec_process(img_dir, train=False, num_thread=1):
rec_dir = os.path.abspath(os.path.join(img_dir, '../rec'))
mkdir(rec_dir)
prefix = 'train' if train else 'val'
print('Building ImageRecord file for ' + prefix + ' ...')
to_path = rec_dir
# download lst file and im2rec script
script_path = os.path.join(rec_dir, 'im2rec.py')
script_url = 'https://raw.githubusercontent.com/apache/incubator-mxnet/master/tools/im2rec.py'
download(script_url, script_path)
# execution
import sys
cmd = [
sys.executable, script_path, rec_dir, img_dir, '--recursive',
'--pass-through', '--pack-label', '--num-thread',
str(num_thread)
]
subprocess.call(cmd)
os.remove(script_path)
print('ImageRecord file for ' + prefix + ' has been built!')
def save(self, checkname=None):
checkname = checkname if checkname else self.checkname
mkdir(os.path.dirname(checkname))
save(self.state_dict(), checkname)