97 lines
4.0 KiB
Python
97 lines
4.0 KiB
Python
import os
|
|
from shutil import copyfile
|
|
|
|
def pytorch_prepare(data_dir, dataset_name):
|
|
dataset_dir = os.path.join(data_dir, dataset_name)
|
|
|
|
if not os.path.isdir(dataset_dir):
|
|
print('please change the download_path')
|
|
|
|
pytorch_path = os.path.join(dataset_dir , 'pytorch')
|
|
|
|
if not os.path.isdir(pytorch_path):
|
|
os.mkdir(pytorch_path)
|
|
#-----------------------------------------
|
|
#query
|
|
print('generatring ' + dataset_name + ' query images.')
|
|
query_dir = os.path.join(dataset_dir , 'query')
|
|
query_save_dir = os.path.join(dataset_dir , 'pytorch', 'query')
|
|
if not os.path.isdir(query_save_dir):
|
|
os.mkdir(query_save_dir)
|
|
|
|
for root, dirs, files in os.walk(query_dir, topdown=True):
|
|
for name in files:
|
|
if not name[-3:]=='jpg':
|
|
continue
|
|
ID = name.split('_')
|
|
src_dir = os.path.join(query_dir , name)
|
|
dst_dir = os.path.join(query_save_dir, ID[0])
|
|
if not os.path.isdir(dst_dir):
|
|
os.mkdir(dst_dir)
|
|
copyfile(src_dir, os.path.join(dst_dir , name))
|
|
#-----------------------------------------
|
|
#gallery
|
|
print('generatring '+dataset_name+' gallery images.')
|
|
gallery_dir = os.path.join(dataset_dir , 'bounding_box_test')
|
|
gallery_save_dir = os.path.join(dataset_dir , 'pytorch' , 'gallery')
|
|
if not os.path.isdir(gallery_save_dir):
|
|
os.mkdir(gallery_save_dir)
|
|
|
|
for root, dirs, files in os.walk(gallery_dir, topdown=True):
|
|
for name in files:
|
|
if not name[-3:]=='jpg':
|
|
continue
|
|
ID = name.split('_')
|
|
src_dir = os.path.join(gallery_dir, name)
|
|
dst_dir = os.path.join(gallery_save_dir, ID[0])
|
|
if not os.path.isdir(dst_dir):
|
|
os.mkdir(dst_dir)
|
|
copyfile(src_dir, os.path.join(dst_dir,name))
|
|
#---------------------------------------
|
|
#train_all
|
|
print('generatring '+dataset_name + ' all training images.')
|
|
train_dir = os.path.join( dataset_dir , 'bounding_box_train')
|
|
train_save_all_dir = os.path.join( dataset_dir , 'pytorch', 'train_all')
|
|
if not os.path.isdir(train_save_all_dir):
|
|
os.mkdir(train_save_all_dir)
|
|
|
|
for root, dirs, files in os.walk(train_dir, topdown=True):
|
|
for name in files:
|
|
if not name[-3:]=='jpg':
|
|
continue
|
|
ID = name.split('_')
|
|
src_dir = os.path.join(train_dir , name)
|
|
dst_dir = os.path.join(train_save_all_dir, ID[0])
|
|
if not os.path.isdir(dst_dir):
|
|
os.mkdir(dst_dir)
|
|
copyfile(src_dir, os.path.join(dst_dir, name))
|
|
|
|
#---------------------------------------
|
|
#train_val
|
|
print('generatring '+ dataset_name+' training and validation images.')
|
|
train_save_dir = os.path.join(dataset_dir, 'pytorch', 'train')
|
|
val_save_dir = os.path.join(dataset_dir , 'pytorch' , 'val')
|
|
if not os.path.isdir(train_save_dir):
|
|
os.mkdir(train_save_dir)
|
|
os.mkdir(val_save_dir)
|
|
|
|
for root, dirs, files in os.walk(train_dir, topdown=True):
|
|
for name in files:
|
|
if not name[-3:]=='jpg':
|
|
continue
|
|
ID = name.split('_')
|
|
src_dir = os.path.join(train_dir , name)
|
|
dst_dir = os.path.join(train_save_dir , ID[0])
|
|
if not os.path.isdir(dst_dir):
|
|
os.mkdir(dst_dir)
|
|
dst_dir = os.path.join(val_save_dir, ID[0]) #first image is used as val image
|
|
os.mkdir(dst_dir)
|
|
copyfile(src_dir, os.path.join(dst_dir , name))
|
|
print('Finished ' + dataset_name)
|
|
else:
|
|
print(dataset_name + ' pytorch directory exists!')
|
|
|
|
def pytorch_prepare_all(data_dir):
|
|
pytorch_prepare('Market1501', data_dir)
|
|
pytorch_prepare('DukeMTMC', data_dir)
|