summaryrefslogtreecommitdiff
path: root/cli/app/commands/biggan/search_class.py
blob: 97573675698c4384658e571974919afc22b33aca (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
50
51
52
53
54
55
56
57
58
59
60
import click

import os
import time
from app.search.search_class import find_nearest_vector_for_images

@click.command('')
@click.option('-f', '--folder_id', 'opt_folder_id', type=int, default=0,
  help='Folder ID to process')
@click.option('-i', '--input', 'opt_fp_in', required=True, 
  help='Path to input image')
@click.option('-d', '--dims', 'opt_dims', default=512, type=int,
  help='Dimensions of BigGAN network (128, 256, 512)')
@click.option('-s', '--steps', 'opt_steps', default=2000, type=int,
  help='Number of optimization iterations')
@click.option('-l', '--limit', 'opt_limit', default=1000, type=int,
  help='Limit the number of images to process')
@click.option('-v', '--video', 'opt_video', is_flag=True,
  help='Export a video for each dataset')
@click.option('-d', '--dims', 'opt_dims', default=512, type=int,
  help='Dimensions of BigGAN network (128, 256, 512)')
@click.option('-t', '--tag', 'opt_tag', default='inverse_' + str(int(time.time() * 1000)), 
  help='Tag this dataset')
@click.option('-sc', '--stochastic_clipping', 'opt_stochastic_clipping', default=0,
  help='Compute feature loss')
@click.option('-lc', '--label_clipping', 'opt_label_clipping', default=0,
  help='Normalize labels every N steps')
@click.option('-feat', '--use_feature_detector', 'opt_use_feature_detector', is_flag=True,
  help='Compute feature loss')
@click.option('-ll', '--feature_layers', 'opt_feature_layers', default="1a,2a,4a,7a",
  help='Feature layers used for loss')
@click.option('-snap', '--snapshot_interval', 'opt_snapshot_interval', default=20,
  help='Interval to store sample images')

@click.pass_context
def cli(ctx, opt_folder_id, opt_fp_in, opt_dims, opt_steps, opt_video, opt_tag,
  opt_stochastic_clipping, opt_label_clipping, opt_use_feature_detector, opt_feature_layers, opt_snapshot_interval):
  """
  Search for an image (class vector) in BigGAN using gradient descent
  """
  if opt_folder_id != 0:
    folder = cortex_folder(opt_folder_id)
    files = download_cortex_files(opt_folder_id)
    unprocessed_files = [file for file in files if file['generated'] == 0 and file['type'] == 'image']
    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)))
    paths = [file['path'] for file in unprocessed_files]
  elif 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]

  opt_feature_layers = opt_feature_layers.split(',')

  find_nearest_vector_for_images(paths, opt_dims, opt_steps, opt_video, opt_tag, opt_limit,
    opt_stochastic_clipping, opt_label_clipping, opt_use_feature_detector, opt_feature_layers)