From 754c11f4596db4b8d71e5f8cc3254523cef45183 Mon Sep 17 00:00:00 2001 From: Tolkachev Aleksei Date: Wed, 7 Aug 2024 14:13:30 +0500 Subject: [PATCH] blocks navigation + import/export + menu highlight --- app/ui/dayu_widgets/push_button.py | 5 + app/ui/main_window.py | 35 +++++-- comic.py | 142 +++++++++++++++++++++++++++-- pipeline.py | 1 + 4 files changed, 170 insertions(+), 13 deletions(-) diff --git a/app/ui/dayu_widgets/push_button.py b/app/ui/dayu_widgets/push_button.py index 7cb6549..66c042f 100644 --- a/app/ui/dayu_widgets/push_button.py +++ b/app/ui/dayu_widgets/push_button.py @@ -98,6 +98,11 @@ def primary(self): self.set_dayu_type(MPushButton.PrimaryType) return self + def default(self): + """Set MPushButton to DefaultType""" + self.set_dayu_type(MPushButton.DefaultType) + return self + def success(self): """Set MPushButton to SuccessType""" self.set_dayu_type(MPushButton.SuccessType) diff --git a/app/ui/main_window.py b/app/ui/main_window.py index 43f1c46..6d4a9cd 100644 --- a/app/ui/main_window.py +++ b/app/ui/main_window.py @@ -325,8 +325,30 @@ def _create_main_content(self): input_layout.addLayout(t_combo_text_layout) + self.load_blocks_state_button = MClickBrowserFileToolButton(multiple=False) + self.load_blocks_state_button.set_dayu_svg("folder_fill.svg") + self.load_blocks_state_button.set_dayu_filters([".txt"]) + self.load_blocks_state_button.setToolTip(self.tr("Import text file with translations")) + self.load_blocks_state_button.sig_file_changed.connect(self.load_blocks_state) + + blocks_layout = QtWidgets.QHBoxLayout() + self.blocks_checker_group = MToolButtonGroup(orientation=QtCore.Qt.Horizontal, exclusive=False) + blocks_checker_buttons = [ + {"text": self.tr("Previous"), "svg": "left_fill.svg", "checkable": False, "tooltip": self.tr("Previous block"), "clicked": self.select_prev_text}, + {"text": self.tr("Next"), "svg": "right_fill.svg", "checkable": False, "tooltip": self.tr("Next block"), "clicked": self.select_next_text}, + {"text": self.tr("Save state"), "svg": "save_fill.svg", "checkable": False, "tooltip": self.tr("Backup all blocks to file"), "clicked": self.save_blocks_state}, + {"text": self.tr("Load saved"), "svg": "folder_fill.svg", "checkable": False, "tooltip": self.tr("Import blocks from file"), "clicked": self.load_blocks_button}, + ] + self.blocks_checker_group.set_button_list(blocks_checker_buttons) + blocks_layout.addWidget(self.blocks_checker_group) + blocks_panel_layout = QtWidgets.QVBoxLayout() + blocks_panel_header = MDivider(self.tr('Actions with text blocks')) + blocks_panel_layout.addWidget(blocks_panel_header) + blocks_panel_layout.addLayout(blocks_layout) + blocks_panel_layout.addSpacing(20) + # Tools Layout - tools_widget = QtWidgets.QWidget() + tools_widget = QtWidgets.QWidget() tools_layout = QtWidgets.QVBoxLayout() misc_lay = QtWidgets.QHBoxLayout() @@ -371,21 +393,21 @@ def _create_main_content(self): self.change_all_blocks_size_dec = self.create_tool_button(svg="minus_line.svg") self.change_all_blocks_size_dec.setToolTip(self.tr("Reduce the size of all blocks")) - + self.change_all_blocks_size_diff = MLineEdit() self.change_all_blocks_size_diff.setFixedWidth(30) self.change_all_blocks_size_diff.setText("3") - + # Set up integer validator int_validator = QIntValidator() self.change_all_blocks_size_diff.setValidator(int_validator) - + # Optional: Ensure the text is center-aligned self.change_all_blocks_size_diff.setAlignment(QtCore.Qt.AlignCenter) - + self.change_all_blocks_size_inc = self.create_tool_button(svg="add_line.svg") self.change_all_blocks_size_inc.setToolTip(self.tr("Increase the size of all blocks")) - + box_tools_lay.addStretch() box_tools_lay.addWidget(self.change_all_blocks_size_dec) box_tools_lay.addWidget(self.change_all_blocks_size_diff) @@ -473,6 +495,7 @@ def _create_main_content(self): right_layout.addLayout(input_layout) right_layout.addLayout(font_panel_layout) + right_layout.addLayout(blocks_panel_layout) right_layout.addWidget(tools_scroll) right_layout.addStretch() diff --git a/comic.py b/comic.py index e707d4b..98d7832 100644 --- a/comic.py +++ b/comic.py @@ -1,5 +1,6 @@ import os import cv2, shutil +import json import tempfile import numpy as np from typing import Callable, Tuple, List @@ -19,6 +20,8 @@ from app.thread_worker import GenericWorker from app.ui.dayu_widgets.message import MMessage +from datetime import datetime + from modules.detection import do_rectangles_overlap from modules.utils.textblock import TextBlock from modules.rendering.render import draw_text @@ -47,6 +50,7 @@ def __init__(self, parent=None): self.image_states = {} self.blk_list = [] + self.cleaned_image = None # Store cleaned image self.image_data = {} # Store the latest version of each image self.image_history = {} # Store undo/redo history for each image self.current_history_index = {} # Current position in the undo/redo history for each image @@ -160,18 +164,126 @@ def set_block_font_settings(self): if blk.max_font_size > 0: self.max_font_spinbox.setValue(blk.max_font_size) + def get_current_block_index(self): + if self.current_text_block: + return self.blk_list.index(self.current_text_block) + return 0 + + def select_prev_text(self): + if len(self.blk_list) == 0: + return + current_block_index = self.get_current_block_index() + if current_block_index == 0: + block_index = -1 + else: + block_index = current_block_index - 1 + rect = self.find_corresponding_rect(self.blk_list[block_index], 0.5) + if rect == None: + return + self.image_viewer.select_rectangle(rect) + + def select_next_text(self): + if len(self.blk_list) == 0: + return + current_block_index = self.get_current_block_index() + if current_block_index == len(self.blk_list) - 1: + block_index = 0 + else: + block_index = current_block_index + 1 + rect = self.find_corresponding_rect(self.blk_list[block_index], 0.5) + if rect == None: + return + self.image_viewer.select_rectangle(rect) + + def save_blocks_state(self): + if len(self.blk_list) == 0: + return + date_time = datetime.now().strftime("_%Y-%m-%d_%H-%M-%S") + file_name_original = self.image_files[self.current_image_index] + file_name = file_name_original[0:-4] + date_time + ".txt" + a = open(file_name, 'w') + + default_init_font_size = self.settings_page.get_min_font_size() + default_min_font_size = self.settings_page.get_max_font_size() + + for blk in self.blk_list: + blk_rect = tuple(blk.xyxy) + blk_rect_export = str(int(blk_rect[0])) + ',' + str(int(blk_rect[1])) + ',' + str(int(blk_rect[2])) + ',' + str(int(blk_rect[3])) + + if blk.min_font_size > 0: + min_font_size = blk.min_font_size + else: + min_font_size = default_min_font_size + if blk.max_font_size > 0: + init_font_size = blk.max_font_size + else: + init_font_size = default_init_font_size + + blk_to_save = { + 'text': blk.text, + 'rect': blk_rect_export, + 'translation': blk.translation, + 'min_font_size': min_font_size, + 'init_font_size': init_font_size, + } + a.write(json.dumps(blk_to_save, ensure_ascii=False) + "\n") + + a.close() + dialog_message = "File " + file_name + " with data saved\n" + + if self.cleaned_image is not None: + print('cleaned_image') + cv2_img = self.cleaned_image #self.image_data[file_name_original] + cv2_img_save = cv2.cvtColor(cv2_img, cv2.COLOR_BGR2RGB) + sv_pth = file_name_original[0:-4] + date_time + file_name_original[-4:] + cv2.imwrite(sv_pth, cv2_img_save) + dialog_message += "File " + sv_pth + " with cleaned image saved\n" + + dialog = QtWidgets.QDialog(self) + dialog.setWindowTitle("Export completed successfully") + label = QtWidgets.QLabel(dialog) + label.setText(dialog_message) + label.setMargin(20) + label.adjustSize() + dialog.exec_() + + def load_blocks_button(self): + self.load_blocks_state_button.click() + + def load_blocks_state(self, file_path: str): + updated_blk_list = [] + with open(file_path, 'r') as f: + for line in f.readlines(): + blk_to_load = json.loads(line) + blk_rect_coord = blk_to_load['rect'].split(',') + new_blk_coord = [int(blk_rect_coord[0]), int(blk_rect_coord[1]), int(blk_rect_coord[2]), int(blk_rect_coord[3])] + new_blk = TextBlock(new_blk_coord) + new_blk.translation = blk_to_load['translation'] + new_blk.text = blk_to_load['text'] + new_blk.min_font_size = blk_to_load['min_font_size'] + new_blk.init_font_size = blk_to_load['init_font_size'] + updated_blk_list.append(new_blk) + + self.blk_list = updated_blk_list + self.pipeline.load_box_coords(self.blk_list) + + def menu_highlight(self, button_index: int): + self.hbutton_group.get_button_group().buttons()[button_index].success() + def batch_mode_selected(self): self.disable_hbutton_group() self.translate_button.setEnabled(True) self.cancel_button.setEnabled(True) self.set_manual_font_settings_enabled(False) + self.blocks_checker_group.setVisible(False) def manual_mode_selected(self): self.enable_hbutton_group() self.translate_button.setEnabled(False) self.cancel_button.setEnabled(False) self.set_manual_font_settings_enabled(True) - + self.blocks_checker_group.setVisible(True) + def on_image_processed(self, index: int, rendered_image: np.ndarray, image_path: str): if index == self.current_image_index: self.set_cv2_image(rendered_image) @@ -257,7 +369,8 @@ def block_detect(self, load_rects: bool = True): self.disable_hbutton_group() self.run_threaded(self.pipeline.detect_blocks, self.pipeline.on_blk_detect_complete, self.default_error_handler, self.on_manual_finished, load_rects) - + self.menu_highlight(0) + def clear_text_edits(self): self.current_text_block = None self.s_text_edit.clear() @@ -278,6 +391,7 @@ def ocr(self): self.loading.setVisible(True) self.disable_hbutton_group() self.run_threaded(self.pipeline.OCR_image, None, self.default_error_handler, self.finish_ocr_translate) + self.menu_highlight(1) def translate_image(self): source_lang = self.s_combo.currentText() @@ -287,6 +401,7 @@ def translate_image(self): self.loading.setVisible(True) self.disable_hbutton_group() self.run_threaded(self.pipeline.translate_image, None, self.default_error_handler, self.finish_ocr_translate) + self.menu_highlight(2) def inpaint_and_set(self): if self.image_viewer.hasPhoto() and self.image_viewer.has_drawn_elements(): @@ -295,6 +410,7 @@ def inpaint_and_set(self): self.disable_hbutton_group() self.run_threaded(self.pipeline.inpaint, self.pipeline.inpaint_complete, self.default_error_handler, self.on_manual_finished) + self.menu_highlight(4) def load_images_threaded(self, file_paths: List[str]): self.file_handler.file_paths = file_paths @@ -349,6 +465,9 @@ def on_images_loaded(self, loaded_images: List[Tuple[str, np.ndarray]]): self.image_viewer.resetTransform() self.image_viewer.fitInView() + for button in self.hbutton_group.get_button_group().buttons(): + button.default() + def update_image_cards(self): # Clear existing cards for i in reversed(range(self.image_card_layout.count())): @@ -393,7 +512,8 @@ def save_image_state(self, file: str): 'target_text': self.t_text_edit.toPlainText(), 'target_lang': self.t_combo.currentText(), 'brush_strokes': self.image_viewer.save_brush_strokes(), - 'blk_list': self.blk_list + 'blk_list': self.blk_list, + 'cleaned_image': self.cleaned_image, } def save_current_image_state(self): @@ -414,6 +534,7 @@ def load_image_state(self, file_path: str): self.t_combo.setCurrentText(state['target_lang']) self.image_viewer.load_brush_strokes(state['brush_strokes']) self.blk_list = state['blk_list'] + self.cleaned_image = state['cleaned_image'] else: self.s_text_edit.clear() self.t_text_edit.clear() @@ -457,8 +578,9 @@ def load_segmentation_points(self): self.image_viewer.draw_segmentation_lines(bboxes) self.enable_hbutton_group() - + self.menu_highlight(3) else: + print('no blk') self.loading.setVisible(True) self.disable_hbutton_group() self.run_threaded(self.pipeline.detect_blocks, self.blk_detect_segment, @@ -490,6 +612,7 @@ def undo_image(self): if self.current_image_index >= 0: file_path = self.image_files[self.current_image_index] current_index = self.current_history_index[file_path] + last_index = current_index while current_index > 0: current_index -= 1 cv2_img = self.image_history[file_path][current_index] @@ -498,6 +621,8 @@ def undo_image(self): self.image_data[file_path] = cv2_img self.image_viewer.display_cv2_image(cv2_img) break + if last_index >= 2: + self.pipeline.load_box_coords(self.blk_list) def redo_image(self): if self.current_image_index >= 0: @@ -534,6 +659,8 @@ def update_blk_list(self): # Add new TextBlocks for any remaining rectangles for new_rect in current_rectangles: + pprint(vars(new_rect)) + pprint(np.array(new_rect)) new_blk = TextBlock(np.array(new_rect)) updated_blk_list.append(new_blk) @@ -636,7 +763,8 @@ def render_text(self): self.run_threaded(draw_text, self.on_render_complete, self.default_error_handler, None, inpaint_image, self.blk_list, font_path, colour=font_color, init_font_size=max_font_size, min_font_size=min_font_size, outline=outline) - + self.menu_highlight(5) + def handle_rectangle_change(self, new_rect: QtCore.QRectF): # Find the corresponding TextBlock in blk_list for blk in self.blk_list: @@ -775,7 +903,7 @@ def get_system_language(): 'ko': '한국어', 'fr': 'Français', 'ja': '日本語', - 'ru': 'русский', + 'ru': 'Русский', 'de': 'Deutsch', 'nl': 'Nederlands', 'es': 'Español', @@ -794,7 +922,7 @@ def load_translation(app, language: str): '日本語': 'ja', '简体中文': 'zh_CN', '繁體中文': 'zh_TW', - 'русский': 'ru', + 'Русский': 'ru', 'Deutsch': 'de', 'Nederlands': 'nl', 'Español': 'es', diff --git a/pipeline.py b/pipeline.py index c88630a..902d440 100644 --- a/pipeline.py +++ b/pipeline.py @@ -83,6 +83,7 @@ def manual_inpaint(self): def inpaint_complete(self, result): inpainted, original_image = result self.main_page.set_cv2_image(inpainted) + self.main_page.cleaned_image = inpainted # get_best_render_area(self.main_page.blk_list, original_image, inpainted) self.load_box_coords(self.main_page.blk_list)