diff --git a/spinetoolbox/spine_db_editor/widgets/custom_qgraphicsviews.py b/spinetoolbox/spine_db_editor/widgets/custom_qgraphicsviews.py index e612cc10c..59ebc48ac 100644 --- a/spinetoolbox/spine_db_editor/widgets/custom_qgraphicsviews.py +++ b/spinetoolbox/spine_db_editor/widgets/custom_qgraphicsviews.py @@ -17,6 +17,7 @@ import sys import tempfile from contextlib import contextmanager +import numpy as np from PySide6.QtCore import Qt, QTimeLine, Signal, Slot, QRectF from PySide6.QtWidgets import QMenu, QGraphicsView, QInputDialog, QColorDialog, QMessageBox, QLineEdit from PySide6.QtGui import QCursor, QPainter, QIcon, QAction, QPageSize, QPixmap @@ -26,6 +27,7 @@ from ...widgets.custom_qgraphicsviews import CustomQGraphicsView from ...widgets.custom_qwidgets import ToolBarWidgetAction, HorizontalSpinBox from ..graphics_items import EntityItem, CrossHairsArcItem, BgItem +from .custom_qwidgets import ExportAsVideoDialog from .select_graph_parameters_dialog import SelectGraphParametersDialog @@ -293,57 +295,6 @@ def populate_context_menu(self): self._rebuild_action = self._menu.addAction("Rebuild", self._spine_db_editor.rebuild_graph) self._menu.aboutToShow.connect(self._update_actions_visibility) - def _save_state(self): - name, ok = QInputDialog.getText( - self, "Save state...", "Enter a name for the state.", QLineEdit.Normal, self._current_state_name - ) - if not ok: - return - db_map_graph_data = self._db_map_graph_data_by_name.get(name) - if db_map_graph_data is not None: - button = QMessageBox.question( - self._spine_db_editor, - self._spine_db_editor.windowTitle(), - f"State {name} already exists. Do you want to overwrite it?", - ) - if button == QMessageBox.StandardButton.Yes: - self._spine_db_editor.overwrite_graph_data(db_map_graph_data) - return - self._spine_db_editor.save_graph_data(name) - - @Slot(QAction) - def _load_state(self, action): - self._current_state_name = name = action.text() - db_map_graph_data = self._db_map_graph_data_by_name.get(name) - self._spine_db_editor.load_graph_data(db_map_graph_data) - - @Slot(QAction) - def _remove_state(self, action): - name = action.text() - self._spine_db_editor.remove_graph_data(name) - - def _find(self): - expr, ok = QInputDialog.getText(self, "Find in graph...", "Enter entity names to find separated by comma.") - if not ok: - return - names = [x.strip() for x in expr.split(",")] - items = [item for item in self.entity_items if any(n == item.entity_name for n in names)] - if not items: - return - color = QColorDialog.getColor(Qt.yellow, self, "Choose highlight color") - for item in items: - item.set_highlight_color(color) - - def increase_arc_length(self): - for item in self.entity_items: - new_pos = 1.1 * item.pos() - item.set_pos(new_pos.x(), new_pos.y()) - - def decrease_arc_length(self): - for item in self.entity_items: - new_pos = item.pos() / 1.1 - item.set_pos(new_pos.x(), new_pos.y()) - @Slot() def _update_actions_visibility(self): """Enables or disables actions according to current selection in the graph.""" @@ -398,6 +349,57 @@ def make_items_menu(self): menu.addAction("Remove", self.remove_selected).setEnabled(bool(self.selected_items)) return menu + def _save_state(self): + name, ok = QInputDialog.getText( + self, "Save state...", "Enter a name for the state.", QLineEdit.Normal, self._current_state_name + ) + if not ok: + return + db_map_graph_data = self._db_map_graph_data_by_name.get(name) + if db_map_graph_data is not None: + button = QMessageBox.question( + self._spine_db_editor, + self._spine_db_editor.windowTitle(), + f"State {name} already exists. Do you want to overwrite it?", + ) + if button == QMessageBox.StandardButton.Yes: + self._spine_db_editor.overwrite_graph_data(db_map_graph_data) + return + self._spine_db_editor.save_graph_data(name) + + @Slot(QAction) + def _load_state(self, action): + self._current_state_name = name = action.text() + db_map_graph_data = self._db_map_graph_data_by_name.get(name) + self._spine_db_editor.load_graph_data(db_map_graph_data) + + @Slot(QAction) + def _remove_state(self, action): + name = action.text() + self._spine_db_editor.remove_graph_data(name) + + def _find(self): + expr, ok = QInputDialog.getText(self, "Find in graph...", "Enter entity names to find separated by comma.") + if not ok: + return + names = [x.strip() for x in expr.split(",")] + items = [item for item in self.entity_items if any(n == item.entity_name for n in names)] + if not items: + return + color = QColorDialog.getColor(Qt.yellow, self, "Choose highlight color") + for item in items: + item.set_highlight_color(color) + + def increase_arc_length(self): + for item in self.entity_items: + new_pos = 1.1 * item.pos() + item.set_pos(new_pos.x(), new_pos.y()) + + def decrease_arc_length(self): + for item in self.entity_items: + new_pos = item.pos() / 1.1 + item.set_pos(new_pos.x(), new_pos.y()) + @Slot(bool) def add_entities_at_position(self, checked=False): self._spine_db_editor.add_entities_at_position(self._context_menu_pos) @@ -700,23 +702,36 @@ def _print_scene(self, printer, source, size, index=None): painter.end() @busy_effect - def _pixmaps(self, fps): + def _pixmaps(self, start, stop, frame_count): + if start == stop: + return [] pixmaps = [] with self._no_zoom(): source = self._get_print_source() size = source.size().toSize() - for index in self._spine_db_editor.ui.time_line_widget.indexes(fps): + incr = (stop - start) / frame_count + index = start + while True: pixmap = QPixmap(size) pixmap.fill(Qt.white) self._spine_db_editor._update_time_line_index(index) # FIXME self._print_scene(pixmap, source, size, index=index) pixmaps.append(pixmap.scaledToWidth(1600)) + index += incr + if index > stop: + break return pixmaps @Slot(bool) def export_as_video(self): try: import cv2 + + def pixmap_to_frame(pixmap, file_path): + ok = pixmap.save(file_path) + assert ok + return cv2.imread(file_path, -1) + except ModuleNotFoundError: self._spine_db_editor.msg_error.emit( "Export as video requires opencv-python" @@ -727,20 +742,20 @@ def export_as_video(self): ) if not file_path: return + start, stop = self._spine_db_editor.ui.time_line_widget.get_index_range() + dialog = ExportAsVideoDialog(str(start), str(stop), parent=self) + if dialog.exec_() == ExportAsVideoDialog.Rejected: + return file_ext = os.path.splitext(file_path)[-1].lower() if not file_ext: file_ext = ".mp4" file_path += file_ext - fps = 1.0 - pixmaps = self._pixmaps(fps) + start, stop, frame_count, fps = dialog.selections() + start = np.datetime64(start) + stop = np.datetime64(stop) + pixmaps = self._pixmaps(start, stop, frame_count) if not pixmaps: return - - def pixmap_to_frame(pixmap, file_path): - ok = pixmap.save(file_path) - assert ok - return cv2.imread(file_path, -1) - pixmap_iter = enumerate(pixmaps) with tempfile.NamedTemporaryFile() as f: buffer_path = f.name + ".png" @@ -748,6 +763,7 @@ def pixmap_to_frame(pixmap, file_path): frame = pixmap_to_frame(pixmap, buffer_path) height, width, _layers = frame.shape video = cv2.VideoWriter(file_path, cv2.VideoWriter_fourcc(*"XVID"), fps, (width, height)) + video.write(frame) self._spine_db_editor.file_exported.emit(file_path, k / len(pixmaps), False) for k, pixmap in pixmap_iter: frame = pixmap_to_frame(pixmap, buffer_path) diff --git a/spinetoolbox/spine_db_editor/widgets/custom_qwidgets.py b/spinetoolbox/spine_db_editor/widgets/custom_qwidgets.py index 8217f3471..de95b51bf 100644 --- a/spinetoolbox/spine_db_editor/widgets/custom_qwidgets.py +++ b/spinetoolbox/spine_db_editor/widgets/custom_qwidgets.py @@ -25,10 +25,25 @@ QSlider, QVBoxLayout, QHBoxLayout, + QFormLayout, QSizePolicy, + QDialog, + QDateTimeEdit, + QSpinBox, ) from PySide6.QtGui import QPainter, QColor, QIcon, QBrush, QPainterPath, QPalette -from PySide6.QtCore import Signal, Slot, QVariantAnimation, QPointF, Qt, QTimeLine, QRectF, QTimer, QEasingCurve +from PySide6.QtCore import ( + Signal, + Slot, + QVariantAnimation, + QPointF, + Qt, + QTimeLine, + QRectF, + QTimer, + QEasingCurve, + QDateTime, +) from sqlalchemy.engine.url import URL from ...helpers import open_url, CharIconEngine, color_from_index @@ -104,11 +119,12 @@ def __init__(self, file_path, progress, db_editor): self.set_progress(progress) def set_progress(self, progress): - self.progress = self._progress_bar.minimum() + progress * ( + self.progress = progress + progress_bar_value = self._progress_bar.minimum() + self.progress * ( self._progress_bar.maximum() - self._progress_bar.minimum() ) - self._progress_bar.setValue(self.progress) - if self.progress == self._progress_bar.maximum(): + self._progress_bar.setValue(progress_bar_value) + if progress_bar_value == self._progress_bar.maximum(): def _show_button(): self._progress_bar.hide() @@ -280,19 +296,8 @@ def set_index_range(self, min_index, max_index): self.index_changed.emit(self._index) self.show() - def indexes(self, fps): - indexes = [] - total_seconds = self._ms_per_step * self._STEP_COUNT / 1000 # in s - frame_count = total_seconds * fps - index_range = self._max_index - self._min_index - incr = index_range / frame_count - index = self._min_index - while True: - indexes.append(index) - index += incr - if index > self._max_index: - break - return indexes + def get_index_range(self): + return (self._min_index, self._max_index) class LegendWidget(QWidget): @@ -382,3 +387,50 @@ def paint(self, painter, rect): cell.setWidth(max_val_cw) painter.drawText(cell, text_flags, max_val) painter.restore() + + +class ExportAsVideoDialog(QDialog): + def __init__(self, start, stop, parent=None): + super().__init__(parent=parent) + self.setWindowTitle("Export as video") + layout = QVBoxLayout(self) + form = QFormLayout() + self._start_edit = QDateTimeEdit() + self._stop_edit = QDateTimeEdit() + for dt_edit in (self._start_edit, self._stop_edit): + dt_edit.setMinimumDateTime(QDateTime.fromString(start, Qt.ISODate)) + dt_edit.setMaximumDateTime(QDateTime.fromString(stop, Qt.ISODate)) + dt_edit.setCalendarPopup(True) + self._start_edit.dateTimeChanged.connect(self._handle_start_dt_changed) + self._stop_edit.dateTimeChanged.connect(self._handle_stop_dt_changed) + self._frame_count_spin_box = QSpinBox() + self._frame_count_spin_box.setRange(1, 16777215) + self._fps_spin_box = QSpinBox() + self._fps_spin_box.setRange(1, 20) + form.addRow("Start:", self._start_edit) + form.addRow("Stop:", self._stop_edit) + form.addRow("Number of frames:", self._frame_count_spin_box) + form.addRow("Frames per second:", self._fps_spin_box) + self._button_box = QDialogButtonBox(QDialogButtonBox.StandardButton.Ok | QDialogButtonBox.StandardButton.Cancel) + self._button_box.accepted.connect(self.accept) + self._button_box.rejected.connect(self.reject) + layout.addLayout(form) + layout.addWidget(self._button_box) + + @Slot(QDateTime) + def _handle_start_dt_changed(self, start_dt): + if start_dt > self._stop_edit.dateTime(): + self._stop_edit.setDateTime(start_dt) + + @Slot(QDateTime) + def _handle_stop_dt_changed(self, stop_dt): + if stop_dt < self._start_edit.dateTime(): + self._start_edit.setDateTime(stop_dt) + + def selections(self): + return ( + self._start_edit.dateTime().toString(Qt.ISODate), + self._stop_edit.dateTime().toString(Qt.ISODate), + self._frame_count_spin_box.value(), + self._fps_spin_box.value(), + )