1from subprocess import check_call, CalledProcessError
2import logging
3import tempfile
5from pyrocko import util
6from pyrocko.guts import Object, String, Float, Bytes, clone, \
7 dump_all, load_all
9from pyrocko.gui.qt_compat import qw, qc, qg, get_em
10from .state import ViewerState, Interpolator, interpolateables
11from vtk.util.numpy_support import vtk_to_numpy
12import vtk
13from . import common
15guts_prefix = 'sparrow'
17logger = logging.getLogger('pyrocko.gui.sparrow.snapshots')
19thumb_size = 128, 72
22def to_rect(r):
23 return [float(x) for x in (r.left(), r.top(), r.width(), r.height())]
26def fit_to_rect(frame, size, halign='center', valign='center'):
27 fl, ft, fw, fh = to_rect(frame)
28 rw, rh = size.width(), size.height()
30 ft += 1
31 fh -= 1
33 fl += 1
34 fw -= 1
36 fa = fh / fw
37 ra = rh / rw
39 if fa <= ra:
40 rh = fh
41 rw = rh / ra
42 if halign == 'left':
43 rl = fl
44 elif halign == 'center':
45 rl = fl + 0.5 * (fw - rw)
46 elif halign == 'right':
47 rl = fl + fw - rw
49 rt = ft
50 else:
51 rw = fw
52 rh = rw * ra
53 rl = fl
54 if valign == 'top':
55 rt = ft
56 elif valign == 'center':
57 rt = ft + 0.5 * (fh - rh)
58 elif valign == 'bottom':
59 rt = ft + fh - rh
61 return qc.QRectF(rl, rt, rw, rh)
64def getitem_or_none(items, i):
65 try:
66 return items[i]
67 except IndexError:
68 return None
71def iround(f):
72 return int(round(f))
75class SnapshotItemDelegate(qw.QStyledItemDelegate):
76 def __init__(self, model, parent):
77 qw.QStyledItemDelegate.__init__(self, parent=parent)
78 self.model = model
80 def sizeHint(self, option, index):
81 item = self.model.get_item_or_none(index)
82 if isinstance(item, Snapshot):
83 return qc.QSize(*thumb_size)
84 else:
85 return qw.QStyledItemDelegate.sizeHint(self, option, index)
87 def paint(self, painter, option, index):
88 app = common.get_app()
89 item = self.model.get_item_or_none(index)
90 em = get_em(painter)
91 frect = option.rect.adjusted(0, 0, 0, 0)
92 nb = iround(em*0.5)
93 trect = option.rect.adjusted(nb, nb, -nb, -nb)
95 if isinstance(item, Snapshot):
97 old_pen = painter.pen()
98 if option.state & qw.QStyle.State_Selected:
99 bg_brush = app.palette().brush(
100 qg.QPalette.Active, qg.QPalette.Highlight)
102 fg_pen = qg.QPen(app.palette().color(
103 qg.QPalette.Active, qg.QPalette.HighlightedText))
105 painter.fillRect(frect, bg_brush)
106 painter.setPen(fg_pen)
108 else:
109 bg_brush = app.palette().brush(
110 qg.QPalette.Active, qg.QPalette.AlternateBase)
112 painter.fillRect(frect, bg_brush)
114 # painter.drawRect(frect)
115 img = item.get_image()
116 if img is not None:
117 prect = fit_to_rect(frect, img.size(), halign='right')
118 painter.drawImage(prect, img)
120 painter.drawText(
121 trect,
122 qc.Qt.AlignLeft | qc.Qt.AlignTop,
123 item.name)
125 painter.setPen(
126 app.palette().brush(
127 qg.QPalette.Disabled
128 if item.duration is None
129 else qg.QPalette.Active,
130 qg.QPalette.Text).color())
132 ed = item.effective_duration
133 painter.drawText(
134 trect,
135 qc.Qt.AlignLeft | qc.Qt.AlignBottom,
136 '%.2f s' % ed if ed != 0.0 else '')
138 painter.setPen(old_pen)
140 else:
141 qw.QStyledItemDelegate.paint(self, painter, option, index)
143 # painter.drawText(
144 # trect,
145 # qc.Qt.AlignRight | qc.Qt.AlignTop,
146 # '%.2f' % item.effective_duration)
148 def editorEvent(self, event, model, option, index):
150 item = self.model.get_item_or_none(index)
152 if isinstance(event, qg.QMouseEvent) \
153 and event.button() == qc.Qt.RightButton:
155 menu = qw.QMenu()
157 for name, duration in [
158 ('Auto', None),
159 ('0 s', 0.0),
160 ('1/2 s', 0.5),
161 ('1 s', 1.0),
162 ('3 s', 3.0),
163 ('5 s', 5.0),
164 ('10 s', 10.0),
165 ('60 s', 60.0)]:
167 def make_triggered(duration):
168 def triggered():
169 item.duration = duration
171 return triggered
173 action = qw.QAction(name, menu)
174 action.triggered.connect(make_triggered(duration))
175 menu.addAction(action)
177 action = qw.QAction('Custom...', menu)
179 def triggered():
180 self.parent().edit(index)
182 action.triggered.connect(triggered)
184 menu.addAction(action)
185 menu.exec_(event.globalPos())
187 return True
189 else:
190 return qw.QStyledItemDelegate.editorEvent(
191 self, event, model, option, index)
193 def createEditor(self, parent, option, index):
194 return qw.QLineEdit(parent=parent)
196 def setModelData(self, editor, model, index):
197 item = self.model.get_item_or_none(index)
198 if item:
199 try:
200 item.duration = max(float(editor.text()), 0.0)
201 except ValueError:
202 item.duration = None
204 def setEditorData(self, editor, index):
205 item = self.model.get_item_or_none(index)
206 if item:
207 editor.setText(
208 'Auto' if item.duration is None else '%g' % item.duration)
211class SnapshotListView(qw.QListView):
213 def startDrag(self, supported):
214 if supported & (qc.Qt.CopyAction | qc.Qt.MoveAction):
215 drag = qg.QDrag(self)
216 selected_indexes = self.selectedIndexes()
217 mime_data = self.model().mimeData(selected_indexes)
218 drag.setMimeData(mime_data)
219 drag.exec(qc.Qt.MoveAction)
221 def dropEvent(self, *args):
222 mod = self.model()
223 selected_items = [
224 mod.get_item_or_none(index) for index in self.selectedIndexes()]
226 selected_items = [item for item in selected_items if item is not None]
228 result = qw.QListView.dropEvent(self, *args)
230 indexes = [mod.get_index_for_item(item) for item in selected_items]
232 smod = self.selectionModel()
233 smod.clear()
234 scroll_index = None
235 for index in indexes:
236 if index is not None:
237 smod.select(index, qc.QItemSelectionModel.Select)
238 if scroll_index is None:
239 scroll_index = index
241 if scroll_index is not None:
242 self.scrollTo(scroll_index)
244 return result
247class SnapshotsPanel(qw.QFrame):
249 def __init__(self, viewer):
250 qw.QFrame.__init__(self)
251 layout = qw.QGridLayout()
252 self.setLayout(layout)
254 self.model = SnapshotsModel()
256 self.viewer = viewer
258 lv = SnapshotListView()
259 lv.sizePolicy().setVerticalPolicy(qw.QSizePolicy.Expanding)
260 lv.setModel(self.model)
261 lv.doubleClicked.connect(self.goto_snapshot)
262 lv.setSelectionMode(qw.QAbstractItemView.ExtendedSelection)
263 lv.setDragDropMode(qw.QAbstractItemView.InternalMove)
264 lv.setEditTriggers(qw.QAbstractItemView.NoEditTriggers)
265 lv.viewport().setAcceptDrops(True)
266 self.item_delegate = SnapshotItemDelegate(self.model, lv)
267 lv.setItemDelegate(self.item_delegate)
268 self.list_view = lv
269 layout.addWidget(lv, 0, 0, 1, 3)
271 pb = qw.QPushButton('New')
272 pb.clicked.connect(self.take_snapshot)
273 layout.addWidget(pb, 1, 0, 1, 1)
275 pb = qw.QPushButton('Replace')
276 pb.clicked.connect(self.replace_snapshot)
277 layout.addWidget(pb, 1, 1, 1, 1)
279 pb = qw.QPushButton('Delete')
280 pb.clicked.connect(self.delete_snapshots)
281 layout.addWidget(pb, 1, 2, 1, 1)
283 self.window_to_image_filter = None
285 def setup_menu(self, menu):
286 menu.addAction(
287 'New',
288 self.take_snapshot,
289 qg.QKeySequence(qc.Qt.CTRL | qc.Qt.Key_N)).setShortcutContext(
290 qc.Qt.ApplicationShortcut)
292 menu.addSeparator()
294 menu.addAction(
295 'Next',
296 self.transition_to_next_snapshot,
297 qg.QKeySequence(qc.Qt.Key_PageDown)).setShortcutContext(
298 qc.Qt.ApplicationShortcut)
300 menu.addAction(
301 'Previous',
302 self.transition_to_previous_snapshot,
303 qg.QKeySequence(qc.Qt.Key_PageUp)).setShortcutContext(
304 qc.Qt.ApplicationShortcut)
306 menu.addSeparator()
308 menu.addAction(
309 'Import...',
310 self.import_snapshots)
312 menu.addAction(
313 'Export...',
314 self.export_snapshots)
316 menu.addAction(
317 'Animate',
318 self.animate_snapshots)
320 menu.addAction(
321 'Export Movie...',
322 self.render_movie)
324 menu.addSeparator()
326 menu.addAction(
327 'Show Panel',
328 self.show_and_raise)
330 def show_and_raise(self):
331 self.viewer.raise_panel(self)
333 def get_snapshot_image(self):
334 if not self.window_to_image_filter:
335 wif = vtk.vtkWindowToImageFilter()
336 wif.SetInput(self.viewer.renwin)
337 wif.SetInputBufferTypeToRGBA()
338 wif.ReadFrontBufferOff()
339 self.window_to_image_filter = wif
341 writer = vtk.vtkPNGWriter()
342 writer.SetInputConnection(wif.GetOutputPort())
343 writer.SetWriteToMemory(True)
344 self.png_writer = writer
346 self.viewer.renwin.Render()
347 self.window_to_image_filter.Modified()
348 self.png_writer.Write()
349 data = vtk_to_numpy(self.png_writer.GetResult()).tobytes()
350 img = qg.QImage()
351 img.loadFromData(data)
352 return img
354 def get_snapshot_thumbnail(self):
355 return self.get_snapshot_image().scaled(
356 thumb_size[0], thumb_size[1],
357 qc.Qt.KeepAspectRatio, qc.Qt.SmoothTransformation)
359 def get_snapshot_thumbnail_png(self):
360 img = self.get_snapshot_thumbnail()
362 ba = qc.QByteArray()
363 buf = qc.QBuffer(ba)
364 buf.open(qc.QIODevice.WriteOnly)
365 img.save(buf, format='PNG')
366 return ba.data()
368 def take_snapshot(self):
369 self.model.add_snapshot(
370 Snapshot(
371 state=clone(self.viewer.state),
372 thumb=self.get_snapshot_thumbnail_png()))
373 self.viewer.raise_panel(self)
375 def replace_snapshot(self):
376 state = clone(self.viewer.state)
377 selected_indexes = self.list_view.selectedIndexes()
379 if len(selected_indexes) == 1:
380 self.model.replace_snapshot(
381 selected_indexes[0],
382 Snapshot(
383 state,
384 thumb=self.get_snapshot_thumbnail_png()))
386 self.list_view.update()
388 def goto_snapshot(self, index):
389 item = self.model.get_item_or_none(index)
390 if isinstance(item, Snapshot):
391 self.viewer.set_state(item.state)
392 elif isinstance(item, Transition):
393 snap1 = self.model.get_item_or_none(index.row()-1)
394 snap2 = self.model.get_item_or_none(index.row()+1)
395 if isinstance(snap1, Snapshot) and isinstance(snap2, Snapshot):
396 ip = Interpolator(
397 [0.0, item.effective_duration],
398 [snap1.state, snap2.state])
400 self.viewer.start_animation(ip)
402 def transition_to_next_snapshot(self, direction=1):
403 index = self.list_view.currentIndex()
404 if index.row() == -1:
405 if direction == 1:
406 index = self.model.createIndex(0, 0)
408 item = self.model.get_item_or_none(index)
409 if item is None:
410 return
412 if isinstance(item, Snapshot):
413 snap1 = item
414 transition = self.model.get_item_or_none(index.row()+1*direction)
415 snap2 = self.model.get_item_or_none(index.row()+2*direction)
416 elif isinstance(item, Transition):
417 snap1 = self.model.get_item_or_none(index.row()-1*direction)
418 transition = item
419 snap2 = self.model.get_item_or_none(index.row()+1*direction)
421 if None not in (snap1, transition, snap2):
422 ip = Interpolator(
423 [0.0, transition.effective_duration],
424 [snap1.state, snap2.state])
426 index = self.model.get_index_for_item(snap2)
427 self.list_view.setCurrentIndex(index)
429 self.viewer.start_animation(ip)
431 elif snap2 is not None:
432 index = self.model.get_index_for_item(snap2)
433 self.list_view.setCurrentIndex(index)
434 self.viewer.set_state(snap2.state)
436 def transition_to_previous_snapshot(self):
437 self.transition_to_next_snapshot(-1)
439 def delete_snapshots(self):
440 selected_indexes = self.list_view.selectedIndexes()
441 self.model.remove_snapshots(selected_indexes)
443 def animate_snapshots(self, **kwargs):
444 selected_indexes = self.list_view.selectedIndexes()
445 items = self.model.get_series(selected_indexes)
447 time_state = []
448 item_previous = None
449 t = 0.0
450 for i, item in enumerate(items):
451 item_next = getitem_or_none(items, i+1)
452 item_previous = getitem_or_none(items, i-1)
454 if isinstance(item, Snapshot):
455 time_state.append((t, item.state))
456 if item.effective_duration > 0:
457 time_state.append((t+item.effective_duration, item.state))
459 t += item.effective_duration
461 elif isinstance(item, Transition):
462 if None not in (item_previous, item_next) \
463 and item.effective_duration != 0.0:
465 t += item.effective_duration
467 item_previous = item
469 if len(time_state) < 2:
470 return
472 ip = Interpolator(*zip(*time_state))
474 self.viewer.start_animation(
475 ip, output_path=kwargs.get('output_path', None))
477 def render_movie(self):
478 try:
479 check_call(['ffmpeg', '-loglevel', 'panic'])
480 except CalledProcessError:
481 pass
482 except (TypeError, FileNotFoundError):
483 logger.warn(
484 'Package ffmpeg needed for movie rendering. Please install it '
485 '(e.g. on linux distr. via sudo apt-get ffmpeg.) and retry.')
486 return
488 caption = 'Export Movie'
489 fn_out, _ = qw.QFileDialog.getSaveFileName(
490 self, caption, 'movie.mp4',
491 options=common.qfiledialog_options)
493 if fn_out:
494 self.animate_snapshots(output_path=fn_out)
496 def export_snapshots(self):
497 caption = 'Export Snapshots'
498 fn, _ = qw.QFileDialog.getSaveFileName(
499 self, caption, options=common.qfiledialog_options)
501 selected_indexes = self.list_view.selectedIndexes()
502 items = self.model.get_series(selected_indexes)
504 if fn:
505 dump_all(items, filename=fn)
507 def add_snapshots(self, snapshots):
508 self.model.append_series(snapshots)
510 def load_snapshots(self, path):
511 items = load_snapshots(path)
512 self.add_snapshots(items)
514 def import_snapshots(self):
515 caption = 'Import Snapshots'
516 path, _ = qw.QFileDialog.getOpenFileName(
517 self, caption, options=common.qfiledialog_options)
519 if path:
520 self.load_snapshots(path)
523class Item(Object):
524 duration = Float.T(optional=True)
526 def __init__(self, **kwargs):
527 Object.__init__(self, **kwargs)
528 self.auto_duration = 0.0
530 @property
531 def effective_duration(self):
532 if self.duration is not None:
533 return self.duration
534 else:
535 return self.auto_duration
538class Snapshot(Item):
539 name = String.T()
540 state = ViewerState.T()
541 thumb = Bytes.T(optional=True)
543 isnapshot = 0
545 def __init__(self, state, name=None, thumb=None, **kwargs):
547 if name is None:
548 Snapshot.isnapshot += 1
549 name = '%i' % Snapshot.isnapshot
551 Item.__init__(self, state=state, name=name, thumb=thumb, **kwargs)
552 self._img = None
554 def get_name(self):
555 return self.name
557 def get_image(self):
558 if self.thumb is not None and not self._img:
559 img = qg.QImage()
560 img.loadFromData(self.thumb)
561 self._img = img
563 return self._img
566class Transition(Item):
568 def __init__(self, **kwargs):
569 Item.__init__(self, **kwargs)
570 self.animate = []
572 def get_name(self):
573 ed = self.effective_duration
574 return '%s %s' % (
575 'T' if self.animate and self.effective_duration > 0.0 else '',
576 '%.2f s' % ed if ed != 0.0 else '')
578 @property
579 def name(self):
580 return self.get_name()
583class SnapshotsModel(qc.QAbstractListModel):
585 def __init__(self):
586 qc.QAbstractListModel.__init__(self)
587 self._items = []
589 def supportedDropActions(self):
590 return qc.Qt.MoveAction
592 def rowCount(self, parent=None):
593 return len(self._items)
595 def insertRows(self, index):
596 pass
598 def mimeTypes(self):
599 return ['text/plain']
601 def mimeData(self, indices):
602 objects = [self._items[i.row()] for i in indices]
603 serialized = dump_all(objects)
604 md = qc.QMimeData()
605 md.setText(serialized)
606 md._item_objects = objects
607 return md
609 def dropMimeData(self, md, action, row, col, index):
610 i = index.row()
611 items = getattr(md, '_item_objects', [])
612 self.beginInsertRows(qc.QModelIndex(), i, i)
613 self._items[i:i] = items
614 self.endInsertRows()
615 n = len(items)
616 joff = 0
617 for j in range(len(self._items)):
618 if (j < i or j >= i+n) and self._items[j+joff] in items:
619 self.beginRemoveRows(qc.QModelIndex(), j+joff, j+joff)
620 self._items[j+joff:j+joff+1] = []
621 self.endRemoveRows()
622 joff -= 1
624 self.repair_transitions()
625 return True
627 def removeRows(self, i, n, parent):
628 return True
630 def flags(self, index):
631 if index.isValid():
632 i = index.row()
633 if isinstance(self._items[i], Snapshot):
634 return qc.Qt.ItemFlags(
635 qc.Qt.ItemIsSelectable
636 | qc.Qt.ItemIsEnabled
637 | qc.Qt.ItemIsDragEnabled
638 | qc.Qt.ItemIsEditable)
640 else:
641 return qc.Qt.ItemFlags(
642 qc.Qt.ItemIsEnabled
643 | qc.Qt.ItemIsEnabled
644 | qc.Qt.ItemIsDropEnabled
645 | qc.Qt.ItemIsEditable)
646 else:
647 return qc.QAbstractListModel.flags(self, index)
649 def data(self, index, role):
650 app = common.get_app()
651 i = index.row()
652 item = self._items[i]
653 is_snap = isinstance(item, Snapshot)
654 if role == qc.Qt.DisplayRole:
655 if is_snap:
656 return qc.QVariant(str(item.get_name()))
657 else:
658 return qc.QVariant(str(item.get_name()))
660 elif role == qc.Qt.ToolTipRole:
661 if is_snap:
662 # return qc.QVariant(str(item.state))
663 return qc.QVariant()
664 else:
665 if item.animate:
666 label = 'Interpolation: %s' % \
667 ', '.join(x[0] for x in item.animate)
668 else:
669 label = 'Not interpolable.'
671 return qc.QVariant(label)
673 elif role == qc.Qt.TextAlignmentRole and not is_snap:
674 return qc.QVariant(qc.Qt.AlignRight)
676 elif role == qc.Qt.ForegroundRole and not is_snap:
677 if item.duration is None:
678 return qc.QVariant(app.palette().brush(
679 qg.QPalette.Disabled, qg.QPalette.Text))
680 else:
681 return qc.QVariant(app.palette().brush(
682 qg.QPalette.Active, qg.QPalette.Text))
684 else:
685 qc.QVariant()
687 def headerData(self):
688 pass
690 def add_snapshot(self, snapshot):
691 self.beginInsertRows(
692 qc.QModelIndex(), self.rowCount(), self.rowCount())
693 self._items.append(snapshot)
694 self.endInsertRows()
695 self.repair_transitions()
697 def replace_snapshot(self, index, snapshot):
698 self._items[index.row()] = snapshot
699 self.dataChanged.emit(index, index)
700 self.repair_transitions()
702 def remove_snapshots(self, indexes):
703 indexes = sorted(indexes, key=lambda index: index.row())
704 ioff = 0
705 for index in indexes:
706 i = index.row()
707 self.beginRemoveRows(qc.QModelIndex(), i+ioff, i+ioff)
708 self._items[i+ioff:i+ioff+1] = []
709 self.endRemoveRows()
710 ioff -= 1
712 self.repair_transitions()
714 def repair_transitions(self):
715 items = self._items
716 i = 0
717 need = 0
718 while i < len(items):
719 if need == 0:
720 if not isinstance(items[i], Transition):
721 self.beginInsertRows(qc.QModelIndex(), i, i)
722 items[i:i] = [Transition()]
723 self.endInsertRows()
724 else:
725 i += 1
726 need = 1
727 elif need == 1:
728 if not isinstance(items[i], Snapshot):
729 self.beginRemoveRows(qc.QModelIndex(), i, i)
730 items[i:i+1] = []
731 self.endRemoveRows()
732 else:
733 i += 1
734 need = 0
736 if len(items) == 1:
737 self.beginRemoveRows(qc.QModelIndex(), 0, 0)
738 items[:] = []
739 self.endRemoveRows()
741 elif len(items) > 1:
742 if not isinstance(items[-1], Transition):
743 self.beginInsertRows(
744 qc.QModelIndex(), self.rowCount(), self.rowCount())
745 items.append(Transition())
746 self.endInsertRows()
748 self.update_auto_durations()
750 def update_auto_durations(self):
751 items = self._items
752 for i, item in enumerate(items):
753 if isinstance(item, Transition):
754 if 0 < i < len(items)-1:
755 item.animate = interpolateables(
756 items[i-1].state, items[i+1].state)
758 if item.animate:
759 item.auto_duration = 1.
760 else:
761 item.auto_duration = 0.
763 for i, item in enumerate(items):
764 if isinstance(item, Snapshot):
765 if 0 < i < len(items)-1:
766 if items[i-1].effective_duration == 0 \
767 and items[i+1].effective_duration == 0:
768 item.auto_duration = 1.
769 else:
770 item.auto_duration = 0.
772 def get_index_for_item(self, item):
773 for i, candidate in enumerate(self._items):
774 if candidate is item:
775 return self.createIndex(i, 0)
777 return None
779 def get_item_or_none(self, index):
780 if not isinstance(index, int):
781 i = index.row()
782 else:
783 i = index
785 if i < 0 or len(self._items) <= i:
786 return None
788 try:
789 return self._items[i]
790 except IndexError:
791 return None
793 def get_series(self, indexes):
794 items = self._items
796 ilist = sorted([index.row() for index in indexes])
797 if len(ilist) <= 1:
798 ilist = list(range(0, len(self._items)))
800 ilist = [i for i in ilist if isinstance(items[i], Snapshot)]
801 if len(ilist) == 0:
802 return []
804 i = ilist[0]
806 series = []
807 while ilist:
808 i = ilist.pop(0)
809 series.append(items[i])
810 if ilist and ilist[0] == i+2:
811 series.append(items[i+1])
813 return series
815 def append_series(self, items):
816 self.beginInsertRows(
817 qc.QModelIndex(),
818 self.rowCount(), self.rowCount() + len(items) - 1)
820 self._items.extend(items)
821 self.endInsertRows()
823 self.repair_transitions()
826def load_snapshots(path):
828 if path.startswith('http://') or path.startswith('https://'):
829 with tempfile.NamedTemporaryFile() as fp:
830 util.download_file(path, fp.name)
831 return load_snapshots(fp.name)
833 items = load_all(filename=path)
834 for i in range(len(items)):
835 if not isinstance(
836 items[i], (ViewerState, Snapshot, Transition)):
838 logger.warn(
839 'Only Snapshot, Transition and ViewerState objects '
840 'are accepted. Object #%i from file %s ignored.'
841 % (i, path))
843 if isinstance(items[i], ViewerState):
844 items[i] = Snapshot(items[i])
846 for item in items:
847 if isinstance(item, Snapshot):
848 item.state.sort_elements()
850 return items