diff options
Diffstat (limited to 'cli/app/search/search_dense.py')
| -rw-r--r-- | cli/app/search/search_dense.py | 8 |
1 files changed, 5 insertions, 3 deletions
diff --git a/cli/app/search/search_dense.py b/cli/app/search/search_dense.py index 6fba70d..066d946 100644 --- a/cli/app/search/search_dense.py +++ b/cli/app/search/search_dense.py @@ -327,20 +327,20 @@ def find_dense_embedding_for_images(params, opt_tag="inverse_" + timestamp(), op sample_labels = in_file['ytrain'][()] sample_fns = in_file['fn'][()] NUM_IMGS = sample_images.shape[0] # number of images to be inverted. + INFILL_IMGS = NUM_IMGS print("Number of images: {}".format(NUM_IMGS)) print("Batch size: {}".format(BATCH_SIZE)) def sample_images_gen(): - for i in range(int(NUM_IMGS / BATCH_SIZE)): + for i in range(int(INFILL_IMGS / BATCH_SIZE)): i_1, i_2 = i*BATCH_SIZE, (i+1)*BATCH_SIZE yield sample_images[i_1:i_2], sample_labels[i_1:i_2] image_gen = sample_images_gen() sample_latents = in_file['latent'] def sample_latent_gen(): - for i in range(int(NUM_IMGS / BATCH_SIZE)): + for i in range(int(INFILL_IMGS / BATCH_SIZE)): i_1, i_2 = i*BATCH_SIZE, (i+1)*BATCH_SIZE yield sample_latents[i_1:i_2] latent_gen = sample_latent_gen() - INFILL_IMGS = NUM_IMGS while INFILL_IMGS % BATCH_SIZE != 0: REMAINDER = 1 # BATCH_SIZE - (NUM_IMGS % BATCH_SIZE) INFILL_IMGS += REMAINDER @@ -444,7 +444,9 @@ def find_dense_embedding_for_images(params, opt_tag="inverse_" + timestamp(), op for i in range(BATCH_SIZE): out_i = out_pos + i if out_i >= NUM_IMGS: + print("{} >= {}, skipping...".format(out_i, NUM_IMGS)) continue + print("{}: {}".format(out_i, sample_fn)) sample_fn, ext = os.path.splitext(sample_fns[out_i]) image = Image.fromarray(images[i]) fp = BytesIO() |