"""Visual Iterative Prompting functions.
Code to implement visual iterative prompting, an approach for querying VLMs.
"""
import copy
import dataclasses
import enum
import io
from typing import Optional, Tuple
import cv2
import matplotlib.pyplot as plt
import numpy as np
import scipy.stats
import vip_utils
@enum.unique
class SupportedEmbodiments(str, enum.Enum):
"""Embodiments supported by VIP."""
HF_DEMO = 'hf_demo'
@dataclasses.dataclass()
class Coordinate:
"""Coordinate with necessary information for visualizing annotation."""
# 2D image coordinates for the target annotation
xy: Tuple[int, int]
# Color and style of the coord.
color: Optional[float] = None
radius: Optional[int] = None
@dataclasses.dataclass()
class Sample:
"""Single Sample mapping actions to Coordinates."""
# 2D or 3D action
action: np.ndarray
# Coordinates for the main annotation
coord: Coordinate
# Coordinates for the text label
text_coord: Coordinate
# Label to display in the text label
label: str
class VisualIterativePrompter:
"""Visual Iterative Prompting class."""
def __init__(self, style, action_spec, embodiment):
self.embodiment = embodiment
self.style = style
self.action_spec = action_spec
self.fig_scale_size = None
# image preparer
# robot_to_image_canonical_coords
def action_to_coord(self, action, image, arm_xy, do_project=False):
"""Converts candidate action to image coordinate."""
return self.navigation_action_to_coord(
action=action, image=image, center_xy=arm_xy, do_project=do_project
)
def navigation_action_to_coord(
self, action, image, center_xy, do_project=False
):
"""Converts a ZXY or XY action to an image coordinate.
Conversion is done based on style['focal_offset'] and action_spec['scale'].
Args:
action: z, y, x action in robot action space
image: image
center_xy: x, y in image space
do_project: whether or not to project actions sampled outside the image to
the edge of the image
Returns:
Dict coordinate with image x, y, arrow color, and circle radius.
"""
if self.action_spec['scale'][0] == 0: # no z dimension
norm_action = [
(action[d] - self.action_spec['loc'][d])
/ (2 * self.action_spec['scale'][d])
for d in range(1, 3)
]
norm_action_y, norm_action_x = norm_action
norm_action_z = 0
else:
norm_action = [
(action[d] - self.action_spec['loc'][d])
/ (2 * self.action_spec['scale'][d])
for d in range(3)
]
norm_action_z, norm_action_y, norm_action_x = norm_action
focal_length = np.max([
0.2, # positive focal lengths only
self.style['focal_offset']
/ (self.style['focal_offset'] + norm_action_z),
])
image_x = center_xy[0] - (
self.action_spec['action_to_coord'] * norm_action_x * focal_length
)
image_y = center_xy[1] - (
self.action_spec['action_to_coord'] * norm_action_y * focal_length
)
if (
vip_utils.coord_outside_image(
Coordinate(xy=(image_x, image_y)), image, self.style['radius']
)
and do_project
):
# project the arrow to the edge of the image if too large
height, width, _ = image.shape
max_x = (
width - center_xy[0] - 2 * self.style['radius']
if norm_action_x < 0
else center_xy[0] - 2 * self.style['radius']
)
max_y = (
height - center_xy[1] - 2 * self.style['radius']
if norm_action_y < 0
else center_xy[1] - 2 * self.style['radius']
)
rescale_ratio = min(
np.abs([
max_x / (self.action_spec['action_to_coord'] * norm_action_x),
max_y / (self.action_spec['action_to_coord'] * norm_action_y),
])
)
image_x = (
center_xy[0]
- self.action_spec['action_to_coord'] * norm_action_x * rescale_ratio
)
image_y = (
center_xy[1]
- self.action_spec['action_to_coord'] * norm_action_y * rescale_ratio
)
return Coordinate(
xy=(int(image_x), int(image_y)),
color=0.1 * self.style['rgb_scale'],
radius=int(self.style['radius']),
)
def sample_actions(
self, image, arm_xy, loc, scale, true_action=None, max_itrs=1000
):
"""Sample actions from distribution.
Args:
image: image
arm_xy: x, y in image space of arm
loc: action distribution mean to sample from
scale: action distribution variance to sample from
true_action: action taken in demonstration if available
max_itrs: number of tries to get a valid sample
Returns:
samples: Samples with associated actions, coords, text_coords, labels.
"""
image = copy.deepcopy(image)
samples = []
actions = []
coords = []
text_coords = []
labels = []
# Keep track of oracle action if available.
true_label = None
if true_action is not None:
actions.append(true_action)
coord = self.action_to_coord(true_action, image, arm_xy)
coords.append(coord)
text_coords.append(
vip_utils.coord_to_text_coord(coords[-1], arm_xy, coord.radius)
)
true_label = np.random.randint(self.style['num_samples'])
# labels.append(str(true_label) + '*')
labels.append(str(true_label))
# Generate all action samples.
for i in range(self.style['num_samples']):
if i == true_label:
continue
itrs = 0
# Generate action scaled appropriately.
action = np.clip(
np.random.normal(loc, scale),
self.action_spec['min'],
self.action_spec['max'],
)
# Convert sampled action to image coordinates.
coord = self.action_to_coord(action, image, arm_xy)
# Resample action if it results in invalid image annotation.
adjusted_scale = np.array(scale)
while (
vip_utils.is_invalid_coord(
coord, coords, self.style['radius'] * 1.5, image
)
or vip_utils.coord_outside_image(coord, image, self.style['radius'])
) and itrs < max_itrs:
action = np.clip(
np.random.normal(loc, adjusted_scale),
self.action_spec['min'],
self.action_spec['max'],
)
coord = self.action_to_coord(action, image, arm_xy)
itrs += 1
# increase sampling range slightly if not finding a good sample
adjusted_scale *= 1.1
if itrs == max_itrs:
# If the final iteration results in invalid annotation, just clip
# to edge of image.
coord = self.action_to_coord(action, image, arm_xy, do_project=True)
# Compute image coordinates of text labels.
radius = coord.radius
text_coord = Coordinate(
xy=vip_utils.coord_to_text_coord(coord, arm_xy, radius)
)
actions.append(action)
coords.append(coord)
text_coords.append(text_coord)
labels.append(str(i))
for i in range(len(actions)):
sample = Sample(
action=actions[i],
coord=coords[i],
text_coord=text_coords[i],
label=str(i),
)
samples.append(sample)
return samples
def add_arrow_overlay_plt(self, image, samples, arm_xy):
"""Add arrows and circles to the image.
Args:
image: image
samples: Samples to visualize.
arm_xy: x, y image coordinates for EEF center.
log_image: Boolean for whether to save to CNS.
Returns:
image: image with visual prompts.
"""
# Add transparent arrows and circles
overlay = image.copy()
(original_image_height, original_image_width, _) = image.shape
white = (
self.style['rgb_scale'],
self.style['rgb_scale'],
self.style['rgb_scale'],
)
# Add arrows.
for sample in samples:
color = sample.coord.color
cv2.arrowedLine(
overlay, arm_xy, sample.coord.xy, color, self.style['thickness']
)
image = cv2.addWeighted(
overlay,
self.style['arrow_alpha'],
image,
1 - self.style['arrow_alpha'],
0,
)
overlay = image.copy()
# Add circles.
for sample in samples:
color = sample.coord.color
radius = sample.coord.radius
cv2.circle(
overlay,
sample.text_coord.xy,
radius,
color,
self.style['thickness'] + 1,
)
cv2.circle(overlay, sample.text_coord.xy, radius, white, -1)
image = cv2.addWeighted(
overlay,
self.style['circle_alpha'],
image,
1 - self.style['circle_alpha'],
0,
)
dpi = plt.rcParams['figure.dpi']
if self.fig_scale_size is None:
# test saving a figure to decide size for text figure
fig_size = (original_image_width / dpi, original_image_height / dpi)
plt.subplots(1, figsize=fig_size)
plt.imshow(image, cmap='binary')
plt.axis('off')
fig = plt.gcf()
fig.tight_layout(pad=0)
buf = io.BytesIO()
plt.savefig(buf, format='png')
plt.close()
buf.seek(0)
test_image = cv2.imdecode(
np.frombuffer(buf.getvalue(), dtype=np.uint8), cv2.IMREAD_COLOR
)
self.fig_scale_size = original_image_width / test_image.shape[1]
# Add text to figure.
fig_size = (
self.fig_scale_size * original_image_width / dpi,
self.fig_scale_size * original_image_height / dpi,
)
plt.subplots(1, figsize=fig_size)
plt.imshow(image, cmap='binary')
for sample in samples:
plt.text(
sample.text_coord.xy[0],
sample.text_coord.xy[1],
sample.label,
ha='center',
va='center',
color='k',
fontsize=self.style['fontsize'],
)
# Compile image.
plt.axis('off')
fig = plt.gcf()
fig.tight_layout(pad=0)
buf = io.BytesIO()
plt.savefig(buf, format='png')
plt.close()
image = cv2.imdecode(
np.frombuffer(buf.getvalue(), dtype=np.uint8), cv2.IMREAD_COLOR
)
image = cv2.resize(image, (original_image_width, original_image_height))
image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
return image
def fit(self, values, samples):
"""Fit a loc and scale to selected actions.
Args:
values: list of selected labels
samples: list of all Samples
Returns:
loc: mean of selected distribution
scale: variance of selected distribution
"""
actions = [sample.action for sample in samples]
labels = [sample.label for sample in samples]
if not values: # revert to initial distribution
print('GPT failed to return integer arrows')
loc = self.action_spec['loc']
scale = self.action_spec['scale']
elif len(values) == 1: # single response, add a distribution over it
index = np.where([label == str(values[-1]) for label in labels])[0][0]
action = actions[index]
print('action', action)
loc = action
scale = self.action_spec['min_scale']
else: # fit distribution
selected_actions = []
for value in values:
idx = np.where([label == str(value) for label in labels])[0][0]
selected_actions.append(actions[idx])
print('selected_actions', selected_actions)
loc_scale = [
scipy.stats.norm.fit([action[d] for action in selected_actions])
for d in range(3)
]
loc = [loc_scale[d][0] for d in range(3)]
scale = np.clip(
[loc_scale[d][1] for d in range(3)],
self.action_spec['min_scale'],
None,
)
print('loc', loc, '\nscale', scale)
return loc, scale