headless mode #413 (677fd20b) · Commits · git-mirror / BallonsTranslator

launch.py

+20 −14

Original line number	Diff line number	Diff line
		@@ -33,6 +33,8 @@ parser.add_argument("--proj-dir", default='', type=str, help='Open project direc
		parser.add_argument("--qt-api", default='', choices=QT_APIS, help='Set qt api')
		parser.add_argument("--debug", action='store_true')
		parser.add_argument("--requirements", default='requirements.txt')
		parser.add_argument("--headless", action='store_true', help='run without GUI')
		parser.add_argument("--exec_dirs", default='', help='translation queue (project directories) separated by comma')
		args, _ = parser.parse_known_args()


		@@ -123,6 +125,7 @@ def main():

		if args.debug:
		os.environ['BALLOONTRANS_DEBUG'] = '1'
		os.environ['BT_HEADLESS'] = '1' if args.headless else '0'

		if not args.qt_api in QT_APIS:
		os.environ['QT_API'] = 'pyqt6'
		@@ -186,7 +189,10 @@ def main():

		load_modules()

		app = QApplication(sys.argv)
		app_args = sys.argv
		if args.headless:
		app_args = sys.argv + ['-platform', 'offscreen']
		app = QApplication(app_args)

		lang = config.display_lang
		langp = osp.join(shared.TRANSLATE_DIR, lang + '.qm')
		@@ -198,11 +204,6 @@ def main():
		LOGGER.warning(f'target display language file {langp} doesnt exist.')
		LOGGER.info(f'set display language to {lang}')

		ps = QGuiApplication.primaryScreen()
		shared.LDPI = ps.logicalDotsPerInch()
		shared.SCREEN_W = ps.geometry().width()
		shared.SCREEN_H = ps.geometry().height()

		# Fonts
		# Load custom fonts if they exist
		for font in os.listdir(PATH_FONTS):
		@@ -226,11 +227,16 @@ def main():

		from ui.mainwindow import MainWindow

		ballontrans = MainWindow(app, config, open_dir=args.proj_dir)
		ballontrans = MainWindow(app, config, open_dir=args.proj_dir, **vars(args))
		global BT
		BT = ballontrans
		BT.restart_signal.connect(restart)

		if not args.headless:
		ps = QGuiApplication.primaryScreen()
		shared.LDPI = ps.logicalDotsPerInch()
		shared.SCREEN_W = ps.geometry().width()
		shared.SCREEN_H = ps.geometry().height()
		if shared.SCREEN_W > 1707 and sys.platform == 'win32': # higher than 2560 (1440p) / 1.5
		# https://github.com/dmMaze/BallonsTranslator/issues/220
		BT.comicTransSplitter.setHandleWidth(10)

ui/mainwindow.py

+53 −8

Original line number	Diff line number	Diff line
		import os.path as osp
		import os, re, traceback, sys
		from typing import List
		from typing import List, Union
		from pathlib import Path
		import subprocess
		from functools import partial
		import time

		from qtpy.QtWidgets import QFileDialog, QMenu, QHBoxLayout, QVBoxLayout, QApplication, QStackedWidget, QSplitter, QListWidget, QShortcut, QListWidgetItem, QMessageBox, QTextEdit, QPlainTextEdit
		from qtpy.QtCore import Qt, QPoint, QSize, QEvent, Signal
		@@ -26,7 +27,7 @@ from .drawingpanel import DrawingPanel
		from .scenetext_manager import SceneTextManager, TextPanel, PasteSrcItemsCommand
		from .mainwindowbars import TitleBar, LeftBar, BottomBar
		from .io_thread import ImgSaveThread, ImportDocThread, ExportDocThread
		from .stylewidgets import FrameLessMessageBox, ImgtransProgressMessageBox
		from .stylewidgets import FrameLessMessageBox, ImgtransProgressMessageBox, Widget
		from .global_search_widget import GlobalSearchWidget
		from .textedit_commands import GlobalRepalceAllCommand
		from .framelesswindow import FramelessWindow
		@@ -52,8 +53,9 @@ class PageListView(QListWidget):

		return super().contextMenuEvent(e)


		class MainWindow(FramelessWindow):
		RUN_HEADLESS = os.environ['BT_HEADLESS'] == '1'
		mainwindow_cls = Widget if RUN_HEADLESS else FramelessWindow
		class MainWindow(mainwindow_cls):

		imgtrans_proj: ProjImgTrans = ProjImgTrans()
		save_on_page_changed = True
		@@ -65,9 +67,9 @@ class MainWindow(FramelessWindow):

		restart_signal = Signal()

		def __init__(self, app: QApplication, config: ProgramConfig, open_dir='', args, *kwargs) -> None:
		def __init__(self, app: QApplication, config: ProgramConfig, open_dir='', **exec_args) -> None:

		super().__init__(args, *kwargs)
		super().__init__()

		self.app = app
		self.setupThread()
		@@ -86,6 +88,9 @@ class MainWindow(FramelessWindow):
		if osp.exists(proj_dir):
		self.OpenProj(proj_dir)

		if RUN_HEADLESS:
		self.run_batch(**exec_args)

		def setStyleSheet(self, styleSheet: str) -> None:
		self.imgtrans_progress_msgbox.setStyleSheet(styleSheet)
		self.export_doc_thread.progress_bar.setStyleSheet(styleSheet)
		@@ -834,8 +839,10 @@ class MainWindow(FramelessWindow):

		def on_imgtrans_pipeline_finished(self):
		self.postprocess_mt_toggle = True
		if pcfg.module.empty_runcache:
		if pcfg.module.empty_runcache and not RUN_HEADLESS:
		self.module_manager.unload_all_models()
		if RUN_HEADLESS:
		self.run_next_dir()

		def postprocess_translations(self, blk_list: List[TextBlock]) -> None:
		src_is_cjk = is_cjk(pcfg.module.translate_source)
		@@ -985,6 +992,9 @@ class MainWindow(FramelessWindow):
		pcfg.display_lang = lang
		self.set_display_lang(lang)

		def run_imgtrans(self):
		self.on_run_imgtrans()

		def on_run_imgtrans(self):

		if self.bottomBar.textblockChecker.isChecked():
		@@ -1179,3 +1189,38 @@ class MainWindow(FramelessWindow):
		self.canvas.scale_tool_mode = False
		self.canvas.end_scale_tool.emit()
		return super().keyReleaseEvent(event)

		def run_batch(self, exec_dirs: Union[List, str], **kwargs):
		if not isinstance(exec_dirs, List):
		exec_dirs = exec_dirs.split(',')
		valid_dirs = []
		for d in exec_dirs:
		if osp.exists(d):
		valid_dirs.append(d)
		else:
		LOGGER.warning(f'target directory {d} does not exist.')
		self.exec_dirs = exec_dirs
		self.run_next_dir()

		def run_next_dir(self):
		if len(self.exec_dirs) == 0:
		LOGGER.info(f'finished translating all dirs, quit app...')
		self.app.quit()
		return
		d = self.exec_dirs.pop(0)
		from tqdm import tqdm

		LOGGER.info(f'translating {d} ...')
		self.openDir(d)
		shared.pbar = {}
		npages = len(self.imgtrans_proj.pages)
		if npages > 0:
		if pcfg.module.enable_detect:
		shared.pbar['detect'] = tqdm(range(npages), desc="Text Detection")
		if pcfg.module.enable_ocr:
		shared.pbar['ocr'] = tqdm(range(npages), desc="OCR")
		if pcfg.module.enable_translate:
		shared.pbar['translate'] = tqdm(range(npages), desc="Translation")
		if pcfg.module.enable_inpaint:
		shared.pbar['inpaint'] = tqdm(range(npages), desc="Inpaint")
		self.run_imgtrans()

ui/module_manager.py

+35 −3

Original line number	Diff line number	Diff line
		@@ -20,7 +20,7 @@ from modules import INPAINTERS, TRANSLATORS, TEXTDETECTORS, OCR, \
		import modules
		modules.translators.SYSTEM_LANG = QLocale.system().name()
		from modules.textdetector import TextBlock

		from utils import shared
		from .stylewidgets import ImgtransProgressMessageBox
		from .configpanel import ConfigPanel
		from utils.config import pcfg
		@@ -276,6 +276,11 @@ class ImgtransThread(QThread):
		finish_blktrans_stage = Signal(str, int)
		finish_blktrans = Signal(int, list)

		detect_counter = 0
		ocr_counter = 0
		translate_counter = 0
		inpaint_counter = 0

		def __init__(self,
		textdetect_thread: TextDetectThread,
		ocr_thread: OCRThread,
		@@ -309,6 +314,7 @@ class ImgtransThread(QThread):

		def runImgtransPipeline(self, imgtrans_proj: ProjImgTrans):
		self.imgtrans_proj = imgtrans_proj
		self.num_pages = len(self.imgtrans_proj.pages)
		self.job = self._imgtrans_pipeline
		self.start()

		@@ -712,6 +718,8 @@ class ModuleManager(QObject):

		def on_update_detect_progress(self, progress: int):
		ri = self.imgtrans_thread.recent_finished_index(progress)
		if 'detect' in shared.pbar:
		shared.pbar['detect'].update(progress)
		progress = int(progress / self.imgtrans_thread.num_pages * 100)
		self.progress_msgbox.updateDetectProgress(progress)
		if ri != self.last_finished_index:
		@@ -722,6 +730,8 @@ class ModuleManager(QObject):

		def on_update_ocr_progress(self, progress: int):
		ri = self.imgtrans_thread.recent_finished_index(progress)
		if 'ocr' in shared.pbar:
		shared.pbar['ocr'].update(progress)
		progress = int(progress / self.imgtrans_thread.num_pages * 100)
		self.progress_msgbox.updateOCRProgress(progress)
		if ri != self.last_finished_index:
		@@ -732,6 +742,8 @@ class ModuleManager(QObject):

		def on_update_translate_progress(self, progress: int):
		ri = self.imgtrans_thread.recent_finished_index(progress)
		if 'translate' in shared.pbar:
		shared.pbar['translate'].update(progress)
		progress = int(progress / self.imgtrans_thread.num_pages * 100)
		self.progress_msgbox.updateTranslateProgress(progress)
		if ri != self.last_finished_index:
		@@ -742,6 +754,8 @@ class ModuleManager(QObject):

		def on_update_inpaint_progress(self, progress: int):
		ri = self.imgtrans_thread.recent_finished_index(progress)
		if 'inpaint' in shared.pbar:
		shared.pbar['inpaint'].update(progress)
		progress = int(progress / self.imgtrans_thread.num_pages * 100)
		self.progress_msgbox.updateInpaintProgress(progress)
		if ri != self.last_finished_index:
		@@ -750,11 +764,29 @@ class ModuleManager(QObject):
		if progress == 100:
		self.finishImgtransPipeline()

		def finishImgtransPipeline(self):
		def progress(self):
		progress = {}
		num_pages = self.imgtrans_thread.num_pages
		if cfg_module.enable_detect:
		progress['detect'] = self.imgtrans_thread.detect_counter / num_pages
		if cfg_module.enable_ocr:
		progress['ocr'] = self.imgtrans_thread.ocr_counter / num_pages
		if cfg_module.enable_inpaint:
		progress['inpaint'] = self.imgtrans_thread.inpaint_counter / num_pages
		if cfg_module.enable_translate:
		progress['translate'] = self.imgtrans_thread.translate_counter / num_pages
		return progress

		def proj_finished(self):
		if self.imgtrans_thread.detect_finished() \
		and self.imgtrans_thread.ocr_finished() \
		and self.imgtrans_thread.translate_finished() \
		and self.imgtrans_thread.inpaint_finished():
		return True
		return False

		def finishImgtransPipeline(self):
		if self.proj_finished():
		self.progress_msgbox.hide()
		self.imgtrans_pipeline_finished.emit()

utils/shared.py

+1 −0

Original line number	Diff line number	Diff line
		@@ -93,6 +93,7 @@ check_local_file_hash = True

		FONT_FAMILIES: set = None
		CUSTOM_FONTS = []
		pbar = {}

		showed_exception = set()