fetch list of images and dedupe

5 files changed, 59 insertions, 19 deletions

diff --git a/cli/app/commands/biggan/fetch.py b/cli/app/commands/biggan/fetch.py
index 94117c9..39f503c 100644
--- a/cli/app/commands/biggan/fetch.py
+++ b/cli/app/commands/biggan/fetch.py
@@ -1,6 +1,7 @@
 import click
+import os
 
-from app.utils.cortex_utils import fetch_cortex_folder
+from app.utils.cortex_utils import fetch_cortex_folder, find_unprocessed_files
 
 @click.command('')
 @click.option('-i', '--folder_id', 'opt_folder_id', type=int,
@@ -10,4 +11,8 @@ def cli(ctx, opt_folder_id):
   """
   Fetch JSON from the server
   """
-  fetch_cortex_folder(opt_folder_id)
+  files = fetch_cortex_folder(opt_folder_id)
+  unprocessed_files = find_unprocessed_files(files)
+  print("Unprocessed files:")
+  for file in unprocessed_files:
+    print(" - {}".format(file['name']))
diff --git a/cli/app/search/search_class.py b/cli/app/search/search_class.py
index eb9ff42..eba61e8 100644
--- a/cli/app/search/search_class.py
+++ b/cli/app/search/search_class.py
@@ -26,7 +26,7 @@ from app.search.vector import truncated_z_sample, truncated_z_single, \
 
 def find_nearest_vector_for_images(opt_fp_in, opt_dims, opt_steps, opt_limit, opt_video, opt_tag):
   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):
