1from subprocess import check_call, CalledProcessError
2import logging
5from pyrocko.guts import Object, String, Float, Bytes, clone, \
6 dump_all, load_all
8from pyrocko.gui.qt_compat import qw, qc, qg, get_em
9from .state import ViewerState, Interpolator, interpolateables
10from vtk.util.numpy_support import vtk_to_numpy
11import vtk
12from . import common
14guts_prefix = 'sparrow'
16logger = logging.getLogger('pyrocko.gui.sparrow.snapshots')
18thumb_size = 128, 72
21def to_rect(r):
22 return [float(x) for x in (r.left(), r.top(), r.width(), r.height())]
25def fit_to_rect(frame, size, halign='center', valign='center'):
26 fl, ft, fw, fh = to_rect(frame)
27 rw, rh = size.width(), size.height()
29 ft += 1
30 fh -= 1
32 fl += 1
33 fw -= 1
35 fa = fh / fw
36 ra = rh / rw
38 if fa <= ra:
39 rh = fh
40 rw = rh / ra
41 if halign == 'left':
42 rl = fl
43 elif halign == 'center':
44 rl = fl + 0.5 * (fw - rw)
45 elif halign == 'right':
46 rl = fl + fw - rw
48 rt = ft
49 else:
50 rw = fw
51 rh = rw * ra
52 rl = fl
53 if valign == 'top':
54 rt = ft
55 elif valign == 'center':
56 rt = ft + 0.5 * (fh - rh)
57 elif valign == 'bottom':
58 rt = ft + fh - rh
60 return qc.QRectF(rl, rt, rw, rh)
63def getitem_or_none(items, i):
64 try:
65 return items[i]
66 except IndexError:
67 return None
70def iround(f):
71 return int(round(f))
74class SnapshotItemDelegate(qw.QStyledItemDelegate):
75 def __init__(self, model, parent):
76 qw.QStyledItemDelegate.__init__(self, parent=parent)
77 self.model = model
79 def sizeHint(self, option, index):
80 item = self.model.get_item_or_none(index)
81 if isinstance(item, Snapshot):
82 return qc.QSize(*thumb_size)
83 else:
84 return qw.QStyledItemDelegate.sizeHint(self, option, index)
86 def paint(self, painter, option, index):
87 app = common.get_app()
88 item = self.model.get_item_or_none(index)
89 em = get_em(painter)
90 frect = option.rect.adjusted(0, 0, 0, 0)
91 nb = iround(em*0.5)
92 trect = option.rect.adjusted(nb, nb, -nb, -nb)
94 if isinstance(item, Snapshot):
96 old_pen = painter.pen()
97 if option.state & qw.QStyle.State_Selected:
98 bg_brush = app.palette().brush(
99 qg.QPalette.Active, qg.QPalette.Highlight)
101 fg_pen = qg.QPen(app.palette().color(
102 qg.QPalette.Active, qg.QPalette.HighlightedText))
104 painter.fillRect(frect, bg_brush)
105 painter.setPen(fg_pen)
107 else:
108 bg_brush = app.palette().brush(
109 qg.QPalette.Active, qg.QPalette.AlternateBase)
111 painter.fillRect(frect, bg_brush)
113 # painter.drawRect(frect)
114 img = item.get_image()
115 if img is not None:
116 prect = fit_to_rect(frect, img.size(), halign='right')
117 painter.drawImage(prect, img)
119 painter.drawText(
120 trect,
121 qc.Qt.AlignLeft | qc.Qt.AlignTop,
122 item.name)
124 painter.setPen(
125 app.palette().brush(
126 qg.QPalette.Disabled
127 if item.duration is None
128 else qg.QPalette.Active,
129 qg.QPalette.Text).color())
131 ed = item.effective_duration
132 painter.drawText(
133 trect,
134 qc.Qt.AlignLeft | qc.Qt.AlignBottom,
135 '%.2f s' % ed if ed != 0.0 else '')
137 painter.setPen(old_pen)
139 else:
140 qw.QStyledItemDelegate.paint(self, painter, option, index)
142 # painter.drawText(
143 # trect,
144 # qc.Qt.AlignRight | qc.Qt.AlignTop,
145 # '%.2f' % item.effective_duration)
147 def editorEvent(self, event, model, option, index):
149 item = self.model.get_item_or_none(index)
151 if isinstance(event, qg.QMouseEvent) \
152 and event.button() == qc.Qt.RightButton:
154 menu = qw.QMenu()
156 for name, duration in [
157 ('Auto', None),
158 ('0 s', 0.0),
159 ('1/2 s', 0.5),
160 ('1 s', 1.0),
161 ('3 s', 3.0),
162 ('5 s', 5.0),
163 ('10 s', 10.0),
164 ('60 s', 60.0)]:
166 def make_triggered(duration):
167 def triggered():
168 item.duration = duration
170 return triggered
172 action = qw.QAction(name, menu)
173 action.triggered.connect(make_triggered(duration))
174 menu.addAction(action)
176 action = qw.QAction('Custom...', menu)
178 def triggered():
179 self.parent().edit(index)
181 action.triggered.connect(triggered)
183 menu.addAction(action)
184 menu.exec_(event.globalPos())
186 return True
188 else:
189 return qw.QStyledItemDelegate.editorEvent(
190 self, event, model, option, index)
192 def createEditor(self, parent, option, index):
193 return qw.QLineEdit(parent=parent)
195 def setModelData(self, editor, model, index):
196 item = self.model.get_item_or_none(index)
197 if item:
198 try:
199 item.duration = max(float(editor.text()), 0.0)
200 except ValueError:
201 item.duration = None
203 def setEditorData(self, editor, index):
204 item = self.model.get_item_or_none(index)
205 if item:
206 editor.setText(
207 'Auto' if item.duration is None else '%g' % item.duration)
210class SnapshotListView(qw.QListView):
212 def startDrag(self, supported):
213 if supported & (qc.Qt.CopyAction | qc.Qt.MoveAction):
214 drag = qg.QDrag(self)
215 selected_indexes = self.selectedIndexes()
216 mime_data = self.model().mimeData(selected_indexes)
217 drag.setMimeData(mime_data)
218 drag.exec(qc.Qt.MoveAction)
220 def dropEvent(self, *args):
221 mod = self.model()
222 selected_items = [
223 mod.get_item_or_none(index) for index in self.selectedIndexes()]
225 selected_items = [item for item in selected_items if item is not None]
227 result = qw.QListView.dropEvent(self, *args)
229 indexes = [mod.get_index_for_item(item) for item in selected_items]
231 smod = self.selectionModel()
232 smod.clear()
233 scroll_index = None
234 for index in indexes:
235 if index is not None:
236 smod.select(index, qc.QItemSelectionModel.Select)
237 if scroll_index is None:
238 scroll_index = index
240 if scroll_index is not None:
241 self.scrollTo(scroll_index)
243 return result
246class SnapshotsPanel(qw.QFrame):
248 def __init__(self, viewer):
249 qw.QFrame.__init__(self)
250 layout = qw.QGridLayout()
251 self.setLayout(layout)
253 self.model = SnapshotsModel()
255 self.viewer = viewer
257 lv = SnapshotListView()
258 lv.sizePolicy().setVerticalPolicy(qw.QSizePolicy.Expanding)
259 lv.setModel(self.model)
260 lv.doubleClicked.connect(self.goto_snapshot)
261 lv.setSelectionMode(qw.QAbstractItemView.ExtendedSelection)
262 lv.setDragDropMode(qw.QAbstractItemView.InternalMove)
263 lv.setEditTriggers(qw.QAbstractItemView.NoEditTriggers)
264 lv.viewport().setAcceptDrops(True)
265 self.item_delegate = SnapshotItemDelegate(self.model, lv)
266 lv.setItemDelegate(self.item_delegate)
267 self.list_view = lv
268 layout.addWidget(lv, 0, 0, 1, 3)
270 pb = qw.QPushButton('New')
271 pb.clicked.connect(self.take_snapshot)
272 layout.addWidget(pb, 1, 0, 1, 1)
274 pb = qw.QPushButton('Replace')
275 pb.clicked.connect(self.replace_snapshot)
276 layout.addWidget(pb, 1, 1, 1, 1)
278 pb = qw.QPushButton('Delete')
279 pb.clicked.connect(self.delete_snapshots)
280 layout.addWidget(pb, 1, 2, 1, 1)
282 self.window_to_image_filter = None
284 def setup_menu(self, menu):
285 menu.addAction(
286 'New',
287 self.take_snapshot,
288 qg.QKeySequence(qc.Qt.CTRL | qc.Qt.Key_N)).setShortcutContext(
289 qc.Qt.ApplicationShortcut)
291 menu.addSeparator()
293 menu.addAction(
294 'Next',
295 self.transition_to_next_snapshot,
296 qg.QKeySequence(qc.Qt.Key_PageDown)).setShortcutContext(
297 qc.Qt.ApplicationShortcut)
299 menu.addAction(
300 'Previous',
301 self.transition_to_previous_snapshot,
302 qg.QKeySequence(qc.Qt.Key_PageUp)).setShortcutContext(
303 qc.Qt.ApplicationShortcut)
305 menu.addSeparator()
307 menu.addAction(
308 'Import...',
309 self.import_snapshots)
311 menu.addAction(
312 'Export...',
313 self.export_snapshots)
315 menu.addAction(
316 'Animate',
317 self.animate_snapshots)
319 menu.addAction(
320 'Export Movie...',
321 self.render_movie)
323 menu.addSeparator()
325 menu.addAction(
326 'Show Panel',
327 self.show_and_raise)
329 def show_and_raise(self):
330 self.viewer.raise_panel(self)
332 def get_snapshot_image(self):
333 if not self.window_to_image_filter:
334 wif = vtk.vtkWindowToImageFilter()
335 wif.SetInput(self.viewer.renwin)
336 wif.SetInputBufferTypeToRGBA()
337 wif.ReadFrontBufferOff()
338 self.window_to_image_filter = wif
340 writer = vtk.vtkPNGWriter()
341 writer.SetInputConnection(wif.GetOutputPort())
342 writer.SetWriteToMemory(True)
343 self.png_writer = writer
345 self.viewer.renwin.Render()
346 self.window_to_image_filter.Modified()
347 self.png_writer.Write()
348 data = vtk_to_numpy(self.png_writer.GetResult()).tobytes()
349 img = qg.QImage()
350 img.loadFromData(data)
351 return img
353 def get_snapshot_thumbnail(self):
354 return self.get_snapshot_image().scaled(
355 thumb_size[0], thumb_size[1],
356 qc.Qt.KeepAspectRatio, qc.Qt.SmoothTransformation)
358 def get_snapshot_thumbnail_png(self):
359 img = self.get_snapshot_thumbnail()
361 ba = qc.QByteArray()
362 buf = qc.QBuffer(ba)
363 buf.open(qc.QIODevice.WriteOnly)
364 img.save(buf, format='PNG')
365 return ba.data()
367 def take_snapshot(self):
368 self.model.add_snapshot(
369 Snapshot(
370 state=clone(self.viewer.state),
371 thumb=self.get_snapshot_thumbnail_png()))
372 self.viewer.raise_panel(self)
374 def replace_snapshot(self):
375 state = clone(self.viewer.state)
376 selected_indexes = self.list_view.selectedIndexes()
378 if len(selected_indexes) == 1:
379 self.model.replace_snapshot(
380 selected_indexes[0],
381 Snapshot(
382 state,
383 thumb=self.get_snapshot_thumbnail_png()))
385 self.list_view.update()
387 def goto_snapshot(self, index):
388 item = self.model.get_item_or_none(index)
389 if isinstance(item, Snapshot):
390 self.viewer.set_state(item.state)
391 elif isinstance(item, Transition):
392 snap1 = self.model.get_item_or_none(index.row()-1)
393 snap2 = self.model.get_item_or_none(index.row()+1)
394 if isinstance(snap1, Snapshot) and isinstance(snap2, Snapshot):
395 ip = Interpolator(
396 [0.0, item.effective_duration],
397 [snap1.state, snap2.state])
399 self.viewer.start_animation(ip)
401 def transition_to_next_snapshot(self, direction=1):
402 index = self.list_view.currentIndex()
403 if index.row() == -1:
404 if direction == 1:
405 index = self.model.createIndex(0, 0)
407 item = self.model.get_item_or_none(index)
408 if item is None:
409 return
411 if isinstance(item, Snapshot):
412 snap1 = item
413 transition = self.model.get_item_or_none(index.row()+1*direction)
414 snap2 = self.model.get_item_or_none(index.row()+2*direction)
415 elif isinstance(item, Transition):
416 snap1 = self.model.get_item_or_none(index.row()-1*direction)
417 transition = item
418 snap2 = self.model.get_item_or_none(index.row()+1*direction)
420 if None not in (snap1, transition, snap2):
421 ip = Interpolator(
422 [0.0, transition.effective_duration],
423 [snap1.state, snap2.state])
425 index = self.model.get_index_for_item(snap2)
426 self.list_view.setCurrentIndex(index)
428 self.viewer.start_animation(ip)
430 elif snap2 is not None:
431 index = self.model.get_index_for_item(snap2)
432 self.list_view.setCurrentIndex(index)
433 self.viewer.set_state(snap2.state)
435 def transition_to_previous_snapshot(self):
436 self.transition_to_next_snapshot(-1)
438 def delete_snapshots(self):
439 selected_indexes = self.list_view.selectedIndexes()
440 self.model.remove_snapshots(selected_indexes)
442 def animate_snapshots(self, **kwargs):
443 selected_indexes = self.list_view.selectedIndexes()
444 items = self.model.get_series(selected_indexes)
446 time_state = []
447 item_previous = None
448 t = 0.0
449 for i, item in enumerate(items):
450 item_next = getitem_or_none(items, i+1)
451 item_previous = getitem_or_none(items, i-1)
453 if isinstance(item, Snapshot):
454 time_state.append((t, item.state))
455 if item.effective_duration > 0:
456 time_state.append((t+item.effective_duration, item.state))
458 t += item.effective_duration
460 elif isinstance(item, Transition):
461 if None not in (item_previous, item_next) \
462 and item.effective_duration != 0.0:
464 t += item.effective_duration
466 item_previous = item
468 if len(time_state) < 2:
469 return
471 ip = Interpolator(*zip(*time_state))
473 self.viewer.start_animation(
474 ip, output_path=kwargs.get('output_path', None))
476 def render_movie(self):
477 try:
478 check_call(['ffmpeg', '-loglevel', 'panic'])
479 except CalledProcessError:
480 pass
481 except (TypeError, FileNotFoundError):
482 logger.warn(
483 'Package ffmpeg needed for movie rendering. Please install it '
484 '(e.g. on linux distr. via sudo apt-get ffmpeg.) and retry.')
485 return
487 caption = 'Export Movie'
488 fn_out, _ = qw.QFileDialog.getSaveFileName(
489 self, caption, 'movie.mp4',
490 options=common.qfiledialog_options)
492 if fn_out:
493 self.animate_snapshots(output_path=fn_out)
495 def export_snapshots(self):
496 caption = 'Export Snapshots'
497 fn, _ = qw.QFileDialog.getSaveFileName(
498 self, caption, options=common.qfiledialog_options)
500 selected_indexes = self.list_view.selectedIndexes()
501 items = self.model.get_series(selected_indexes)
503 if fn:
504 dump_all(items, filename=fn)
506 def add_snapshots(self, snapshots):
507 self.model.append_series(snapshots)
509 def load_snapshots(self, path):
510 items = load_snapshots(path)
511 self.add_snapshots(items)
513 def import_snapshots(self):
514 caption = 'Import Snapshots'
515 path, _ = qw.QFileDialog.getOpenFileName(
516 self, caption, options=common.qfiledialog_options)
518 if path:
519 self.load_snapshots(path)
522class Item(Object):
523 duration = Float.T(optional=True)
525 def __init__(self, **kwargs):
526 Object.__init__(self, **kwargs)
527 self.auto_duration = 0.0
529 @property
530 def effective_duration(self):
531 if self.duration is not None:
532 return self.duration
533 else:
534 return self.auto_duration
537class Snapshot(Item):
538 name = String.T()
539 state = ViewerState.T()
540 thumb = Bytes.T(optional=True)
542 isnapshot = 0
544 def __init__(self, state, name=None, thumb=None, **kwargs):
546 if name is None:
547 Snapshot.isnapshot += 1
548 name = '%i' % Snapshot.isnapshot
550 Item.__init__(self, state=state, name=name, thumb=thumb, **kwargs)
551 self._img = None
553 def get_name(self):
554 return self.name
556 def get_image(self):
557 if self.thumb is not None and not self._img:
558 img = qg.QImage()
559 img.loadFromData(self.thumb)
560 self._img = img
562 return self._img
565class Transition(Item):
567 def __init__(self, **kwargs):
568 Item.__init__(self, **kwargs)
569 self.animate = []
571 def get_name(self):
572 ed = self.effective_duration
573 return '%s %s' % (
574 'T' if self.animate and self.effective_duration > 0.0 else '',
575 '%.2f s' % ed if ed != 0.0 else '')
577 @property
578 def name(self):
579 return self.get_name()
582class SnapshotsModel(qc.QAbstractListModel):
584 def __init__(self):
585 qc.QAbstractListModel.__init__(self)
586 self._items = []
588 def supportedDropActions(self):
589 return qc.Qt.MoveAction
591 def rowCount(self, parent=None):
592 return len(self._items)
594 def insertRows(self, index):
595 pass
597 def mimeTypes(self):
598 return ['text/plain']
600 def mimeData(self, indices):
601 objects = [self._items[i.row()] for i in indices]
602 serialized = dump_all(objects)
603 md = qc.QMimeData()
604 md.setText(serialized)
605 md._item_objects = objects
606 return md
608 def dropMimeData(self, md, action, row, col, index):
609 i = index.row()
610 items = getattr(md, '_item_objects', [])
611 self.beginInsertRows(qc.QModelIndex(), i, i)
612 self._items[i:i] = items
613 self.endInsertRows()
614 n = len(items)
615 joff = 0
616 for j in range(len(self._items)):
617 if (j < i or j >= i+n) and self._items[j+joff] in items:
618 self.beginRemoveRows(qc.QModelIndex(), j+joff, j+joff)
619 self._items[j+joff:j+joff+1] = []
620 self.endRemoveRows()
621 joff -= 1
623 self.repair_transitions()
624 return True
626 def removeRows(self, i, n, parent):
627 return True
629 def flags(self, index):
630 if index.isValid():
631 i = index.row()
632 if isinstance(self._items[i], Snapshot):
633 return qc.Qt.ItemFlags(
634 qc.Qt.ItemIsSelectable
635 | qc.Qt.ItemIsEnabled
636 | qc.Qt.ItemIsDragEnabled
637 | qc.Qt.ItemIsEditable)
639 else:
640 return qc.Qt.ItemFlags(
641 qc.Qt.ItemIsEnabled
642 | qc.Qt.ItemIsEnabled
643 | qc.Qt.ItemIsDropEnabled
644 | qc.Qt.ItemIsEditable)
645 else:
646 return qc.QAbstractListModel.flags(self, index)
648 def data(self, index, role):
649 app = common.get_app()
650 i = index.row()
651 item = self._items[i]
652 is_snap = isinstance(item, Snapshot)
653 if role == qc.Qt.DisplayRole:
654 if is_snap:
655 return qc.QVariant(str(item.get_name()))
656 else:
657 return qc.QVariant(str(item.get_name()))
659 elif role == qc.Qt.ToolTipRole:
660 if is_snap:
661 # return qc.QVariant(str(item.state))
662 return qc.QVariant()
663 else:
664 if item.animate:
665 label = 'Interpolation: %s' % \
666 ', '.join(x[0] for x in item.animate)
667 else:
668 label = 'Not interpolable.'
670 return qc.QVariant(label)
672 elif role == qc.Qt.TextAlignmentRole and not is_snap:
673 return qc.QVariant(qc.Qt.AlignRight)
675 elif role == qc.Qt.ForegroundRole and not is_snap:
676 if item.duration is None:
677 return qc.QVariant(app.palette().brush(
678 qg.QPalette.Disabled, qg.QPalette.Text))
679 else:
680 return qc.QVariant(app.palette().brush(
681 qg.QPalette.Active, qg.QPalette.Text))
683 else:
684 qc.QVariant()
686 def headerData(self):
687 pass
689 def add_snapshot(self, snapshot):
690 self.beginInsertRows(
691 qc.QModelIndex(), self.rowCount(), self.rowCount())
692 self._items.append(snapshot)
693 self.endInsertRows()
694 self.repair_transitions()
696 def replace_snapshot(self, index, snapshot):
697 self._items[index.row()] = snapshot
698 self.dataChanged.emit(index, index)
699 self.repair_transitions()
701 def remove_snapshots(self, indexes):
702 indexes = sorted(indexes, key=lambda index: index.row())
703 ioff = 0
704 for index in indexes:
705 i = index.row()
706 self.beginRemoveRows(qc.QModelIndex(), i+ioff, i+ioff)
707 self._items[i+ioff:i+ioff+1] = []
708 self.endRemoveRows()
709 ioff -= 1
711 self.repair_transitions()
713 def repair_transitions(self):
714 items = self._items
715 i = 0
716 need = 0
717 while i < len(items):
718 if need == 0:
719 if not isinstance(items[i], Transition):
720 self.beginInsertRows(qc.QModelIndex(), i, i)
721 items[i:i] = [Transition()]
722 self.endInsertRows()
723 else:
724 i += 1
725 need = 1
726 elif need == 1:
727 if not isinstance(items[i], Snapshot):
728 self.beginRemoveRows(qc.QModelIndex(), i, i)
729 items[i:i+1] = []
730 self.endRemoveRows()
731 else:
732 i += 1
733 need = 0
735 if len(items) == 1:
736 self.beginRemoveRows(qc.QModelIndex(), 0, 0)
737 items[:] = []
738 self.endRemoveRows()
740 elif len(items) > 1:
741 if not isinstance(items[-1], Transition):
742 self.beginInsertRows(
743 qc.QModelIndex(), self.rowCount(), self.rowCount())
744 items.append(Transition())
745 self.endInsertRows()
747 self.update_auto_durations()
749 def update_auto_durations(self):
750 items = self._items
751 for i, item in enumerate(items):
752 if isinstance(item, Transition):
753 if 0 < i < len(items)-1:
754 item.animate = interpolateables(
755 items[i-1].state, items[i+1].state)
757 if item.animate:
758 item.auto_duration = 1.
759 else:
760 item.auto_duration = 0.
762 for i, item in enumerate(items):
763 if isinstance(item, Snapshot):
764 if 0 < i < len(items)-1:
765 if items[i-1].effective_duration == 0 \
766 and items[i+1].effective_duration == 0:
767 item.auto_duration = 1.
768 else:
769 item.auto_duration = 0.
771 def get_index_for_item(self, item):
772 for i, candidate in enumerate(self._items):
773 if candidate is item:
774 return self.createIndex(i, 0)
776 return None
778 def get_item_or_none(self, index):
779 if not isinstance(index, int):
780 i = index.row()
781 else:
782 i = index
784 if i < 0 or len(self._items) <= i:
785 return None
787 try:
788 return self._items[i]
789 except IndexError:
790 return None
792 def get_series(self, indexes):
793 items = self._items
795 ilist = sorted([index.row() for index in indexes])
796 if len(ilist) <= 1:
797 ilist = list(range(0, len(self._items)))
799 ilist = [i for i in ilist if isinstance(items[i], Snapshot)]
800 if len(ilist) == 0:
801 return []
803 i = ilist[0]
805 series = []
806 while ilist:
807 i = ilist.pop(0)
808 series.append(items[i])
809 if ilist and ilist[0] == i+2:
810 series.append(items[i+1])
812 return series
814 def append_series(self, items):
815 self.beginInsertRows(
816 qc.QModelIndex(),
817 self.rowCount(), self.rowCount() + len(items) - 1)
819 self._items.extend(items)
820 self.endInsertRows()
822 self.repair_transitions()
825def load_snapshots(path):
826 items = load_all(filename=path)
827 for i in range(len(items)):
828 if not isinstance(
829 items[i], (ViewerState, Snapshot, Transition)):
831 logger.warn(
832 'Only Snapshot, Transition and ViewerState objects '
833 'are accepted. Object #%i from file %s ignored.'
834 % (i, path))
836 if isinstance(items[i], ViewerState):
837 items[i] = Snapshot(items[i])
839 for item in items:
840 if isinstance(item, Snapshot):
841 item.state.sort_elements()
843 return items