diff options
Diffstat (limited to 'cli/app/search/search_class.py')
| -rw-r--r-- | cli/app/search/search_class.py | 11 |
1 files changed, 2 insertions, 9 deletions
diff --git a/cli/app/search/search_class.py b/cli/app/search/search_class.py index eba61e8..7eab4bc 100644 --- a/cli/app/search/search_class.py +++ b/cli/app/search/search_class.py @@ -24,23 +24,16 @@ from app.search.image import image_to_uint8, imconvert_uint8, imconvert_float32, from app.search.vector import truncated_z_sample, truncated_z_single, \ create_labels, create_labels_uniform -def find_nearest_vector_for_images(opt_fp_in, opt_dims, opt_steps, opt_limit, opt_video, opt_tag): +def find_nearest_vector_for_images(paths, opt_dims, opt_steps, opt_video, opt_tag, opt_limit=-1): sess = tf.compat.v1.Session() tf.reset_default_graph() generator = hub.Module('https://tfhub.dev/deepmind/biggan-512/2') - if os.path.isdir(opt_fp_in): - paths = glob(os.path.join(opt_fp_in, '*.jpg')) + \ - glob(os.path.join(opt_fp_in, '*.jpeg')) + \ - glob(os.path.join(opt_fp_in, '*.png')) - else: - paths = [opt_fp_in] - fp_inverses = os.path.join(app_cfg.DIR_INVERSES, opt_tag) os.makedirs(fp_inverses, exist_ok=True) save_params_latent(fp_inverses, opt_tag) save_params_dense(fp_inverses, opt_tag) - out_file = h5py.File(join(fp_inverses, 'dataset.hdf5'), 'w') + out_file = h5py.File(join(fp_inverses, 'dataset.latent.hdf5'), 'w') out_images = out_file.create_dataset('xtrain', (len(paths), 3, 512, 512,), dtype='float32') out_labels = out_file.create_dataset('ytrain', (len(paths), 1000,), dtype='float32') out_latent = out_file.create_dataset('latent', (len(paths), 128,), dtype='float32') |