00001
00002
00003
00004
00005
00006
00007
00008
00009
00010
00011
00012
00013
00014
00015
00016
00017
00018
00019
00020
00021
00022
00023
00024
00025
00026
00027
00028
00029
00030
00031
00032
00033 PKG = 'rxbag_plugins'
00034 import roslib; roslib.load_manifest(PKG)
00035
00036 import os
00037 import shutil
00038 import sys
00039 import threading
00040 import time
00041
00042 import numpy
00043 import wx
00044 import wx.lib.masked
00045 import wx.lib.wxcairo
00046
00047 import Image
00048
00049 from rxbag import bag_helper, TopicMessageView
00050 import image_helper
00051
00052 class ImageView(TopicMessageView):
00053 name = 'Image'
00054
00055 def __init__(self, timeline, parent):
00056 TopicMessageView.__init__(self, timeline, parent)
00057
00058 self._image_lock = threading.RLock()
00059 self._image = None
00060 self._image_topic = None
00061 self._image_stamp = None
00062
00063 self._image_surface = None
00064
00065 self._overlay_font_size = 14.0
00066 self._overlay_indent = (4, 4)
00067 self._overlay_color = (0.2, 0.2, 1.0)
00068
00069 self._size_set = False
00070
00071 self._next_frame_num = 1
00072
00073 tb = self.parent.GetToolBar()
00074 icons_dir = roslib.packages.get_pkg_dir(PKG) + '/icons/'
00075 tb.AddSeparator()
00076
00077 self.save_frame_tool = tb.AddLabelTool(wx.ID_ANY, '', wx.Bitmap(icons_dir + 'picture_save.png'), shortHelp='Save frame', longHelp='Save frame using filename spec')
00078 self.file_spec_tool = wx.TextCtrl(tb, wx.ID_ANY, 'frame%04d.png', size=(200, 22))
00079 tb.AddControl(self.file_spec_tool)
00080 self.save_selected_frames_tool = tb.AddLabelTool(wx.ID_ANY, '', wx.Bitmap(icons_dir + 'pictures_save.png'), shortHelp='Save frames', longHelp='Save frames from selected region')
00081
00082 tb.Bind(wx.EVT_TOOL, lambda e: self.save_frame(), self.save_frame_tool)
00083 tb.Bind(wx.EVT_TOOL, lambda e: self.save_selected_frames(), self.save_selected_frames_tool)
00084
00085 self.parent.Bind(wx.EVT_SIZE, self.on_size)
00086 self.parent.Bind(wx.EVT_PAINT, self.on_paint)
00087 self.parent.Bind(wx.EVT_RIGHT_DOWN, self.on_right_down)
00088
00089
00090
00091 def message_viewed(self, bag, msg_details):
00092 TopicMessageView.message_viewed(self, bag, msg_details)
00093
00094 topic, msg, t = msg_details[:3]
00095
00096 if not msg:
00097 self.set_image(None, topic, stamp)
00098 else:
00099 self.set_image(msg, topic, msg.header.stamp)
00100
00101 if not self._size_set:
00102 self._size_set = True
00103 wx.CallAfter(self.reset_size)
00104
00105 def message_cleared(self):
00106 TopicMessageView.message_cleared(self)
00107
00108 self.set_image(None, None, None)
00109
00110
00111
00112 def set_image(self, image_msg, image_topic, image_stamp):
00113 with self._image_lock:
00114 self._image_msg = image_msg
00115 if image_msg:
00116 self._image = image_helper.imgmsg_to_pil(image_msg)
00117 else:
00118 self._image = None
00119 self._image_surface = None
00120
00121 self._image_topic = image_topic
00122 self._image_stamp = image_stamp
00123
00124 wx.CallAfter(self.parent.Refresh)
00125
00126 def reset_size(self):
00127 with self._image_lock:
00128 if self._image:
00129 self.parent.ClientSize = self._image.size
00130
00131
00132
00133 def on_size(self, event):
00134 with self._image_lock:
00135 self._image_surface = None
00136
00137 self.parent.Refresh()
00138
00139 def on_paint(self, event):
00140 dc = wx.lib.wxcairo.ContextFromDC(wx.PaintDC(self.parent))
00141
00142 with self._image_lock:
00143 if not self._image:
00144 return
00145
00146 ix, iy, iw, ih = 0, 0, self.parent.ClientSize[0], self.parent.ClientSize[1]
00147
00148
00149 if not self._image_surface:
00150 if self._image.size[0] != iw or self._image.size[1] != ih:
00151 self._image_surface = image_helper.pil_to_cairo(self._image.resize((iw, ih), Image.NEAREST))
00152 else:
00153 self._image_surface = image_helper.pil_to_cairo(self._image)
00154
00155
00156 dc.set_source_surface(self._image_surface, ix, iy)
00157 dc.rectangle(ix, iy, iw, ih)
00158 dc.fill()
00159
00160
00161 dc.set_font_size(self._overlay_font_size)
00162 font_height = dc.font_extents()[2]
00163 dc.set_source_rgb(*self._overlay_color)
00164 dc.move_to(self._overlay_indent[0], self._overlay_indent[1] + font_height)
00165 dc.show_text(bag_helper.stamp_to_str(self._image_stamp))
00166
00167 def on_right_down(self, event):
00168 if self._image:
00169 self.parent.PopupMenu(ImagePopupMenu(self.parent, self), event.GetPosition())
00170
00171
00172
00173 def save_frame(self):
00174 if not self._image:
00175 return
00176
00177 file_spec = self.file_spec_tool.Value
00178 try:
00179 filename = file_spec % self._next_frame_num
00180 except Exception:
00181 wx.MessageDialog(None, 'Error with filename specification.\n\nPlease include a frame number format, e.g. frame%04d.png', 'Error', wx.OK | wx.ICON_ERROR).ShowModal()
00182 return
00183
00184 with self._image_lock:
00185 if self._image_msg:
00186 _save_image_msg(self._image_msg, filename)
00187 self._next_frame_num += 1
00188
00189 self.parent.SetStatusText('%s written' % filename)
00190
00191 def save_frame_as(self):
00192 if not self._image:
00193 return
00194
00195 dialog = wx.FileDialog(self.parent.Parent, 'Save frame to...', wildcard='PNG files (*.png)|*.png', style=wx.FD_SAVE)
00196 if dialog.ShowModal() == wx.ID_OK:
00197 with self._image_lock:
00198 if self._image_msg:
00199 _save_image_msg(self._image_msg, dialog.Path)
00200 dialog.Destroy()
00201
00202 def save_selected_frames(self):
00203 if not self._image or not self.timeline.has_selected_region:
00204 return
00205
00206 wx.CallAfter(self._show_export_dialog)
00207
00208 def _show_export_dialog(self):
00209 dialog = ExportFramesDialog(None, -1, 'Export frames', self.timeline, self._image_topic)
00210 dialog.Show()
00211
00212 """
00213 def export_video(self):
00214 bag_file = self.timeline.bag_file
00215
00216 msg_positions = bag_index.msg_positions[self._image_topic]
00217 if len(msg_positions) == 0:
00218 return
00219
00220 dialog = wx.FileDialog(self.parent.Parent, 'Save video to...', wildcard='AVI files (*.avi)|*.avi', style=wx.FD_SAVE)
00221 if dialog.ShowModal() != wx.ID_OK:
00222 return
00223 video_filename = dialog.Path
00224
00225 import tempfile
00226 tmpdir = tempfile.mkdtemp()
00227
00228 frame_count = 0
00229 w, h = None, None
00230
00231 total_frames = len(bag_file.read_messages(self._image_topic, raw=True))
00232
00233 for i, (topic, msg, t) in enumerate(bag_file.read_messages(self._image_topic)):
00234 img = image_helper.imgmsg_to_wx(msg)
00235 if img:
00236 frame_filename = '%s/frame-%s.png' % (tmpdir, str(t))
00237 print '[%d / %d]' % (i + 1, total_frames)
00238 img.SaveFile(frame_filename, wx.BITMAP_TYPE_PNG)
00239 frame_count += 1
00240
00241 if w is None:
00242 w, h = img.GetWidth(), img.GetHeight()
00243
00244 if frame_count > 0:
00245 positions = numpy.array([stamp for (stamp, pos) in msg_positions])
00246 if len(positions) > 1:
00247 spacing = positions[1:] - positions[:-1]
00248 fps = 1.0 / numpy.median(spacing)
00249
00250 print 'Encoding %dx%d at %d fps' % (w, h, fps)
00251
00252 try:
00253 command = ('mencoder',
00254 'mf://' + tmpdir + '/*.png',
00255 '-mf',
00256 'type=png:w=%d:h=%d:fps=%d' % (w, h, fps),
00257 '-ovc',
00258 'lavc',
00259 '-lavcopts',
00260 'vcodec=mpeg4',
00261 '-oac',
00262 'copy',
00263 '-o',
00264 video_filename)
00265 os.spawnvp(os.P_WAIT, 'mencoder', command)
00266 finally:
00267 shutil.rmtree(tmpdir)
00268
00269 dialog.Destroy()
00270 """
00271
00272 class ImagePopupMenu(wx.Menu):
00273 def __init__(self, parent, image_view):
00274 wx.Menu.__init__(self)
00275
00276 self.image_view = image_view
00277
00278
00279 reset_item = wx.MenuItem(self, wx.NewId(), 'Reset Size')
00280 self.AppendItem(reset_item)
00281 self.Bind(wx.EVT_MENU, lambda e: self.image_view.reset_size(), id=reset_item.Id)
00282
00283
00284 save_frame_as_item = wx.MenuItem(self, wx.NewId(), 'Save Frame...')
00285 self.AppendItem(save_frame_as_item)
00286 self.Bind(wx.EVT_MENU, lambda e: self.image_view.save_frame_as(), id=save_frame_as_item.Id)
00287
00288
00289 save_selected_frames_item = wx.MenuItem(self, wx.NewId(), 'Save Selected Frames...')
00290 self.AppendItem(save_selected_frames_item)
00291 self.Bind(wx.EVT_MENU, lambda e: self.image_view.save_selected_frames(), id=save_selected_frames_item.Id)
00292 if not self.image_view.timeline.has_selected_region:
00293 save_selected_frames_item.Enable(False)
00294
00295
00296
00297
00298
00299
00300 class ExportFramesDialog(wx.Dialog):
00301 def __init__(self, parent, id, title, timeline, topic):
00302 wx.Dialog.__init__(self, parent, id, title, size=(280, 90))
00303
00304 self.timeline = timeline
00305 self.topic = topic
00306
00307 panel = wx.Panel(self, -1)
00308
00309 self.file_spec_label = wx.StaticText(panel, -1, 'File spec:', (5, 10))
00310 self.file_spec_text = wx.TextCtrl (panel, -1, 'frame%04d.png', (65, 7), (210, 22))
00311
00312 self.steps_label = wx.StaticText(panel, -1, 'Every:', (5, 35))
00313 self.steps_text = wx.lib.masked.NumCtrl(panel, -1, pos=(65, 32), size=(34, 22), value=1, integerWidth=3, allowNegative=False)
00314 self.steps_text.SetMin(1)
00315
00316 self.frames_label = wx.StaticText(panel, -1, 'frame(s)', (102, 36))
00317
00318 self.export_button = wx.Button(self, -1, 'Export', size=(68, 26))
00319 self.cancel_button = wx.Button(self, -1, 'Cancel', size=(68, 26))
00320
00321 hbox = wx.BoxSizer(wx.HORIZONTAL)
00322 hbox.Add(self.export_button, 1)
00323 hbox.Add(self.cancel_button, 1, wx.LEFT, 5)
00324
00325 vbox = wx.BoxSizer(wx.VERTICAL)
00326 vbox.Add(panel)
00327 vbox.Add(hbox, 1, wx.ALIGN_CENTER | wx.TOP | wx.BOTTOM, 5)
00328 self.SetSizer(vbox)
00329
00330 self.Bind(wx.EVT_BUTTON, self._on_export, id=self.export_button.Id)
00331 self.Bind(wx.EVT_BUTTON, self._on_cancel, id=self.cancel_button.Id)
00332
00333 self.progress = None
00334 self.export_thread = None
00335
00336 def _on_export(self, event):
00337 file_spec = self.file_spec_text.Value
00338 try:
00339 test_filename = file_spec % 0
00340 except Exception:
00341 wx.MessageDialog(None, 'Error with filename specification.\n\nPlease include a frame number format, e.g. frame%04d.png', 'Error', wx.OK | wx.ICON_ERROR).ShowModal()
00342 return
00343
00344 if not self.timeline.start_background_task('Exporting frames to "%s"' % file_spec):
00345 return
00346
00347 wx.CallAfter(wx.GetApp().GetTopWindow().StatusBar.gauge.Show)
00348
00349 try:
00350 step = max(1, int(self.steps_text.Value))
00351 except ValueError:
00352 step = 1
00353
00354 bag_entries = list(self.timeline.get_entries_with_bags(self.topic, self.timeline.play_region[0], self.timeline.play_region[1]))[::step]
00355 total_frames = len(bag_entries)
00356
00357 if not self.timeline.background_task_cancel:
00358 self.export_thread = threading.Thread(target=self._run, args=(file_spec, bag_entries))
00359 self.export_thread.setDaemon(True)
00360 self.export_thread.start()
00361
00362 self.Close()
00363
00364 def _run(self, file_spec, bag_entries):
00365 total_frames = len(bag_entries)
00366
00367 progress = 0
00368
00369 def update_progress(v):
00370 wx.GetApp().TopWindow.StatusBar.progress = v
00371
00372 frame_num = 1
00373 for i, (bag, entry) in enumerate(bag_entries):
00374 if self.timeline.background_task_cancel:
00375 break
00376
00377 try:
00378 topic, msg, t = self.timeline.read_message(bag, entry.position)
00379
00380 filename = file_spec % frame_num
00381 try:
00382 _save_image_msg(msg, filename)
00383 frame_num += 1
00384 except Exception, ex:
00385 print >> sys.stderr, 'Error saving frame at %s to disk: %s' % (str(t), str(ex))
00386
00387 new_progress = int(100.0 * (float(frame_num) / total_frames))
00388 if new_progress != progress:
00389 progress = new_progress
00390 wx.CallAfter(update_progress, progress)
00391 except Exception, ex:
00392 print >> sys.stderr, 'Error saving frame %d: %s' % (i, str(ex))
00393
00394 wx.CallAfter(wx.GetApp().TopWindow.StatusBar.gauge.Hide)
00395
00396 wx.CallAfter(self.Destroy)
00397
00398 self.timeline.stop_background_task()
00399
00400 def _on_cancel(self, event):
00401 self.Destroy()
00402
00403 def _save_image_msg(img_msg, filename):
00404 pil_img = image_helper.imgmsg_to_pil(img_msg, rgba=False)
00405
00406 if pil_img.mode in ('RGB', 'RGBA'):
00407 pil_img = pil_img.convert('RGB')
00408 pil_img = image_helper.pil_bgr2rgb(pil_img)
00409 pil_img = pil_img.convert('RGB')
00410
00411 pil_img.save(filename)
