from typing import List, Callable from modules import GET_VALID_INPAINTERS, GET_VALID_TEXTDETECTORS, GET_VALID_TRANSLATORS, GET_VALID_OCR, \ BaseTranslator, DEFAULT_DEVICE, GPUINTENSIVE_SET from utils.logger import logger as LOGGER from .custom_widget import ConfigComboBox, ParamComboBox, NoBorderPushBtn, ParamNameLabel from utils.shared import CONFIG_COMBOBOX_LONG, size2width, CONFIG_COMBOBOX_SHORT, CONFIG_COMBOBOX_HEIGHT from utils.config import pcfg from qtpy.QtWidgets import QPlainTextEdit, QHBoxLayout, QVBoxLayout, QWidget, QLabel, QCheckBox, QLineEdit, QGridLayout, QPushButton from qtpy.QtCore import Qt, Signal from qtpy.QtGui import QDoubleValidator class ParamCheckGroup(QWidget): paramwidget_edited = Signal(str, dict) def __init__(self, param_key, check_group: dict, parent=None) -> None: super().__init__(parent=parent) self.param_key = param_key layout = QHBoxLayout(self) self.label2widget = {} for k, v in check_group.items(): checker = QCheckBox(text=k, parent=self) checker.setChecked(v) layout.addWidget(checker) self.label2widget[k] = checker checker.clicked.connect(self.on_checker_clicked) def on_checker_clicked(self): new_state_dict = {} w = QCheckBox() for k, w in self.label2widget.items(): new_state_dict[k] = w.isChecked() self.paramwidget_edited.emit(self.param_key, new_state_dict) class ParamLineEditor(QLineEdit): paramwidget_edited = Signal(str, str) def __init__(self, param_key: str, force_digital, size='short', *args, **kwargs) -> None: super().__init__( *args, **kwargs) self.param_key = param_key self.setFixedWidth(size2width(size)) self.setFixedHeight(CONFIG_COMBOBOX_HEIGHT) self.textChanged.connect(self.on_text_changed) if force_digital: validator = QDoubleValidator() self.setValidator(validator) def on_text_changed(self): self.paramwidget_edited.emit(self.param_key, self.text()) class ParamEditor(QPlainTextEdit): paramwidget_edited = Signal(str, str) def __init__(self, param_key: str, *args, **kwargs) -> None: super().__init__( *args, **kwargs) self.param_key = param_key if param_key == 'chat sample': self.setFixedWidth(int(CONFIG_COMBOBOX_LONG * 1.2)) self.setFixedHeight(200) else: self.setFixedWidth(CONFIG_COMBOBOX_LONG) self.setFixedHeight(100) # self.setFixedHeight(CONFIG_COMBOBOX_HEIGHT) self.textChanged.connect(self.on_text_changed) def on_text_changed(self): self.paramwidget_edited.emit(self.param_key, self.text()) def setText(self, text: str): self.setPlainText(text) def text(self): return self.toPlainText() class ParamCheckerBox(QWidget): checker_changed = Signal(bool) paramwidget_edited = Signal(str, str) def __init__(self, param_key: str, *args, **kwargs): super().__init__(*args, **kwargs) self.param_key = param_key self.checker = QCheckBox() name_label = ParamNameLabel(param_key) hlayout = QHBoxLayout(self) hlayout.addWidget(name_label) hlayout.addWidget(self.checker) hlayout.setAlignment(Qt.AlignmentFlag.AlignLeft) self.checker.stateChanged.connect(self.on_checker_changed) def on_checker_changed(self): is_checked = self.checker.isChecked() self.checker_changed.emit(is_checked) checked = 'true' if is_checked else 'false' self.paramwidget_edited.emit(self.param_key, checked) class ParamCheckBox(QCheckBox): paramwidget_edited = Signal(str, bool) def __init__(self, param_key: str, *args, **kwargs): super().__init__(*args, **kwargs) self.param_key = param_key self.stateChanged.connect(self.on_checker_changed) def on_checker_changed(self): self.paramwidget_edited.emit(self.param_key, self.isChecked()) def get_param_display_name(param_key: str, param_dict: dict = None): if param_dict is not None and isinstance(param_dict, dict): if 'display_name' in param_dict: return param_dict['display_name'] return param_key class ParamPushButton(QPushButton): paramwidget_edited = Signal(str, str) def __init__(self, param_key: str, param_dict: dict = None, *args, **kwargs): super().__init__(*args, **kwargs) self.param_key = param_key self.setText(get_param_display_name(param_key, param_dict)) self.clicked.connect(self.on_clicked) def on_clicked(self): self.paramwidget_edited.emit(self.param_key, '') class ParamWidget(QWidget): paramwidget_edited = Signal(str, dict) def __init__(self, params, scrollWidget: QWidget = None, *args, **kwargs) -> None: super().__init__(*args, **kwargs) layout = QHBoxLayout(self) self.param_layout = param_layout = QGridLayout() param_layout.setAlignment(Qt.AlignmentFlag.AlignLeft) param_layout.setContentsMargins(0, 0, 0, 0) param_layout.setAlignment(Qt.AlignmentFlag.AlignLeft) layout.addLayout(param_layout) layout.addStretch(-1) if 'description' in params: self.setToolTip(params['description']) for ii, param_key in enumerate(params): if param_key == 'description': continue display_param_name = param_key require_label = True is_str = isinstance(params[param_key], str) is_digital = isinstance(params[param_key], float) or isinstance(params[param_key], int) param_widget = None if isinstance(params[param_key], bool): param_widget = ParamCheckBox(param_key) val = params[param_key] param_widget.setChecked(val) param_widget.paramwidget_edited.connect(self.on_paramwidget_edited) elif is_str or is_digital: param_widget = ParamLineEditor(param_key, force_digital=is_digital) val = params[param_key] if is_digital: val = str(val) param_widget.setText(val) param_widget.paramwidget_edited.connect(self.on_paramwidget_edited) elif isinstance(params[param_key], dict): param_dict = params[param_key] display_param_name = get_param_display_name(param_key, param_dict) value = params[param_key]['value'] param_widget = None # Ensure initialization param_type = param_dict['type'] if 'type' in param_dict else 'line_editor' flush_btn = param_dict.get('flush_btn', False) path_selector = param_dict.get('path_selector', False) param_size = param_dict.get('size', 'short') if param_type == 'selector': if 'url' in param_key: size = size2width('median') else: size = size2width(param_size) param_widget = ParamComboBox( param_key, param_dict['options'], size=size, scrollWidget=scrollWidget, flush_btn=flush_btn, path_selector=path_selector) if param_key == 'device' and DEFAULT_DEVICE == 'cpu': param_dict['value'] = 'cpu' for ii, device in enumerate(param_dict['options']): if device in GPUINTENSIVE_SET: model = param_widget.model() item = model.item(ii, 0) item.setEnabled(False) param_widget.setCurrentText(str(value)) param_widget.setEditable(param_dict.get('editable', False)) elif param_type == 'editor': param_widget = ParamEditor(param_key) param_widget.setText(value) elif param_type == 'checkbox': param_widget = ParamCheckBox(param_key) if isinstance(value, str): value = value.lower().strip() == 'true' params[param_key]['value'] = value param_widget.setChecked(value) elif param_type == 'pushbtn': param_widget = ParamPushButton(param_key, param_dict) require_label = False elif param_type == 'line_editor': param_widget = ParamLineEditor(param_key, force_digital=is_digital) param_widget.setText(str(value)) elif param_type == 'check_group': param_widget = ParamCheckGroup(param_key, check_group=value) if param_widget is not None: param_widget.paramwidget_edited.connect(self.on_paramwidget_edited) if 'description' in param_dict: param_widget.setToolTip(param_dict['description']) widget_idx = 0 if require_label: param_label = ParamNameLabel(display_param_name) param_layout.addWidget(param_label, ii, 0) widget_idx = 1 if param_widget is not None: pw_lo = None if hasattr(param_widget, 'flush_btn') or hasattr(param_widget, 'path_select_btn'): pw_lo = QHBoxLayout() pw_lo.addWidget(param_widget) if hasattr(param_widget, 'flush_btn'): pw_lo.addWidget(param_widget.flush_btn) param_widget.flushbtn_clicked.connect(self.on_flushbtn_clicked) if hasattr(param_widget, 'path_select_btn'): pw_lo.addWidget(param_widget.path_select_btn) param_widget.pathbtn_clicked.connect(self.on_pathbtn_clicked) if pw_lo is None: param_layout.addWidget(param_widget, ii, widget_idx) else: param_layout.addLayout(pw_lo, ii, widget_idx) else: v = params[param_key] raise ValueError(f"Failed to initialize widget for key-value pair: {param_key}-{v}") def on_flushbtn_clicked(self): paramw: ParamComboBox = self.sender() content_dict = {'content': '', 'widget': paramw, 'flush': True} self.paramwidget_edited.emit(paramw.param_key, content_dict) def on_pathbtn_clicked(self): paramw: ParamComboBox = self.sender() content_dict = {'content': '', 'widget': paramw, 'select_path': True} self.paramwidget_edited.emit(paramw.param_key, content_dict) def on_paramwidget_edited(self, param_key, param_content): content_dict = {'content': param_content} self.paramwidget_edited.emit(param_key, content_dict) class ModuleParseWidgets(QWidget): def addModulesParamWidgets(self, ocr_instance): self.params = ocr_instance.get_params() self.on_module_changed() def on_module_changed(self): self.updateModuleParamWidget() def updateModuleParamWidget(self): widget = ParamWidget(self.params, scrollWidget=self) layout = QVBoxLayout() layout.addWidget(widget) self.setLayout(layout) class ModuleConfigParseWidget(QWidget): module_changed = Signal(str) paramwidget_edited = Signal(str, dict) def __init__(self, module_name: str, get_valid_module_keys: Callable, scrollWidget: QWidget, add_from: int = 1, *args, **kwargs) -> None: super().__init__( *args, **kwargs) self.get_valid_module_keys = get_valid_module_keys self.module_combobox = ConfigComboBox(scrollWidget=scrollWidget) self.params_layout = QHBoxLayout() self.params_layout.setContentsMargins(0, 0, 0, 0) p_layout = QHBoxLayout() p_layout.setAlignment(Qt.AlignmentFlag.AlignLeft | Qt.AlignmentFlag.AlignVCenter) self.module_label = ParamNameLabel(module_name) p_layout.addWidget(self.module_label) p_layout.addWidget(self.module_combobox) p_layout.addStretch(-1) self.p_layout = p_layout layout = QVBoxLayout(self) self.param_widget_map = {} layout.addLayout(p_layout) layout.addLayout(self.params_layout) layout.setSpacing(30) self.vlayout = layout self.visibleWidget: QWidget = None self.module_dict: dict = {} def addModulesParamWidgets(self, module_dict: dict): invalid_module_keys = [] valid_modulekeys = self.get_valid_module_keys() num_widgets_before = len(self.param_widget_map) for module in module_dict: if module not in valid_modulekeys: invalid_module_keys.append(module) continue if module in self.param_widget_map: LOGGER.warning(f'duplicated module key: {module}') continue self.module_combobox.addItem(module) params = module_dict[module] if params is not None: self.param_widget_map[module] = None if len(invalid_module_keys) > 0: LOGGER.warning(F'Invalid module keys: {invalid_module_keys}') for ik in invalid_module_keys: module_dict.pop(ik) self.module_dict = module_dict num_widgets_after = len(self.param_widget_map) if num_widgets_before == 0 and num_widgets_after > 0: self.on_module_changed() self.module_combobox.currentTextChanged.connect(self.on_module_changed) def setModule(self, module: str): self.blockSignals(True) self.module_combobox.setCurrentText(module) self.updateModuleParamWidget() self.blockSignals(False) def updateModuleParamWidget(self): module = self.module_combobox.currentText() if self.visibleWidget is not None: self.visibleWidget.hide() if module in self.param_widget_map: widget: QWidget = self.param_widget_map[module] if widget is None: # lazy load widgets params = self.module_dict[module] widget = ParamWidget(params, scrollWidget=self) widget.paramwidget_edited.connect(self.paramwidget_edited) self.param_widget_map[module] = widget self.params_layout.addWidget(widget) else: widget.show() self.visibleWidget = widget def on_module_changed(self): self.updateModuleParamWidget() self.module_changed.emit(self.module_combobox.currentText()) class TranslatorConfigPanel(ModuleConfigParseWidget): show_pre_MT_keyword_window = Signal() show_MT_keyword_window = Signal() show_OCR_keyword_window = Signal() def __init__(self, module_name, scrollWidget: QWidget = None, *args, **kwargs) -> None: super().__init__(module_name, GET_VALID_TRANSLATORS, scrollWidget=scrollWidget, *args, **kwargs) self.translator_changed = self.module_changed self.source_combobox = ConfigComboBox(scrollWidget=scrollWidget) self.target_combobox = ConfigComboBox(scrollWidget=scrollWidget) self.replacePreMTkeywordBtn = NoBorderPushBtn(self.tr("Keyword substitution for machine translation source text"), self) self.replacePreMTkeywordBtn.clicked.connect(self.show_pre_MT_keyword_window) self.replacePreMTkeywordBtn.setFixedWidth(500) self.replaceMTkeywordBtn = NoBorderPushBtn(self.tr("Keyword substitution for machine translation"), self) self.replaceMTkeywordBtn.clicked.connect(self.show_MT_keyword_window) self.replaceMTkeywordBtn.setFixedWidth(500) self.replaceOCRkeywordBtn = NoBorderPushBtn(self.tr("Keyword substitution for source text"), self) self.replaceOCRkeywordBtn.clicked.connect(self.show_OCR_keyword_window) self.replaceOCRkeywordBtn.setFixedWidth(500) st_layout = QHBoxLayout() st_layout.setSpacing(15) st_layout.setAlignment(Qt.AlignmentFlag.AlignLeft | Qt.AlignmentFlag.AlignVCenter) st_layout.addWidget(ParamNameLabel(self.tr('Source'))) st_layout.addWidget(self.source_combobox) st_layout.addWidget(ParamNameLabel(self.tr('Target'))) st_layout.addWidget(self.target_combobox) self.vlayout.insertLayout(1, st_layout) self.vlayout.addWidget(self.replaceOCRkeywordBtn) self.vlayout.addWidget(self.replacePreMTkeywordBtn) self.vlayout.addWidget(self.replaceMTkeywordBtn) def finishSetTranslator(self, translator: BaseTranslator): self.source_combobox.blockSignals(True) self.target_combobox.blockSignals(True) self.module_combobox.blockSignals(True) self.source_combobox.clear() self.target_combobox.clear() self.source_combobox.addItems(translator.supported_src_list) self.target_combobox.addItems(translator.supported_tgt_list) self.module_combobox.setCurrentText(translator.name) self.source_combobox.setCurrentText(translator.lang_source) self.target_combobox.setCurrentText(translator.lang_target) self.updateModuleParamWidget() self.source_combobox.blockSignals(False) self.target_combobox.blockSignals(False) self.module_combobox.blockSignals(False) class InpaintConfigPanel(ModuleConfigParseWidget): def __init__(self, module_name: str, scrollWidget: QWidget = None, *args, **kwargs) -> None: super().__init__(module_name, GET_VALID_INPAINTERS, scrollWidget = scrollWidget, *args, **kwargs) self.inpainter_changed = self.module_changed self.setInpainter = self.setModule self.needInpaintChecker = ParamCheckerBox(self.tr('Let the program decide whether it is necessary to use the selected inpaint method.')) self.vlayout.addWidget(self.needInpaintChecker) def showEvent(self, e) -> None: self.p_layout.insertWidget(1, self.module_combobox) super().showEvent(e) def hideEvent(self, e) -> None: self.p_layout.removeWidget(self.module_combobox) return super().hideEvent(e) class TextDetectConfigPanel(ModuleConfigParseWidget): def __init__(self, module_name: str, scrollWidget: QWidget = None, *args, **kwargs) -> None: super().__init__(module_name, GET_VALID_TEXTDETECTORS, scrollWidget = scrollWidget, *args, **kwargs) self.detector_changed = self.module_changed self.setDetector = self.setModule self.keep_existing_checker = QCheckBox(text=self.tr('Keep Existing Lines')) self.p_layout.insertWidget(2, self.keep_existing_checker) class OCRConfigPanel(ModuleConfigParseWidget): def __init__(self, module_name: str, scrollWidget: QWidget = None, *args, **kwargs) -> None: super().__init__(module_name, GET_VALID_OCR, scrollWidget = scrollWidget, *args, **kwargs) self.ocr_changed = self.module_changed self.setOCR = self.setModule self.restoreEmptyOCRChecker = QCheckBox(self.tr("Delete and restore region where OCR return empty string."), self) self.restoreEmptyOCRChecker.clicked.connect(self.on_restore_empty_ocr) self.vlayout.addWidget(self.restoreEmptyOCRChecker) def on_restore_empty_ocr(self): pcfg.restore_ocr_empty = self.restoreEmptyOCRChecker.isChecked()