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 delete_snapshots(self):
369 selected_indexes = self.list_view.selectedIndexes()
370 self.model.remove_snapshots(selected_indexes)
372 def animate_snapshots(self, **kwargs):
373 selected_indexes = self.list_view.selectedIndexes()
374 items = self.model.get_series(selected_indexes)
376 time_state = []
377 item_previous = None
378 t = 0.0
379 for i, item in enumerate(items):
380 item_next = getitem_or_none(items, i+1)
381 item_previous = getitem_or_none(items, i-1)
383 if isinstance(item, Snapshot):
384 time_state.append((t, item.state))
385 if item.effective_duration > 0:
386 time_state.append((t+item.effective_duration, item.state))
388 t += item.effective_duration
390 elif isinstance(item, Transition):
391 if None not in (item_previous, item_next) \
392 and item.effective_duration != 0.0:
394 t += item.effective_duration
396 item_previous = item
398 if len(time_state) < 2:
399 return
401 ip = Interpolator(*zip(*time_state))
403 self.viewer.start_animation(
404 ip, output_path=kwargs.get('output_path', None))
406 def render_movie(self):
407 try:
408 check_call(['ffmpeg', '-loglevel', 'panic'])
409 except CalledProcessError:
410 pass
411 except (TypeError, FileNotFoundError):
412 logger.warn(
413 'Package ffmpeg needed for movie rendering. Please install it '
414 '(e.g. on linux distr. via sudo apt-get ffmpeg.) and retry.')
415 return
417 caption = 'Export Movie'
418 fn_out, _ = qw.QFileDialog.getSaveFileName(
419 self, caption, 'movie.mp4',
420 options=common.qfiledialog_options)
422 if fn_out:
423 self.animate_snapshots(output_path=fn_out)
425 def export_snapshots(self):
426 caption = 'Export Snapshots'
427 fn, _ = qw.QFileDialog.getSaveFileName(
428 self, caption, options=common.qfiledialog_options)
430 selected_indexes = self.list_view.selectedIndexes()
431 items = self.model.get_series(selected_indexes)
433 if fn:
434 dump_all(items, filename=fn)
436 def add_snapshots(self, snapshots):
437 self.model.append_series(snapshots)
439 def load_snapshots(self, path):
440 items = load_snapshots(path)
441 self.add_snapshots(items)
443 def import_snapshots(self):
444 caption = 'Import Snapshots'
445 path, _ = qw.QFileDialog.getOpenFileName(
446 self, caption, options=common.qfiledialog_options)
448 if path:
449 self.load_snapshots(path)
452class Item(Object):
453 duration = Float.T(optional=True)
455 def __init__(self, **kwargs):
456 Object.__init__(self, **kwargs)
457 self.auto_duration = 0.0
459 @property
460 def effective_duration(self):
461 if self.duration is not None:
462 return self.duration
463 else:
464 return self.auto_duration
467class Snapshot(Item):
468 name = String.T()
469 state = ViewerState.T()
470 thumb = Bytes.T(optional=True)
472 isnapshot = 0
474 def __init__(self, state, name=None, thumb=None, **kwargs):
476 if name is None:
477 Snapshot.isnapshot += 1
478 name = '%i' % Snapshot.isnapshot
480 Item.__init__(self, state=state, name=name, thumb=thumb, **kwargs)
481 self._img = None
483 def get_name(self):
484 return self.name
486 def get_image(self):
487 if self.thumb is not None and not self._img:
488 img = qg.QImage()
489 img.loadFromData(self.thumb)
490 self._img = img
492 return self._img
495class Transition(Item):
497 def __init__(self, **kwargs):
498 Item.__init__(self, **kwargs)
499 self.animate = []
501 def get_name(self):
502 ed = self.effective_duration
503 return '%s %s' % (
504 'T' if self.animate and self.effective_duration > 0.0 else '',
505 '%.2f s' % ed if ed != 0.0 else '')
507 @property
508 def name(self):
509 return self.get_name()
512class SnapshotsModel(qc.QAbstractListModel):
514 def __init__(self):
515 qc.QAbstractListModel.__init__(self)
516 self._items = []
518 def supportedDropActions(self):
519 return qc.Qt.MoveAction
521 def rowCount(self, parent=None):
522 return len(self._items)
524 def insertRows(self, index):
525 pass
527 def mimeTypes(self):
528 return ['text/plain']
530 def mimeData(self, indices):
531 objects = [self._items[i.row()] for i in indices]
532 serialized = dump_all(objects)
533 md = qc.QMimeData()
534 md.setText(serialized)
535 md._item_objects = objects
536 return md
538 def dropMimeData(self, md, action, row, col, index):
539 i = index.row()
540 items = getattr(md, '_item_objects', [])
541 self.beginInsertRows(qc.QModelIndex(), i, i)
542 self._items[i:i] = items
543 self.endInsertRows()
544 n = len(items)
545 joff = 0
546 for j in range(len(self._items)):
547 if (j < i or j >= i+n) and self._items[j+joff] in items:
548 self.beginRemoveRows(qc.QModelIndex(), j+joff, j+joff)
549 self._items[j+joff:j+joff+1] = []
550 self.endRemoveRows()
551 joff -= 1
553 self.repair_transitions()
554 return True
556 def removeRows(self, i, n, parent):
557 return True
559 def flags(self, index):
560 if index.isValid():
561 i = index.row()
562 if isinstance(self._items[i], Snapshot):
563 return qc.Qt.ItemFlags(
564 qc.Qt.ItemIsSelectable
565 | qc.Qt.ItemIsEnabled
566 | qc.Qt.ItemIsDragEnabled
567 | qc.Qt.ItemIsEditable)
569 else:
570 return qc.Qt.ItemFlags(
571 qc.Qt.ItemIsEnabled
572 | qc.Qt.ItemIsEnabled
573 | qc.Qt.ItemIsDropEnabled
574 | qc.Qt.ItemIsEditable)
575 else:
576 return qc.QAbstractListModel.flags(self, index)
578 def data(self, index, role):
579 app = common.get_app()
580 i = index.row()
581 item = self._items[i]
582 is_snap = isinstance(item, Snapshot)
583 if role == qc.Qt.DisplayRole:
584 if is_snap:
585 return qc.QVariant(str(item.get_name()))
586 else:
587 return qc.QVariant(str(item.get_name()))
589 elif role == qc.Qt.ToolTipRole:
590 if is_snap:
591 return qc.QVariant(str(item.state))
592 else:
593 if item.animate:
594 label = 'Interpolation: %s' % \
595 ', '.join(x[0] for x in item.animate)
596 else:
597 label = 'Not interpolable.'
599 return qc.QVariant(label)
601 elif role == qc.Qt.TextAlignmentRole and not is_snap:
602 return qc.QVariant(qc.Qt.AlignRight)
604 elif role == qc.Qt.ForegroundRole and not is_snap:
605 if item.duration is None:
606 return qc.QVariant(app.palette().brush(
607 qg.QPalette.Disabled, qg.QPalette.Text))
608 else:
609 return qc.QVariant(app.palette().brush(
610 qg.QPalette.Active, qg.QPalette.Text))
612 else:
613 qc.QVariant()
615 def headerData(self):
616 pass
618 def add_snapshot(self, snapshot):
619 self.beginInsertRows(
620 qc.QModelIndex(), self.rowCount(), self.rowCount())
621 self._items.append(snapshot)
622 self.endInsertRows()
623 self.repair_transitions()
625 def replace_snapshot(self, index, snapshot):
626 self._items[index.row()] = snapshot
627 self.dataChanged.emit(index, index)
628 self.repair_transitions()
630 def remove_snapshots(self, indexes):
631 indexes = sorted(indexes, key=lambda index: index.row())
632 ioff = 0
633 for index in indexes:
634 i = index.row()
635 self.beginRemoveRows(qc.QModelIndex(), i+ioff, i+ioff)
636 self._items[i+ioff:i+ioff+1] = []
637 self.endRemoveRows()
638 ioff -= 1
640 self.repair_transitions()
642 def repair_transitions(self):
643 items = self._items
644 i = 0
645 need = 0
646 while i < len(items):
647 if need == 0:
648 if not isinstance(items[i], Transition):
649 self.beginInsertRows(qc.QModelIndex(), i, i)
650 items[i:i] = [Transition()]
651 self.endInsertRows()
652 else:
653 i += 1
654 need = 1
655 elif need == 1:
656 if not isinstance(items[i], Snapshot):
657 self.beginRemoveRows(qc.QModelIndex(), i, i)
658 items[i:i+1] = []
659 self.endRemoveRows()
660 else:
661 i += 1
662 need = 0
664 if len(items) == 1:
665 self.beginRemoveRows(qc.QModelIndex(), 0, 0)
666 items[:] = []
667 self.endRemoveRows()
669 elif len(items) > 1:
670 if not isinstance(items[-1], Transition):
671 self.beginInsertRows(
672 qc.QModelIndex(), self.rowCount(), self.rowCount())
673 items.append(Transition())
674 self.endInsertRows()
676 self.update_auto_durations()
678 def update_auto_durations(self):
679 items = self._items
680 for i, item in enumerate(items):
681 if isinstance(item, Transition):
682 if 0 < i < len(items)-1:
683 item.animate = interpolateables(
684 items[i-1].state, items[i+1].state)
686 if item.animate:
687 item.auto_duration = 3.
688 else:
689 item.auto_duration = 0.
691 for i, item in enumerate(items):
692 if isinstance(item, Snapshot):
693 if 0 < i < len(items)-1:
694 if items[i-1].effective_duration == 0 \
695 and items[i+1].effective_duration == 0:
696 item.auto_duration = 3.
697 else:
698 item.auto_duration = 0.
700 def get_index_for_item(self, item):
701 for i, candidate in enumerate(self._items):
702 if candidate is item:
703 return self.createIndex(i, 0)
705 return None
707 def get_item_or_none(self, index):
708 if not isinstance(index, int):
709 i = index.row()
710 else:
711 i = index
713 try:
714 return self._items[i]
715 except IndexError:
716 return None
718 def get_series(self, indexes):
719 items = self._items
721 ilist = sorted([index.row() for index in indexes])
722 if len(ilist) <= 1:
723 ilist = list(range(0, len(self._items)))
725 ilist = [i for i in ilist if isinstance(items[i], Snapshot)]
726 if len(ilist) == 0:
727 return []
729 i = ilist[0]
731 series = []
732 while ilist:
733 i = ilist.pop(0)
734 series.append(items[i])
735 if ilist and ilist[0] == i+2:
736 series.append(items[i+1])
738 return series
740 def append_series(self, items):
741 self.beginInsertRows(
742 qc.QModelIndex(),
743 self.rowCount(), self.rowCount() + len(items) - 1)
745 self._items.extend(items)
746 self.endInsertRows()
748 self.repair_transitions()
751def load_snapshots(path):
752 items = load_all(filename=path)
753 for i in range(len(items)):
754 if not isinstance(
755 items[i], (ViewerState, Snapshot, Transition)):
757 logger.warn(
758 'Only Snapshot, Transition and ViewerState objects '
759 'are accepted. Object #%i from file %s ignored.'
760 % (i, path))
762 if isinstance(items[i], ViewerState):
763 items[i] = Snapshot(items[i])
765 for item in items:
766 if isinstance(item, Snapshot):
767 item.state.sort_elements()
769 return items