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 pb = qw.QPushButton('Import')
283 pb.clicked.connect(self.import_snapshots)
284 layout.addWidget(pb, 2, 0, 1, 1)
286 pb = qw.QPushButton('Export')
287 pb.clicked.connect(self.export_snapshots)
288 layout.addWidget(pb, 2, 1, 1, 1)
290 pb = qw.QPushButton('Animate')
291 pb.clicked.connect(self.animate_snapshots)
292 layout.addWidget(pb, 2, 2, 1, 1)
294 pb = qw.QPushButton('Movie')
295 pb.clicked.connect(self.render_movie)
296 layout.addWidget(pb, 3, 1, 1, 1)
298 self.window_to_image_filter = None
300 def get_snapshot_image(self):
301 if not self.window_to_image_filter:
302 wif = vtk.vtkWindowToImageFilter()
303 wif.SetInput(self.viewer.renwin)
304 wif.SetInputBufferTypeToRGBA()
305 wif.ReadFrontBufferOff()
306 self.window_to_image_filter = wif
308 writer = vtk.vtkPNGWriter()
309 writer.SetInputConnection(wif.GetOutputPort())
310 writer.SetWriteToMemory(True)
311 self.png_writer = writer
313 self.viewer.renwin.Render()
314 self.window_to_image_filter.Modified()
315 self.png_writer.Write()
316 data = vtk_to_numpy(self.png_writer.GetResult()).tobytes()
317 img = qg.QImage()
318 img.loadFromData(data)
319 return img
321 def get_snapshot_thumbnail(self):
322 return self.get_snapshot_image().scaled(
323 thumb_size[0], thumb_size[1],
324 qc.Qt.KeepAspectRatio, qc.Qt.SmoothTransformation)
326 def get_snapshot_thumbnail_png(self):
327 img = self.get_snapshot_thumbnail()
329 ba = qc.QByteArray()
330 buf = qc.QBuffer(ba)
331 buf.open(qc.QIODevice.WriteOnly)
332 img.save(buf, format='PNG')
333 return ba.data()
335 def take_snapshot(self):
336 self.model.add_snapshot(
337 Snapshot(
338 state=clone(self.viewer.state),
339 thumb=self.get_snapshot_thumbnail_png()))
341 def replace_snapshot(self):
342 state = clone(self.viewer.state)
343 selected_indexes = self.list_view.selectedIndexes()
345 if len(selected_indexes) == 1:
346 self.model.replace_snapshot(
347 selected_indexes[0],
348 Snapshot(
349 state,
350 thumb=self.get_snapshot_thumbnail_png()))
352 self.list_view.update()
354 def goto_snapshot(self, index):
355 item = self.model.get_item_or_none(index)
356 if isinstance(item, Snapshot):
357 self.viewer.set_state(item.state)
358 elif isinstance(item, Transition):
359 snap1 = self.model.get_item_or_none(index.row()-1)
360 snap2 = self.model.get_item_or_none(index.row()+1)
361 if isinstance(snap1, Snapshot) and isinstance(snap2, Snapshot):
362 ip = Interpolator(
363 [0.0, item.effective_duration],
364 [snap1.state, snap2.state])
366 self.viewer.start_animation(ip)
368 def transition_to_next_snapshot(self, direction=1):
369 index = self.list_view.currentIndex()
370 if index.row() == -1:
371 if direction == 1:
372 index = self.model.createIndex(0, 0)
374 item = self.model.get_item_or_none(index)
376 if isinstance(item, Snapshot):
377 snap1 = item
378 transition = self.model.get_item_or_none(index.row()+1*direction)
379 snap2 = self.model.get_item_or_none(index.row()+2*direction)
380 elif isinstance(item, Transition):
381 snap1 = self.model.get_item_or_none(index.row()-1*direction)
382 transition = item
383 snap2 = self.model.get_item_or_none(index.row()+1*direction)
385 if None not in (snap1, transition, snap2):
386 ip = Interpolator(
387 [0.0, transition.effective_duration],
388 [snap1.state, snap2.state])
390 index = self.model.get_index_for_item(snap2)
391 self.list_view.setCurrentIndex(index)
393 self.viewer.start_animation(ip)
395 elif snap2 is not None:
396 index = self.model.get_index_for_item(snap2)
397 self.list_view.setCurrentIndex(index)
398 self.viewer.set_state(snap2.state)
400 def transition_to_previous_snapshot(self):
401 self.transition_to_next_snapshot(-1)
403 def delete_snapshots(self):
404 selected_indexes = self.list_view.selectedIndexes()
405 self.model.remove_snapshots(selected_indexes)
407 def animate_snapshots(self, **kwargs):
408 selected_indexes = self.list_view.selectedIndexes()
409 items = self.model.get_series(selected_indexes)
411 time_state = []
412 item_previous = None
413 t = 0.0
414 for i, item in enumerate(items):
415 item_next = getitem_or_none(items, i+1)
416 item_previous = getitem_or_none(items, i-1)
418 if isinstance(item, Snapshot):
419 time_state.append((t, item.state))
420 if item.effective_duration > 0:
421 time_state.append((t+item.effective_duration, item.state))
423 t += item.effective_duration
425 elif isinstance(item, Transition):
426 if None not in (item_previous, item_next) \
427 and item.effective_duration != 0.0:
429 t += item.effective_duration
431 item_previous = item
433 if len(time_state) < 2:
434 return
436 ip = Interpolator(*zip(*time_state))
438 self.viewer.start_animation(
439 ip, output_path=kwargs.get('output_path', None))
441 def render_movie(self):
442 try:
443 check_call(['ffmpeg', '-loglevel', 'panic'])
444 except CalledProcessError:
445 pass
446 except (TypeError, FileNotFoundError):
447 logger.warn(
448 'Package ffmpeg needed for movie rendering. Please install it '
449 '(e.g. on linux distr. via sudo apt-get ffmpeg.) and retry.')
450 return
452 caption = 'Export Movie'
453 fn_out, _ = qw.QFileDialog.getSaveFileName(
454 self, caption, 'movie.mp4',
455 options=common.qfiledialog_options)
457 if fn_out:
458 self.animate_snapshots(output_path=fn_out)
460 def export_snapshots(self):
461 caption = 'Export Snapshots'
462 fn, _ = qw.QFileDialog.getSaveFileName(
463 self, caption, options=common.qfiledialog_options)
465 selected_indexes = self.list_view.selectedIndexes()
466 items = self.model.get_series(selected_indexes)
468 if fn:
469 dump_all(items, filename=fn)
471 def add_snapshots(self, snapshots):
472 self.model.append_series(snapshots)
474 def load_snapshots(self, path):
475 items = load_snapshots(path)
476 self.add_snapshots(items)
478 def import_snapshots(self):
479 caption = 'Import Snapshots'
480 path, _ = qw.QFileDialog.getOpenFileName(
481 self, caption, options=common.qfiledialog_options)
483 if path:
484 self.load_snapshots(path)
487class Item(Object):
488 duration = Float.T(optional=True)
490 def __init__(self, **kwargs):
491 Object.__init__(self, **kwargs)
492 self.auto_duration = 0.0
494 @property
495 def effective_duration(self):
496 if self.duration is not None:
497 return self.duration
498 else:
499 return self.auto_duration
502class Snapshot(Item):
503 name = String.T()
504 state = ViewerState.T()
505 thumb = Bytes.T(optional=True)
507 isnapshot = 0
509 def __init__(self, state, name=None, thumb=None, **kwargs):
511 if name is None:
512 Snapshot.isnapshot += 1
513 name = '%i' % Snapshot.isnapshot
515 Item.__init__(self, state=state, name=name, thumb=thumb, **kwargs)
516 self._img = None
518 def get_name(self):
519 return self.name
521 def get_image(self):
522 if self.thumb is not None and not self._img:
523 img = qg.QImage()
524 img.loadFromData(self.thumb)
525 self._img = img
527 return self._img
530class Transition(Item):
532 def __init__(self, **kwargs):
533 Item.__init__(self, **kwargs)
534 self.animate = []
536 def get_name(self):
537 ed = self.effective_duration
538 return '%s %s' % (
539 'T' if self.animate and self.effective_duration > 0.0 else '',
540 '%.2f s' % ed if ed != 0.0 else '')
542 @property
543 def name(self):
544 return self.get_name()
547class SnapshotsModel(qc.QAbstractListModel):
549 def __init__(self):
550 qc.QAbstractListModel.__init__(self)
551 self._items = []
553 def supportedDropActions(self):
554 return qc.Qt.MoveAction
556 def rowCount(self, parent=None):
557 return len(self._items)
559 def insertRows(self, index):
560 pass
562 def mimeTypes(self):
563 return ['text/plain']
565 def mimeData(self, indices):
566 objects = [self._items[i.row()] for i in indices]
567 serialized = dump_all(objects)
568 md = qc.QMimeData()
569 md.setText(serialized)
570 md._item_objects = objects
571 return md
573 def dropMimeData(self, md, action, row, col, index):
574 i = index.row()
575 items = getattr(md, '_item_objects', [])
576 self.beginInsertRows(qc.QModelIndex(), i, i)
577 self._items[i:i] = items
578 self.endInsertRows()
579 n = len(items)
580 joff = 0
581 for j in range(len(self._items)):
582 if (j < i or j >= i+n) and self._items[j+joff] in items:
583 self.beginRemoveRows(qc.QModelIndex(), j+joff, j+joff)
584 self._items[j+joff:j+joff+1] = []
585 self.endRemoveRows()
586 joff -= 1
588 self.repair_transitions()
589 return True
591 def removeRows(self, i, n, parent):
592 return True
594 def flags(self, index):
595 if index.isValid():
596 i = index.row()
597 if isinstance(self._items[i], Snapshot):
598 return qc.Qt.ItemFlags(
599 qc.Qt.ItemIsSelectable
600 | qc.Qt.ItemIsEnabled
601 | qc.Qt.ItemIsDragEnabled
602 | qc.Qt.ItemIsEditable)
604 else:
605 return qc.Qt.ItemFlags(
606 qc.Qt.ItemIsEnabled
607 | qc.Qt.ItemIsEnabled
608 | qc.Qt.ItemIsDropEnabled
609 | qc.Qt.ItemIsEditable)
610 else:
611 return qc.QAbstractListModel.flags(self, index)
613 def data(self, index, role):
614 app = common.get_app()
615 i = index.row()
616 item = self._items[i]
617 is_snap = isinstance(item, Snapshot)
618 if role == qc.Qt.DisplayRole:
619 if is_snap:
620 return qc.QVariant(str(item.get_name()))
621 else:
622 return qc.QVariant(str(item.get_name()))
624 elif role == qc.Qt.ToolTipRole:
625 if is_snap:
626 # return qc.QVariant(str(item.state))
627 return qc.QVariant()
628 else:
629 if item.animate:
630 label = 'Interpolation: %s' % \
631 ', '.join(x[0] for x in item.animate)
632 else:
633 label = 'Not interpolable.'
635 return qc.QVariant(label)
637 elif role == qc.Qt.TextAlignmentRole and not is_snap:
638 return qc.QVariant(qc.Qt.AlignRight)
640 elif role == qc.Qt.ForegroundRole and not is_snap:
641 if item.duration is None:
642 return qc.QVariant(app.palette().brush(
643 qg.QPalette.Disabled, qg.QPalette.Text))
644 else:
645 return qc.QVariant(app.palette().brush(
646 qg.QPalette.Active, qg.QPalette.Text))
648 else:
649 qc.QVariant()
651 def headerData(self):
652 pass
654 def add_snapshot(self, snapshot):
655 self.beginInsertRows(
656 qc.QModelIndex(), self.rowCount(), self.rowCount())
657 self._items.append(snapshot)
658 self.endInsertRows()
659 self.repair_transitions()
661 def replace_snapshot(self, index, snapshot):
662 self._items[index.row()] = snapshot
663 self.dataChanged.emit(index, index)
664 self.repair_transitions()
666 def remove_snapshots(self, indexes):
667 indexes = sorted(indexes, key=lambda index: index.row())
668 ioff = 0
669 for index in indexes:
670 i = index.row()
671 self.beginRemoveRows(qc.QModelIndex(), i+ioff, i+ioff)
672 self._items[i+ioff:i+ioff+1] = []
673 self.endRemoveRows()
674 ioff -= 1
676 self.repair_transitions()
678 def repair_transitions(self):
679 items = self._items
680 i = 0
681 need = 0
682 while i < len(items):
683 if need == 0:
684 if not isinstance(items[i], Transition):
685 self.beginInsertRows(qc.QModelIndex(), i, i)
686 items[i:i] = [Transition()]
687 self.endInsertRows()
688 else:
689 i += 1
690 need = 1
691 elif need == 1:
692 if not isinstance(items[i], Snapshot):
693 self.beginRemoveRows(qc.QModelIndex(), i, i)
694 items[i:i+1] = []
695 self.endRemoveRows()
696 else:
697 i += 1
698 need = 0
700 if len(items) == 1:
701 self.beginRemoveRows(qc.QModelIndex(), 0, 0)
702 items[:] = []
703 self.endRemoveRows()
705 elif len(items) > 1:
706 if not isinstance(items[-1], Transition):
707 self.beginInsertRows(
708 qc.QModelIndex(), self.rowCount(), self.rowCount())
709 items.append(Transition())
710 self.endInsertRows()
712 self.update_auto_durations()
714 def update_auto_durations(self):
715 items = self._items
716 for i, item in enumerate(items):
717 if isinstance(item, Transition):
718 if 0 < i < len(items)-1:
719 item.animate = interpolateables(
720 items[i-1].state, items[i+1].state)
722 if item.animate:
723 item.auto_duration = 3.
724 else:
725 item.auto_duration = 0.
727 for i, item in enumerate(items):
728 if isinstance(item, Snapshot):
729 if 0 < i < len(items)-1:
730 if items[i-1].effective_duration == 0 \
731 and items[i+1].effective_duration == 0:
732 item.auto_duration = 3.
733 else:
734 item.auto_duration = 0.
736 def get_index_for_item(self, item):
737 for i, candidate in enumerate(self._items):
738 if candidate is item:
739 return self.createIndex(i, 0)
741 return None
743 def get_item_or_none(self, index):
744 if not isinstance(index, int):
745 i = index.row()
746 else:
747 i = index
749 if i < 0 or len(self._items) <= i:
750 return None
752 try:
753 return self._items[i]
754 except IndexError:
755 return None
757 def get_series(self, indexes):
758 items = self._items
760 ilist = sorted([index.row() for index in indexes])
761 if len(ilist) <= 1:
762 ilist = list(range(0, len(self._items)))
764 ilist = [i for i in ilist if isinstance(items[i], Snapshot)]
765 if len(ilist) == 0:
766 return []
768 i = ilist[0]
770 series = []
771 while ilist:
772 i = ilist.pop(0)
773 series.append(items[i])
774 if ilist and ilist[0] == i+2:
775 series.append(items[i+1])
777 return series
779 def append_series(self, items):
780 self.beginInsertRows(
781 qc.QModelIndex(),
782 self.rowCount(), self.rowCount() + len(items) - 1)
784 self._items.extend(items)
785 self.endInsertRows()
787 self.repair_transitions()
790def load_snapshots(path):
791 items = load_all(filename=path)
792 for i in range(len(items)):
793 if not isinstance(
794 items[i], (ViewerState, Snapshot, Transition)):
796 logger.warn(
797 'Only Snapshot, Transition and ViewerState objects '
798 'are accepted. Object #%i from file %s ignored.'
799 % (i, path))
801 if isinstance(items[i], ViewerState):
802 items[i] = Snapshot(items[i])
804 for item in items:
805 if isinstance(item, Snapshot):
806 item.state.sort_elements()
808 return items