class VispyUI(object):
_CUSTOM_LAYERS = {'CircleLayer': CircleLayer, 'LineLayer': LineLayer}
# Variables possible for kwargs, use these defaults if missing from kwargs
_DEFAULTS = {
'n_cores': 0,
'epochs_per_cycle': 10,
'display_multiplier': 3,
'border': 1.0
}
def __init__(self,
state_file=None,
image=None,
suffix=None,
out_dir="output",
**kwargs):
# epochs_per_cycle=None, n_cores=None, display_multiplier=None, border=None):
if int(image is None) + int(state_file is None) != 1:
raise Exception("Need input image, or state to restore.")
# state
self._tempdir = tempfile.mkdtemp(prefix="imageStretch_")
if state_file is not None:
self._load_state(state_file)
else:
self._suffix = "" if suffix is None else "_%s" % (suffix, )
self._out_dir = out_dir
self._image = image
n_cores = kwargs[
'n_cores'] if 'n_cores' in kwargs else self._DEFAULTS['n_cores']
self._sim = ScaleInvariantImage(image, n_cores=n_cores, state=None)
self._index = 0
self._increment_run_number()
# let these be overridden by command line if they're still none
for arg in self._DEFAULTS:
if arg in kwargs and kwargs[arg] is not None:
setattr(self, privatize(arg), kwargs[arg])
elif not hasattr(self, privatize(arg)):
setattr(self, privatize(arg), self._DEFAULTS[arg])
logging.info("Using args:")
for arg in self._DEFAULTS:
logging.info("\t%s: %s" % (arg, getattr(self, privatize(arg))))
self._interactive = not kwargs['just_image']
if kwargs['just_image']:
self._write_image()
return
# set up VISPY display
self._canvas = vispy.scene.SceneCanvas(
keys='interactive', show=True, title="Route editor -- H for help")
self._viewbox = self._canvas.central_widget.add_view()
self._viewbox.camera = vispy.scene.cameras.PanZoomCamera(
parent=self._viewbox.scene, aspect=1)
self._add_graphics_elements()
self.set_text_box_state(False)
self._canvas.events.key_press.connect(self._on_key_press)
self._canvas.events.close.connect(self._on_close)
self._closing = False
self._set_image(self._image / 255.0)
self._worker = Thread(target=self._train_thread)
self._worker.start()
logging.info("Started worker thread.")
def is_interactive(self):
return self._interactive
def _on_close(self, event):
self._closing = True
logging.info("Shutting down...")
def _increment_run_number(self):
if not os.path.exists(self._out_dir):
os.mkdir(self._out_dir)
files = [
f for f in os.listdir(self._out_dir)
if f.startswith('output') and f.endswith('.png')
]
numbers = [
int(re.search('^output_([0-9]+)_.+.png', f).groups()[0])
for f in files
]
self._run = np.max(numbers) + 1 if len(numbers) > 0 else 0
def _save_state(self):
logging.info("Saving state...")
model_filename = "model_%i%s.tf" % (self._run, self._suffix)
model_filepath = os.path.join(self._out_dir, model_filename)
temp_filepath = os.path.join(self._tempdir, model_filename)
self._sim.get_model().save(temp_filepath)
with open(temp_filepath, 'r') as infile:
model_bin = infile.read()
model = {
'model_bin': model_bin,
'epc': self._epochs_per_cycle,
'dm': self._display_multiplier,
'suffix': self._suffix,
'out_dir': self._out_dir,
'image': self._image,
'border': self._border,
'index': self._index,
'run': self._run
}
with open(model_filepath, 'w') as outfile:
cp.dump(model, outfile)
logging.info("Wrote: %s" % (model_filepath, ))
def _load_state(self, model_filepath):
logging.info("loading state...")
with open(model_filepath, 'r') as infile:
state = cp.load(infile)
self._epochs_per_cycle = state['epc']
self._display_multiplier = state['dm']
self._suffix = state['suffix']
self._out_dir = state['out_dir']
self._image = state['image']
self._index = state['index'] + 1
self._run = state['run']
self._border = state['border']
model_filename = "model_%i%s.tf" % (self._run, self._suffix)
model_filepath = os.path.join(self._tempdir, model_filename)
with open(model_filepath, 'w') as outfile:
outfile.write(state['model_bin'])
state['model'] = keras.models.load_model(
model_filepath, custom_objects=self._CUSTOM_LAYERS)
self._sim = ScaleInvariantImage(state=state)
def _train_thread(self):
tf_session, tf_graph = self._sim.get_session_and_graph()
with tf_session.as_default():
with tf_graph.as_default():
while not self._closing:
for ep in range(self._epochs_per_cycle):
if self._closing:
break
self._sim.train_epochs(1)
self._save_state()
if self._closing:
break
img = self._write_image()
self._set_image(img)
self._index += 1
if self._closing:
break
self._save_state()
logging.info("Close detected, shutting down worker thread.")
logging.info("Deleting work dir: %s" % (self._tempdir, ))
shutil.rmtree(self._tempdir)
def _write_image(self):
img = self._sim.get_display_image(
np.array(self._sim.get_image().shape) * self._display_multiplier,
margin=self._border)
out_filename = "output_%i%s_%.8i.png" % (self._run, self._suffix,
self._index)
out_path = os.path.join(self._out_dir, out_filename)
while os.path.exists(out_path):
self._index += 1
out_filename = "output_%i%s_%.8i.png" % (self._run, self._suffix,
self._index)
out_path = os.path.join(self._out_dir, out_filename)
cv2.imwrite(out_path, np.uint8(255 * img[:, :, ::-1]))
logging.info("Wrote: %s" % (out_path, ))
return img
def _add_graphics_elements(self):
"""
Create the VISPY graphics objects
"""
# Image
self._image_object = Image(self._image, parent=self._viewbox.scene)
self._image_object.set_gl_state('translucent', depth_test=False)
self._image_object.order = 1
self._image_object.visible = True
self._text_box_width = 150
self._text_box_height = 40
self._text_box_offset = 10
# Text background box in upper-left corner
self._text_bkg_rect = vispy.scene.visuals.Rectangle(
[
self._text_box_width / 2 + self._text_box_offset,
self._text_box_height / 2 + self._text_box_offset
],
color=[0.1, 0.0, 0.0, .8],
border_color=[0.1, 0.0, 0.0],
border_width=2,
height=self._text_box_height,
width=self._text_box_width,
radius=10.0,
parent=self._canvas.scene)
self._text_bkg_rect.set_gl_state('translucent', depth_test=False)
self._text_bkg_rect.visible = True
self._text_bkg_rect.order = 2
# Text
self._text = "?"
self._text_pos = [
self._text_box_offset + 10, self._text_box_offset + 10
]
self._text_obj = vispy.scene.visuals.Text(self._text,
parent=self._canvas.scene,
color=[0.9, 0.8, 0.8],
anchor_x='left',
anchor_y='top')
self._text_obj.pos = self._text_pos
self._text_obj.font_size = 18
self._text_obj.visible = True
self._text_obj.order = 3
def set_text_box_state(self, state):
self._text_bkg_rect.visible = state
self._text_obj.visible = state
def _change_text(self, new_text):
if new_text is not None:
self._text = new_text
else:
self._text = ""
self._text_obj.text = self._text
def _set_image(self, image=None):
xmin, ymin = 0, 0
xmax, ymax = image.shape[1], image.shape[0]
self._viewbox.camera.set_range(x=(xmin, xmax), y=(ymin, ymax), z=None)
self._image_object.set_data(flip(image))
self._image_object.visible = True
def _on_key_press(self, ev):
if ev.key.name == 'Escape':
self._canvas.close()
else:
self._change_text(ev.key.name)
logging.info("Unknown keypress: %s" % (ev.key.name, ))
class VispyUI(object):
"""
Display an image, rectangle, text, enable pan/tilt, keyboard
"""
def __init__(self):
self._make_new_random_image()
# state
self._last_mouse_pos = np.array([0.0, 0.0])
self._load_data() # Load data in after the view has been setup
self._canvas = vispy.scene.SceneCanvas(
keys='interactive', show=True, title="Route editor -- H for help")
self._viewbox = self._canvas.central_widget.add_view()
self._viewbox.camera = vispy.scene.cameras.PanZoomCamera(
parent=self._viewbox.scene, aspect=1)
self._add_graphics_elements()
self._canvas.events.key_press.connect(self._on_key_press)
self._canvas.events.mouse_move.connect(self._on_mouse_move)
self._canvas.events.mouse_press.connect(self._on_mouse_press)
self._canvas.events.mouse_release.connect(self._on_mouse_release)
self._set_map_image(self._image)
def _make_new_random_image(self):
self._h, self._w = 100 + int(np.random.rand(1) * 100), 100 + int(
np.random.rand(1) * 100)
self._image = np.uint8(
np.zeros((self._h, self._w, 3)) +
255 * np.random.rand(self._h * self._w * 3).reshape(
(self._h, self._w, 3)))
def _add_graphics_elements(self):
"""
Create all the graphics objects (VISPY objects), put them on the canvas.
"""
# Image
self._image_object = Image(self._image, parent=self._viewbox.scene)
self._image_object.set_gl_state('translucent', depth_test=False)
self._image_object.order = 1
self._image_object.visible = True
self._text_box_width = 150
self._text_box_height = 40
self._text_box_offset = 10
# Text background box in upper-left corner
self._text_bkg_rect = vispy.scene.visuals.Rectangle(
[
self._text_box_width / 2 + self._text_box_offset,
self._text_box_height / 2 + self._text_box_offset
],
color=[0.1, 0.0, 0.0, .8],
border_color=[0.1, 0.0, 0.0],
border_width=2,
height=self._text_box_height,
width=self._text_box_width,
radius=10.0,
parent=self._canvas.scene)
self._text_bkg_rect.set_gl_state('translucent', depth_test=False)
self._text_bkg_rect.visible = True
self._text_bkg_rect.order = 2
# Text
self._text = "?"
self._text_pos = [
self._text_box_offset + 10, self._text_box_offset + 10
]
self._text_obj = vispy.scene.visuals.Text(self._text,
parent=self._canvas.scene,
color=[0.9, 0.8, 0.8],
anchor_x='left',
anchor_y='top')
self._text_obj.pos = self._text_pos
self._text_obj.font_size = 18
self._text_obj.visible = True
self._text_obj.order = 3
def _change_text(self, new_text):
if new_text is not None:
self._text = new_text
else:
self._text = ""
self._text_obj.text = self._text
def _load_data(self):
"""
Downlaod (if necessary) and load into memory the costmap, route and scrubber deck state controls.
"""
self._images = []
def _set_map_image(self, disp_image=None):
"""
Extract an image to display from the current cost map.
"""
if disp_image is not None:
self._image = disp_image
xmin, ymin = 0, 0
xmax, ymax = self._image.shape[1], self._image.shape[0]
self._viewbox.camera.set_range(x=(xmin, xmax), y=(ymin, ymax), z=None)
self._image_object.set_data(self._image)
self._image_object.visible = True
def _on_key_press(self, ev):
"""
vispy keyboard callback
"""
if not ev or not ev.key:
# on Mac sometimes this callback is triggered with
# NoneType key when clicking between desktops or full screen apps
return
control_pressed = 'Control' in [e.name for e in ev.modifiers]
shift_pressed = 'Shift' in [e.name for e in ev.modifiers]
print "Shift: %s\nControl:%s" % (shift_pressed, control_pressed)
if ev.key.name == "Space":
self._make_new_random_image()
self._set_map_image()
elif ev.key.name == "Z":
print "Z"
elif ev.key.name == "M":
print "M"
else:
self._change_text(ev.key.name)
logging.info("Unknown keypress: %s" % (ev.key.name, ))
def _on_mouse_move(self, event):
"""
vispy mouse motion callback
All mouse events are in map coordinate frame
"""
self._last_mouse_pos = event.pos
def _on_mouse_press(self, event):
'''
Vispy mouse button press callback.
All mouse events are in map coordinate frame
'''
self._last_mouse_pos = event.pos
def _on_mouse_release(self, event):
"""
Vispy mouse button release callback.
All mouse events are in map coordinate frame
"""
self._last_mouse_pos = event.pos
class ImageSorter(object):
"""
Display an image, rectangle, text, enable pan/tilt, keyboard
"""
def __init__(self, image_dir, out_file=None, max_load=10):
self._in_dir = os.path.abspath(os.path.expanduser(image_dir))
self._out_file = out_file if out_file is not None else os.path.join(
self._in_dir, "results.json")
self._images = LabeledImageSet(image_dir, self._out_file,
self._out_file)
self._current_filename, self._current_index = self._images.get_unlabeled(
)
# state
self._last_mouse_pos = np.array([0.0, 0.0])
self._canvas = vispy.scene.SceneCanvas(
keys='interactive', show=True, title="Route editor -- H for help")
self._viewbox = self._canvas.central_widget.add_view()
self._viewbox.camera = vispy.scene.cameras.PanZoomCamera(
parent=self._viewbox.scene, aspect=1)
self._add_graphics_elements()
self._canvas.events.key_press.connect(self._on_key_press)
self._canvas.events.mouse_move.connect(self._on_mouse_move)
self._canvas.events.mouse_press.connect(self._on_mouse_press)
self._canvas.events.mouse_release.connect(self._on_mouse_release)
self._canvas.events.resize.connect(self._on_resize)
self._update_image()
def _save_results(self):
self._images.save()
def _add_graphics_elements(self):
"""
Create all the graphics objects (VISPY objects), put them on the canvas.
"""
# Image
self._image_object = Image(None, parent=self._viewbox.scene)
self._image_object.set_gl_state('translucent', depth_test=False)
self._image_object.order = 1
self._image_object.visible = True
self._text_box_width = 150
self._text_box_height = 60
self._text_box_offset = 10
# Text background box in upper-left corner
self._text_bkg_rect = vispy.scene.visuals.Rectangle(
[
self._text_box_width / 2 + self._text_box_offset,
self._text_box_height / 2 + self._text_box_offset
],
color=[0.1, 0.0, 0.0, .8],
border_color=[0.1, 0.0, 0.0],
border_width=2,
height=self._text_box_height,
width=self._text_box_width,
radius=10.0,
parent=self._canvas.scene)
self._text_bkg_rect.set_gl_state('translucent', depth_test=False)
self._text_bkg_rect.visible = True
self._text_bkg_rect.order = 2
self._resize_text_bkg_box()
# Text
self._font1_size = 10
self._font2_size = 18
self._vspace = self._font1_size * 2.2
self._text_pos = [
self._text_box_offset + 10, self._text_box_offset + 10
]
self._text_obj = vispy.scene.visuals.Text("",
parent=self._canvas.scene,
color=[0.9, 0.8, 0.8],
anchor_x='left',
anchor_y='top')
self._text_obj.pos = self._text_pos
self._text_obj.font_size = self._font1_size
self._text_obj.visible = True
self._text_obj.order = 3
self._text2_pos = [self._text_pos[0], self._vspace + self._text_pos[1]]
self._text2_obj = vispy.scene.visuals.Text("",
parent=self._canvas.scene,
color=[0.9, 0.8, 0.8],
anchor_x='left',
anchor_y='top')
self._text2_obj.pos = self._text2_pos
self._text2_obj.font_size = self._font2_size
self._text2_obj.visible = True
self._text2_obj.order = 3
def _resize_text_bkg_box(self):
logging.info('resize')
self._text_bkg_rect.center = np.array([
self._canvas.size[0] / 2,
self._text_box_height / 2 + self._text_box_offset
])
self._text_bkg_rect.width = self._canvas.size[
0] - 2 * self._text_box_offset
def _on_resize(self, event):
self._resize_text_bkg_box()
def _update_image(self):
"""
Extract an image to display from the current cost map.
"""
self._current_filename = self._images.get_i(self._current_index)
self._current_image = self._images.get_image(self._current_filename)
logging.info("Got new image: %s, %s" %
(self._current_image.shape, self._current_filename))
xmin, ymin = 0, 0
xmax, ymax = self._current_image.shape[1], self._current_image.shape[0]
self._viewbox.camera.set_range(x=(xmin, xmax), y=(ymin, ymax), z=None)
self._image_object.set_data(self._current_image)
self._image_object.visible = True
self._update_text()
def _update_text(self):
result = self._images.get_label(self._current_filename)
result_str = "%i" % (result, ) if result is not None else "(unlabeled)"
text1 = self._current_filename
text2 = "%i / %i - label: %s" % (self._current_index,
self._images.get_n(), result_str)
self._text_obj.text = text1
self._text2_obj.text = text2
self._text_obj.visible = True
self._text2_obj.visible = True
def _on_key_press(self, ev):
"""
vispy keyboard callback
"""
if not ev or not ev.key:
# on Mac sometimes this callback is triggered with
# NoneType key when clicking between desktops or full screen apps
return
# control_pressed = 'Control' in [e.name for e in ev.modifiers]
shift_pressed = 'Shift' in [e.name for e in ev.modifiers]
# print "Shift: %s\nControl:%s" % (shift_pressed, control_pressed)
if ev.key.name == "Left":
self._advance_index(-1)
elif ev.key.name == "Right":
self._advance_index(1)
elif ev.key.name == "X":
self._classify_and_go_to_next(
0, skip_to_next_unlabeled=not shift_pressed)
elif ev.key.name == "O":
self._classify_and_go_to_next(
1, skip_to_next_unlabeled=not shift_pressed)
else:
logging.info("Unknown keypress: %s" % (ev.key.name, ))
def _advance_index(self, direction):
new_index = self._current_index + direction
if new_index >= 0 and new_index <= self._images.get_n() - 1:
self._current_index = new_index
self._update_image()
self._update_text()
def _classify_and_go_to_next(self, label, skip_to_next_unlabeled=True):
self._images.apply_label(self._current_filename, label)
self._images.save()
if skip_to_next_unlabeled:
self._current_filename, self._current_index = self._images.get_unlabeled(
)
self._update_image()
self._update_text()
else:
self._advance_index(1)
def _on_mouse_move(self, event):
self._last_mouse_pos = event.pos
def _on_mouse_press(self, event):
self._last_mouse_pos = event.pos
def _on_mouse_release(self, event):
self._last_mouse_pos = event.pos