pivot-iterative-visual-optimization/pivot-demo

1"""Utils for visual iterative prompting.2 3A number of utility functions for VIP.4"""5 6import re7 8import matplotlib.pyplot as plt9import numpy as np10import scipy.spatial.distance as distance11 12 13def min_dist(coord, coords):14  if not coords:15    return np.inf16  xys = np.asarray([[coord.xy] for coord in coords])17  return np.linalg.norm(xys - np.asarray(coord.xy), axis=-1).min()18 19 20def coord_outside_image(coord, image, radius):21  (height, image_width, _) = image.shape22  x, y = coord.xy23  x_outside = x > image_width - 2 * radius or x < 2 * radius24  y_outside = y > height - 2 * radius or y < 2 * radius25  return x_outside or y_outside26 27 28def is_invalid_coord(coord, coords, radius, image):29  # invalid if too close to others or outside of the image30  pos_overlaps = min_dist(coord, coords) < 1.5 * radius31  return pos_overlaps or coord_outside_image(coord, image, radius)32 33 34def angle_mag_2_x_y(angle, mag, arm_coord, is_circle=False, radius=40):35  x, y = arm_coord36  x += int(np.cos(angle) * mag)37  y += int(np.sin(angle) * mag)38  if is_circle:39    x += int(np.cos(angle) * radius * np.sign(mag))40    y += int(np.sin(angle) * radius * np.sign(mag))41  return x, y42 43 44def coord_to_text_coord(coord, arm_coord, radius):45  delta_coord = np.asarray(coord.xy) - arm_coord46  if np.linalg.norm(delta_coord) == 0:47    return arm_coord48  return (49      int(coord.xy[0] + radius * delta_coord[0] / np.linalg.norm(delta_coord)),50      int(coord.xy[1] + radius * delta_coord[1] / np.linalg.norm(delta_coord)),51  )52 53 54def parse_response(response, answer_key='Arrow: ['):55  values = []56  if answer_key in response:57    print('parse_response from answer_key')58    arrow_response = response.split(answer_key)[-1].split(']')[0]59    for val in map(int, re.findall(r'\d+', arrow_response)):60      values.append(val)61  else:62    print('parse_response for all ints')63    for val in map(int, re.findall(r'\d+', response)):64      values.append(val)65  return values66 67 68def compute_errors(action, true_action, verbose=False):69  """Compute errors between a predicted action and true action."""70  l2_error = np.linalg.norm(action - true_action)71  cos_sim = 1 - distance.cosine(action, true_action)72  l2_xy_error = np.linalg.norm(action[-2:] - true_action[-2:])73  cos_xy_sim = 1 - distance.cosine(action[-2:], true_action[-2:])74  z_error = np.abs(action[0] - true_action[0])75  errors = {76      'l2': l2_error,77      'cos_sim': cos_sim,78      'l2_xy_error': l2_xy_error,79      'cos_xy_sim': cos_xy_sim,80      'z_error': z_error,81  }82 83  if verbose:84    print('action: \t', [f'{a:.3f}' for a in action])85    print('true_action \t', [f'{a:.3f}' for a in true_action])86    print(f'l2: \t\t{l2_error:.3f}')87    print(f'l2_xy_error: \t{l2_xy_error:.3f}')88    print(f'cos_sim: \t{cos_sim:.3f}')89    print(f'cos_xy_sim: \t{cos_xy_sim:.3f}')90    print(f'z_error: \t{z_error:.3f}')91 92  return errors93 94 95def plot_errors(all_errors, error_types=None):96  """Plot errors across iterations."""97  if error_types is None:98    error_types = [99        'l2',100        'l2_xy_error',101        'z_error',102        'cos_sim',103        'cos_xy_sim',104    ]105 106  _, axs = plt.subplots(2, 3, figsize=(15, 8))107  for i, error_type in enumerate(error_types):  # go through each error type108    all_iter_errors = {}109    for error_by_iter in all_errors:  # go through each call110      for itr in error_by_iter:  # go through each iteration111        if itr in all_iter_errors:  # add error to the iteration it happened112          all_iter_errors[itr].append(error_by_iter[itr][error_type])113        else:114          all_iter_errors[itr] = [error_by_iter[itr][error_type]]115 116    mean_iter_errors = [117        np.mean(all_iter_errors[itr]) for itr in all_iter_errors118    ]119 120    axs[i // 3, i % 3].plot(all_iter_errors.keys(), mean_iter_errors)121    axs[i // 3, i % 3].set_title(error_type)122  plt.show()123