@@ -52,6 +52,7 @@ def find_nearest_vector_for_images(opt_fp_in, opt_dims, opt_steps, opt_limit, op
     fp_frames = find_nearest_vector(sess, generator, path, opt_dims, out_images, out_labels, out_latent, opt_steps, index)
     if opt_video:
       export_video(fp_frames)
+  sess.close()
 
 def find_nearest_vector(sess, generator, opt_fp_in, opt_dims, out_images, out_labels, out_latent, opt_steps, index):
   """
diff --git a/cli/app/search/search_dense.py b/cli/app/search/search_dense.py
index 0086db5..392fa70 100644
--- a/cli/app/search/search_dense.py
+++ b/cli/app/search/search_dense.py
@@ -100,6 +100,8 @@ def find_dense_embedding_for_images(params):
   # --------------------------
   # Load Graph.
   # --------------------------
+  tf.reset_default_graph()
+
   generator = hub.Module(str(params.generator_path))
 
   gen_signature = 'generator'
@@ -419,6 +421,9 @@ def find_dense_embedding_for_images(params):
 
   out_fns[:] = sample_fns[:NUM_IMGS_TO_PROCESS]
 
+  vector_dir = os.path.join(app_cfg.INVERSES_DIR, "vectors")
+  os.makedirs(vector_dir, exist_ok=True)
+
   # Gradient descent w.r.t. generator's inputs.
   it = 0
   out_pos = 0
@@ -485,19 +490,22 @@ def find_dense_embedding_for_images(params):
     # write encoding, latent to pkl file
     for i in range(BATCH_SIZE):
       out_i = out_pos + i
-      fn, ext = os.path.splitext(sample_fns[out_i])
-      fp_out_pkl = os.path.join(app_cfg.INVERSES_DIR, fn  ".pkl")
-      out_data = {
-        'id': fn,
-        'latent': out_lat[out_i],
-        'encoding': out_enc[out_i],
-        'label': out_labels[out_i],
-      }
-      write_pickle(out_data, fp_out_pkl)
+      sample_fn, ext = os.path.splitext(sample_fns[out_i])
       image = Image.fromarray(images[i])
       fp = BytesIO()
       image.save(fp, format='png')
-      upload_bytes_to_cortex(params.folder_id, fn, fp, 'image/png')
+      data = upload_bytes_to_cortex(params.folder_id, sample_fn + "-inverse.png", fp, "image/png")
+      if data is not None:
+        file_id = data['id']
+        fp_out_pkl = os.path.join(vector_dir, "file_{}.pkl".format(file_id))
+        out_data = {
+          'id': file_id,
+          'sample_fn': sample_fn,
+          'label': out_labels[out_i],
+          'latent': out_lat[out_i],
+          'encoding': out_enc[out_i],
+        }
+        write_pickle(out_data, fp_out_pkl)
 
     out_pos += BATCH_SIZE
     if params.max_batches > 0 and (out_pos / BATCH_SIZE) >= params.max_batches:
diff --git a/cli/app/settings/app_cfg.py b/cli/app/settings/app_cfg.py
index 7e739df..bade59c 100644
--- a/cli/app/settings/app_cfg.py
+++ b/cli/app/settings/app_cfg.py
@@ -21,7 +21,6 @@ CLICK_GROUPS = {
   'biggan': 'app/commands/biggan',
   'bigbigan': 'app/commands/bigbigan',
   'cortex': 'app/commands/cortex',
-  'process': 'app/commands/process',
 }
 
 # -----------------------------------------------------------------------------
diff --git a/cli/app/utils/cortex_utils.py b/cli/app/utils/cortex_utils.py
index 328b93a..8076fda 100644
--- a/cli/app/utils/cortex_utils.py
+++ b/cli/app/utils/cortex_utils.py
@@ -7,24 +7,49 @@ urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
 from app.settings import app_cfg
 
 def api_url(path):
+  """Generate a base API path"""
   return "https://lens.neural.garden/api/{}/".format(path)
 
 def fetch_cortex_folder(opt_folder_id):
+  """Fetch all new, non-generated files in a Cortex folder"""
   rows = fetch_json(api_url('file'), folder_id=opt_folder_id)
   fp_out_dir = join(app_cfg.DIR_INPUTS, "cortex", str(opt_folder_id))
   os.makedirs(fp_out_dir, exist_ok=True)
   for row in rows:
     if row['generated'] == 0 and row['processed'] != 1:
       fn, ext = os.path.splitext(row['name'])
-      fp_out_image = join(fp_out_dir, "{}{}".format(row['id'], ext))
+      fp_out_image = join(fp_out_dir, row['name'])
       if not os.path.exists(fp_out_image):
+        row['path'] = fp_out_image
         fetch_file(row['url'], fp_out_image)
+  return rows
+
+def find_unprocessed_files(files):
+  """Find files that haven't been processed yet.
+  This is implied if no matching generated file is found.
+  """
+  datasets = {}
+  unprocessed_files = []
+  for file in files:
+    if file['generated'] == 1:
+      fn, ext = os.path.splitext(file['name'])
+      dataset = fn.split('-')[0]
+      datasets[dataset] = file['id']
+  for file in files:
+    if file['generated'] == 0 and file['processed'] != 1:
+      fn, ext = os.path.splitext(file['name'])
+      dataset = fn.split('-')[0]
+      if dataset not in datasets:
+        unprocessed_files.append(file)
+  return unprocessed_files
 
 def fetch_json(url, **kwargs):
+  """HTTP GET some JSON"""
   resp = requests.get(url, params=kwargs, verify=False, timeout=10)
   return None if resp.status_code != 200 else resp.json()
 
 def fetch_file(url, fn, **kwargs):
+  """HTTP GET a binary file and write it to disk"""
   print("Fetch {} => {}".format(url, fn))
   try:
     resp = requests.get(url, params=kwargs, verify=False, timeout=10)
@@ -41,6 +66,7 @@ def fetch_file(url, fn, **kwargs):
   return size
 
 def upload_fp_to_cortex(opt_folder_id, fp):
+  """Upload a open file/BytesIO object"""
   files = {
     'file': fp
   }
@@ -52,13 +78,14 @@ def upload_fp_to_cortex(opt_folder_id, fp):
     'datatype': 'image',
   }
   url = os.path.join(api_url('folder'), opt_folder_id, 'upload/')
-  print(url)
   r = requests.post(url, files=files, data=data)
-  print(r.json())
+  return None if resp.status_code != 200 else resp.json()
 
 def upload_bytes_to_cortex(opt_folder_id, fn, fp, mimetype):
-  upload_fp_to_cortex(opt_folder_id, (fn, fp.getvalue(), mimetype,))
+  """Upload a BytesIO object"""
+  return upload_fp_to_cortex(opt_folder_id, (fn, fp.getvalue(), mimetype,))
 
 def upload_file_to_cortex(opt_folder_id, fn):
+  """Upload a file from disk"""
   with open(fn, 'rb') as fp:
-    upload_fp_to_cortex(opt_folder_id, fp)
+    return upload_fp_to_cortex(opt_folder_id, fp)