Exemple #1
0
def main():
    """Main function."""
    args = parse_args()
    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id
    assert os.path.isfile(args.image_path)
    output_dir = args.output_dir or f'results/inversion/test'
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    inverter = StyleGANInverter(
        args.model_name,
        mode=args.mode,
        learning_rate=args.learning_rate,
        iteration=args.num_iterations,
        reconstruction_loss_weight=1.0,
        perceptual_loss_weight=args.loss_weight_feat,
        regularization_loss_weight=args.loss_weight_enc,
        clip_loss_weight=args.loss_weight_clip,
        description=args.description,
        logger=None)
    image_size = inverter.G.resolution

    # Invert the given image.
    image = resize_image(load_image(args.image_path), (image_size, image_size))
    _, viz_results = inverter.easy_invert(image, num_viz=args.num_results)

    if args.mode == 'man':
        image_name = os.path.splitext(os.path.basename(args.image_path))[0]
    else:
        image_name = 'gen'
    save_image(f'{output_dir}/{image_name}_enc.png', viz_results[1])
    save_image(f'{output_dir}/{image_name}_inv.png', viz_results[-1])
    print(f'save {image_name} in {output_dir}')
Exemple #2
0
def main():
    """Main function."""

    inverter = StyleGANInverter(model_name,
                                mode=mode,
                                learning_rate=ini_lr,
                                iteration=step,
                                reconstruction_loss_weight=lambda_l2,
                                perceptual_loss_weight=lambda_feat,
                                regularization_loss_weight=lambda_enc,
                                clip_loss_weight=lambda_clip,
                                description=description)
    image_size = inverter.G.resolution

    text_inputs = torch.cat([clip.tokenize(description)]).cuda()

    # Invert images.
    # uploaded_file = uploaded_file.read()
    if uploaded_file is not None:
        image = Image.open(uploaded_file)
        # st.image(image, caption='Uploaded Image.', use_column_width=True)
        # st.write("")
        st.write("Just a second...")

        image = resize_image(np.array(image), (image_size, image_size))
        _, viz_results = inverter.easy_invert(image, 1)
        if mode == 'man':
            final_result = np.hstack([image, viz_results[-1]])
        else:
            final_result = np.hstack([viz_results[1], viz_results[-1]])

        # return final_result
        with st.beta_container():
            st.image(final_result, use_column_width=True)
Exemple #3
0
    def predict(self, image, description):

        self.args.description = description
        self.args.image_path = image

        inverter = StyleGANInverter(
            self.args.model_name,
            mode=self.args.mode,
            learning_rate=self.args.learning_rate,
            iteration=self.args.num_iterations,
            reconstruction_loss_weight=1.0,
            perceptual_loss_weight=self.args.loss_weight_feat,
            regularization_loss_weight=self.args.loss_weight_enc,
            clip_loss_weight=self.args.loss_weight_clip,
            description=self.args.description,
            logger=None,
        )
        image_size = inverter.G.resolution

        # Invert the given image.
        image = resize_image(
            load_image(str(self.args.image_path)), (image_size, image_size)
        )
        _, viz_results = inverter.easy_invert(image, num_viz=self.args.num_results)

        out_path = Path(tempfile.mkdtemp()) / "out.png"
        save_image(str(out_path), viz_results[-1])
        return out_path
def main():
  """Main function."""
  args = parse_args()
  os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id
  assert os.path.exists(args.image_list)
  image_list_name = os.path.splitext(os.path.basename(args.image_list))[0]
  output_dir = args.output_dir or f'results/inversion/{image_list_name}'
  logger = setup_logger(output_dir, 'inversion.log', 'inversion_logger')

  logger.info(f'Loading model.')
  inverter = StyleGANInverter(
      args.model_name,
      learning_rate=args.learning_rate,
      iteration=args.num_iterations,
      reconstruction_loss_weight=1.0,
      perceptual_loss_weight=args.loss_weight_feat,
      regularization_loss_weight=args.loss_weight_enc,
      logger=logger)
  image_size = inverter.G.resolution

  # Load image list.
  logger.info(f'Loading image list.')
  image_list = []
  with open(args.image_list, 'r') as f:
    for line in f:
      image_list.append(line.strip())

  # Initialize visualizer.
  save_interval = args.num_iterations // args.num_results
  headers = ['Name', 'Original Image', 'Encoder Output']
  for step in range(1, args.num_iterations + 1):
    if step == args.num_iterations or step % save_interval == 0:
      headers.append(f'Step {step:06d}')
  viz_size = None if args.viz_size == 0 else args.viz_size
  visualizer = HtmlPageVisualizer(
      num_rows=len(image_list), num_cols=len(headers), viz_size=viz_size)
  visualizer.set_headers(headers)

  # Invert images.
  logger.info(f'Start inversion.')
  latent_codes = []
  for img_idx in tqdm(range(len(image_list)), leave=False):
    image_path = image_list[img_idx]
    image_name = os.path.splitext(os.path.basename(image_path))[0]
    image = resize_image(load_image(image_path), (image_size, image_size))
    code, viz_results = inverter.easy_invert(image, num_viz=args.num_results)
    latent_codes.append(code)
    save_image(f'{output_dir}/{image_name}_ori.png', image)
    save_image(f'{output_dir}/{image_name}_enc.png', viz_results[1])
    save_image(f'{output_dir}/{image_name}_inv.png', viz_results[-1])
    visualizer.set_cell(img_idx, 0, text=image_name)
    visualizer.set_cell(img_idx, 1, image=image)
    for viz_idx, viz_img in enumerate(viz_results[1:]):
      visualizer.set_cell(img_idx, viz_idx + 2, image=viz_img)

  # Save results.
  os.system(f'cp {args.image_list} {output_dir}/image_list.txt')
  np.save(f'{output_dir}/inverted_codes.npy',
          np.concatenate(latent_codes, axis=0))
  visualizer.save(f'{output_dir}/inversion.html')
