| from typing import Callable, List, Dict |
| import time |
| import datetime |
|
|
| from qtpy.QtWidgets import QDialog, QLabel, QHBoxLayout, QVBoxLayout, QMessageBox, QSizePolicy, QProgressBar |
| from qtpy.QtGui import QCloseEvent, QShowEvent |
| from qtpy.QtCore import Qt, Signal |
|
|
| from utils.shared import remove_from_runtime_widget_set, add_to_runtime_widget_set |
| from .widget import Widget |
|
|
|
|
| class MessageBox(QMessageBox): |
|
|
| def __init__(self, info_msg: str = None, btn_type = QMessageBox.StandardButton.Ok, frame_less: bool = False, modal: bool = False, signal_slot_map_list: List[Dict] = None, *args, **kwargs): |
| super().__init__(text=info_msg, *args, **kwargs) |
| self.register_signal_slot_map = [] |
| add_to_runtime_widget_set(self) |
|
|
| if frame_less: |
| self.setWindowFlags(Qt.WindowType.FramelessWindowHint) |
| if modal: |
| self.setModal(modal) |
| if btn_type is not None: |
| self.setStandardButtons(btn_type) |
|
|
| if signal_slot_map_list is not None: |
| self.connect_signals(signal_slot_map_list) |
|
|
| def connect_signals(self, signal_slot_map_list): |
| if signal_slot_map_list is None: |
| return |
| if isinstance(signal_slot_map_list, dict): |
| signal_slot_map_list = [signal_slot_map_list] |
| for signal_slot_map in signal_slot_map_list: |
| slot = signal_slot_map['slot'] |
| if isinstance(slot, Callable): |
| slot_func = slot |
| else: |
| assert isinstance(slot, str) |
| slot_func = getattr(self, slot) |
| signal_slot_map['signal'].connect(slot_func) |
| signal_slot_map['slot_func'] = slot_func |
| self.register_signal_slot_map.append(signal_slot_map) |
|
|
| def disconnect_all(self): |
| |
| for signal_slot_map in self.register_signal_slot_map: |
| signal_slot_map['signal'].disconnect(signal_slot_map['slot_func']) |
| self.register_signal_slot_map.clear() |
|
|
| def clear_before_close(self): |
| remove_from_runtime_widget_set(self) |
| self.disconnect_all() |
|
|
| def done(self, v: int = 0): |
| self.clear_before_close() |
| super().done(v) |
|
|
| def closeEvent(self, event: QCloseEvent) -> None: |
| self.clear_before_close() |
| return super().closeEvent(event) |
| |
|
|
| class TaskProgressBar(Widget): |
| def __init__(self, description: str = '', verbose=False, *args, **kwargs) -> None: |
| super().__init__(*args, **kwargs) |
|
|
| self.progressbar = QProgressBar(self) |
| self.progressbar.setTextVisible(False) |
| self.textlabel = QLabel(self) |
| self.description = description |
| self.text_len = 89 |
| layout = QVBoxLayout(self) |
|
|
| self.verbose = verbose |
| |
| |
| if verbose: |
| self.start_time = 0 |
| self.verbose_label = QLabel(self) |
| hl = QHBoxLayout() |
| hl.addWidget(self.textlabel) |
| hl.addStretch(1) |
| hl.addWidget(self.verbose_label) |
| layout.addLayout(hl) |
| else: |
| layout.addWidget(self.textlabel) |
| |
| layout.addWidget(self.progressbar) |
| self.updateProgress(0) |
|
|
| def updateProgress(self, progress: int, msg: str = ''): |
| self.progressbar.setValue(progress) |
| if self.description: |
| msg = self.description + msg |
| if len(msg) > self.text_len - 3: |
| msg = msg[:self.text_len - 3] + '...' |
| elif len(msg) < self.text_len: |
| pads = self.text_len - len(msg) |
| msg = msg + ' ' * pads |
| self.textlabel.setText(msg) |
| self.progressbar.setValue(progress) |
|
|
| if self.verbose: |
| if progress == 0: |
| self.verbose_label.setText('') |
| self.start_time = time.time() |
| elif progress == 100: |
| self.verbose_label.setText('') |
| else: |
| cur_time = time.time() |
| left_progress = 100 - progress |
| eta = left_progress / progress * (cur_time - self.start_time + 1e-6) |
| eta = datetime.timedelta(seconds=int(round(eta))) |
| added_str = f'{progress}% ETA {eta}' |
| self.verbose_label.setText(added_str) |
|
|
|
|
| class FrameLessMessageBox(QMessageBox): |
| def __init__(self, *args, **kwargs) -> None: |
| super().__init__(*args, **kwargs) |
| self.setWindowFlags(Qt.WindowType.FramelessWindowHint) |
| |
|
|
| class ProgressMessageBox(QDialog): |
| showed = Signal() |
| def __init__(self, task_name: str = None, *args, **kwargs) -> None: |
| super().__init__(*args, **kwargs) |
| self.setWindowFlags(Qt.WindowType.FramelessWindowHint) |
| self.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding) |
| self.setModal(True) |
|
|
| layout = QVBoxLayout(self) |
| layout.setSpacing(0) |
| layout.setContentsMargins(20, 10, 20, 30) |
|
|
| self.task_progress_bar: TaskProgressBar = None |
| if task_name is not None: |
| self.task_progress_bar = TaskProgressBar(task_name) |
| layout.addWidget(self.task_progress_bar) |
|
|
| def updateTaskProgress(self, value: int, msg: str = ''): |
| if self.task_progress_bar is not None: |
| self.task_progress_bar.updateProgress(value, msg) |
|
|
| def setTaskName(self, task_name: str): |
| if self.task_progress_bar is not None: |
| self.task_progress_bar.description = task_name |
|
|
| def showEvent(self, e: QShowEvent) -> None: |
| self.showed.emit() |
| return super().showEvent(e) |
|
|
|
|
| class ImgtransProgressMessageBox(ProgressMessageBox): |
| def __init__(self, *args, **kwargs) -> None: |
| super().__init__(None, *args, **kwargs) |
| |
| self.detect_bar = TaskProgressBar(self.tr('Detecting: '), True, self) |
| self.ocr_bar = TaskProgressBar(self.tr('OCR: '), True, self) |
| self.inpaint_bar = TaskProgressBar(self.tr('Inpainting: '), True, self) |
| self.translate_bar = TaskProgressBar(self.tr('Translating: '), True, self) |
|
|
| layout = self.layout() |
| layout.addWidget(self.detect_bar) |
| layout.addWidget(self.ocr_bar) |
| layout.addWidget(self.inpaint_bar) |
| layout.addWidget(self.translate_bar) |
|
|
| self.setFixedWidth(self.sizeHint().width()) |
|
|
| def updateDetectProgress(self, value: int, msg: str = ''): |
| self.detect_bar.updateProgress(value, msg) |
|
|
| def updateOCRProgress(self, value: int, msg: str = ''): |
| self.ocr_bar.updateProgress(value, msg) |
|
|
| def updateInpaintProgress(self, value: int, msg: str = ''): |
| self.inpaint_bar.updateProgress(value, msg) |
|
|
| def updateTranslateProgress(self, value: int, msg: str = ''): |
| self.translate_bar.updateProgress(value, msg) |
| |
| def zero_progress(self): |
| self.updateDetectProgress(0) |
| self.updateOCRProgress(0) |
| self.updateInpaintProgress(0) |
| self.updateTranslateProgress(0) |
|
|
| def show_all_bars(self): |
| self.detect_bar.show() |
| self.ocr_bar.show() |
| self.translate_bar.show() |
| self.inpaint_bar.show() |
|
|
| def hide_all_bars(self): |
| self.detect_bar.hide() |
| self.ocr_bar.hide() |
| self.translate_bar.hide() |
| self.inpaint_bar.hide() |
|
|
