mrq/DL-Art-School
1
3
Fork 5
Code Issues Pull Requests Packages Projects Releases Wiki Activity

classify_into_folders script

Browse Source
This commit is contained in:
James Betker 2021-10-27 14:56:16 -06:00
parent d91dcbd404
commit bb0a0c8264
1 changed files with 32 additions and 24 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

44
codes/scripts/classify_into_folders.py
View File

@ -1,5 +1,6 @@
import os.path as osp import os.path as osp
import logging import logging
import shutil
import time import time
import argparse import argparse
@ -20,7 +21,7 @@ if __name__ == "__main__":
torch.backends.cudnn.benchmark = True torch.backends.cudnn.benchmark = True
want_metrics = False want_metrics = False
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument('-opt', type=str, help='Path to options YAML file.', default='../options/train_imgset_structural_classifier.yml') parser.add_argument('-opt', type=str, help='Path to options YAML file.', default='../options/test_noisy_audio_clips_classifier.yml')
opt = option.parse(parser.parse_args().opt, is_train=False) opt = option.parse(parser.parse_args().opt, is_train=False)
opt = option.dict_to_nonedict(opt) opt = option.dict_to_nonedict(opt)
utils.util.loaded_options = opt utils.util.loaded_options = opt
@ -36,33 +37,40 @@ if __name__ == "__main__":
#### Create test dataset and dataloader #### Create test dataset and dataloader
test_loaders = [] test_loaders = []
for phase, dataset_opt in sorted(opt['datasets'].items()): for phase, dataset_opt in sorted(opt['datasets'].items()):
dataset_opt['dataset']['includes_labels'] = False if 'test' in phase:
del dataset_opt['dataset']['labeler']
test_set = create_dataset(dataset_opt) test_set = create_dataset(dataset_opt)
if hasattr(test_set, 'wrapped_dataset'): if hasattr(test_set, 'wrapped_dataset'):
test_set = test_set.wrapped_dataset test_set = test_set.wrapped_dataset
test_loader = create_dataloader(test_set, dataset_opt, opt) test_loader = create_dataloader(test_set, dataset_opt, opt)
logger.info('Number of test images: {:d}'.format(len(test_set))) logger.info('Number of test images: {:d}'.format(len(test_set)))
test_loaders.append(test_loader) test_loaders.append((dataset_opt['name'], test_loader))
model = ExtensibleTrainer(opt) model = ExtensibleTrainer(opt)
gen = model.netsG['generator'] # Remove all losses, since often labels will not be provided and this is not implied in test.
label_to_search_for = 4 for s in model.steps:
s.losses = {}
for test_loader in test_loaders: output_key = opt['eval']['classifier_logits_key']
test_set_name = test_loader.dataset.opt['name'] output_base_dir = opt['eval']['output_dir']
labels = opt['eval']['output_labels']
path_key = opt['eval']['path_key']
step = 0
for test_set_name, test_loader in test_loaders:
test_start_time = time.time() test_start_time = time.time()
dataset_dir = osp.join(opt['path']['results_root'], opt['name']) dataset_dir = osp.join(opt['path']['results_root'], opt['name'])
util.mkdir(dataset_dir) util.mkdir(dataset_dir)
tq = tqdm(test_loader) for data in tqdm(test_loader):
step = 1 with torch.no_grad():
for data in tq: model.feed_data(data, 0)
hq = data['hq'].to('cuda') model.test()
res = gen(hq)
res = torch.nn.functional.interpolate(res, size=hq.shape[2:], mode="nearest") lbls = torch.nn.functional.softmax(model.eval_state[output_key][0].cpu(), dim=-1)
res_lbl = res[:, label_to_search_for, :, :].unsqueeze(1) for k, lbl in enumerate(lbls):
res_lbl_mask = (1.0 * (res_lbl > .5))*.5 + .5 lbl = torch.argmax(lbl, dim=0)
hq = hq * res_lbl_mask src_path = data[path_key][k]
torchvision.utils.save_image(hq, os.path.join(dataset_dir, "%i.png" % (step,))) dest = os.path.join(output_base_dir, labels[lbl])
os.makedirs(dest, exist_ok=True)
shutil.copy(str(src_path), os.path.join(dest, f'{step}_{os.path.basename(str(src_path))}'))
step += 1 step += 1