Exemple #5
0
def setup_inverter(model_name,
                   num_iterations=100,
                   regularization_loss_weight=2):
    global inverter
    setup_model(model_name)
    inverter = StyleGANInverter(
        models[model_name]['name'],
        learning_rate=0.01,
        iteration=num_iterations,
        reconstruction_loss_weight=1.0,
        perceptual_loss_weight=5e-5,
        regularization_loss_weight=regularization_loss_weight)
def main():
  """Main function."""
  args = parse_args()
  os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id
  assert os.path.exists(args.target_list)
  target_list_name = os.path.splitext(os.path.basename(args.target_list))[0]
  assert os.path.exists(args.context_list)
  context_list_name = os.path.splitext(os.path.basename(args.context_list))[0]
  output_dir = args.output_dir or f'results/diffusion'
  job_name = f'{target_list_name}_TO_{context_list_name}'
  logger = setup_logger(output_dir, f'{job_name}.log', f'{job_name}_logger')

  logger.info(f'Loading model.')
  inverter = StyleGANInverter(
      args.model_name,
      learning_rate=args.learning_rate,
      iteration=args.num_iterations,
      reconstruction_loss_weight=1.0,
      perceptual_loss_weight=args.loss_weight_feat,
      regularization_loss_weight=0.0,
      logger=logger)
  image_size = inverter.G.resolution

  # Load image list.
  logger.info(f'Loading target images and context images.')
  target_list = []
  with open(args.target_list, 'r') as f:
    for line in f:
      target_list.append(line.strip())
  num_targets = len(target_list)
  context_list = []
  with open(args.context_list, 'r') as f:
    for line in f:
      context_list.append(line.strip())
  num_contexts = len(context_list)
  num_pairs = num_targets * num_contexts

  # Initialize visualizer.
  save_interval = args.num_iterations // args.num_results
  headers = ['Target Image', 'Context Image', 'Stitched Image',
             'Encoder Output']
  for step in range(1, args.num_iterations + 1):
    if step == args.num_iterations or step % save_interval == 0:
      headers.append(f'Step {step:06d}')
  viz_size = None if args.viz_size == 0 else args.viz_size
  visualizer = HtmlPageVisualizer(
      num_rows=num_pairs, num_cols=len(headers), viz_size=viz_size)
  visualizer.set_headers(headers)

  # Diffuse images.
  logger.info(f'Start diffusion.')
  latent_codes = []
  for target_idx in tqdm(range(num_targets), desc='Target ID', leave=False):
    # Load target.
    target_image = resize_image(load_image(target_list[target_idx]),
                                (image_size, image_size))
    visualizer.set_cell(target_idx * num_contexts, 0, image=target_image)
    for context_idx in tqdm(range(num_contexts), desc='Context ID',
                            leave=False):
      row_idx = target_idx * num_contexts + context_idx
      context_image = resize_image(load_image(context_list[context_idx]),
                                   (image_size, image_size))
      visualizer.set_cell(row_idx, 1, image=context_image)
      code, viz_results = inverter.easy_diffuse(target=target_image,
                                                context=context_image,
                                                center_x=args.center_x,
                                                center_y=args.center_y,
                                                crop_x=args.crop_size,
                                                crop_y=args.crop_size,
                                                num_viz=args.num_results)
      for viz_idx, viz_img in enumerate(viz_results):
        visualizer.set_cell(row_idx, viz_idx + 2, image=viz_img)
      latent_codes.append(code)

  # Save results.

  os.system(f'cp {args.target_list} {output_dir}/target_list.txt')
  os.system(f'cp {args.context_list} {output_dir}/context_list.txt')
  np.save(f'{output_dir}/{job_name}_inverted_codes.npy',
          np.concatenate(latent_codes, axis=0))
  visualizer.save(f'{output_dir}/{job_name}.html')