summaryrefslogtreecommitdiff
path: root/cli/app/commands/biggan/extract_dense_vectors.py
blob: 2436ce63ec52c9d44482fe1a8080a175334d5b7b (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
import click
import os

from app.utils.cortex_utils import fetch_cortex_folder, find_unprocessed_files
from app.search.search_class import find_nearest_vector_for_images
from app.search.search_dense import find_dense_embedding_for_images
from app.search.json import params_dense_dict

@click.command('')
@click.option('-f', '--folder_id', 'opt_folder_id', type=int,
  help='Folder ID to process')
@click.option('-ls', '--latent_steps', 'opt_latent_steps', default=2000, type=int,
  help='Number of optimization iterations')
@click.option('-ds', '--dense_steps', 'opt_dense_steps', default=2000, type=int,
  help='Number of optimization iterations')
@click.option('-v', '--video', 'opt_video', is_flag=True,
  help='Export a video for each dataset')
@click.pass_context
def cli(ctx, opt_folder_id, opt_latent_steps, opt_dense_steps, opt_video):
  """
  The full process:
  - Fetch new images from the cortex
  - Extract labels and base latents
  - Extract dense embeddings
  - Upload extract images to the cortex
  """
  folder = cortex_folder(opt_folder_id)
  files = download_cortex_files(opt_folder_id)
  unprocessed_files = find_unprocessed_files(files)
  if len(unprocessed_files) == 0:
    print("All files processed, nothing to do")
    return
  
  print("Processing folder {} ({}), {} new files".format(folder['name'], folder['id'], len(unprocessed_files)))

  tag = "folder_{}".format(folder['id'])
  paths = [file['path'] for file in unprocessed_files]
  
  find_nearest_vector_for_images(
    paths=paths,
    opt_dims=512,
    opt_steps=opt_dense_steps,
    opt_video=opt_video,
    opt_tag=tag,
    opt_limit=-1
  )

  params = params_dense_dict(tag)
  find_dense_embedding_for_images(params)