diff --git a/images/images.cb7 b/images/images.cb7 new file mode 100644 index 00000000..a04b5335 Binary files /dev/null and b/images/images.cb7 differ diff --git a/requirements.txt b/requirements.txt index da222130..32d0989f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -29,3 +29,6 @@ numpy==1.26.4 # WD Tagger huggingface-hub==0.23.2 onnxruntime==1.18.0 + +rarfile +py7zr diff --git a/taggui/dialogs/settings_dialog.py b/taggui/dialogs/settings_dialog.py index 56cc28fd..4cd4de52 100644 --- a/taggui/dialogs/settings_dialog.py +++ b/taggui/dialogs/settings_dialog.py @@ -1,9 +1,10 @@ from PySide6.QtCore import Qt, Slot from PySide6.QtWidgets import (QDialog, QFileDialog, QGridLayout, QLabel, - QLineEdit, QPushButton, QVBoxLayout) + QLineEdit, QPushButton, QVBoxLayout, QTabWidget, + QWidget, QComboBox) from utils.settings import DEFAULT_SETTINGS, get_settings -from utils.settings_widgets import (SettingsBigCheckBox, SettingsLineEdit, +from utils.settings_widgets import (SettingsBigCheckBox, SettingsComboBox, SettingsLineEdit, SettingsSpinBox) @@ -16,97 +17,115 @@ def __init__(self, parent): layout.setContentsMargins(20, 20, 20, 20) layout.setSpacing(20) - grid_layout = QGridLayout() - grid_layout.addWidget(QLabel('Font size (pt)'), 0, 0, - Qt.AlignmentFlag.AlignRight) - grid_layout.addWidget(QLabel('File types to show in image list'), 1, 0, - Qt.AlignmentFlag.AlignRight) - grid_layout.addWidget(QLabel('Image width in image list (px)'), 2, 0, - Qt.AlignmentFlag.AlignRight) - grid_layout.addWidget(QLabel('Tag separator'), 3, 0, - Qt.AlignmentFlag.AlignRight) - grid_layout.addWidget(QLabel('Insert space after tag separator'), 4, 0, - Qt.AlignmentFlag.AlignRight) - grid_layout.addWidget(QLabel('Show tag autocomplete suggestions'), - 5, 0, Qt.AlignmentFlag.AlignRight) - grid_layout.addWidget(QLabel('Auto-captioning models directory'), 6, 0, - Qt.AlignmentFlag.AlignRight) - - font_size_spin_box = SettingsSpinBox( + # Create tab widget and tabs + tabs: QTabWidget = QTabWidget() + appearance_tab: QWidget = QWidget() + tagging_tab: QWidget = QWidget() + directories_tab: QWidget = QWidget() + comics_tab: QWidget = QWidget() # Future use + + tabs.addTab(appearance_tab, "Appearance") + tabs.addTab(tagging_tab, "Tagging") + tabs.addTab(directories_tab, "Directories") + tabs.addTab(comics_tab, "Comics") # Future use + + # Appearance Tab + appearance_layout: QGridLayout = QGridLayout(appearance_tab) + appearance_layout.addWidget(QLabel('Font size (pt)'), 0, 0, Qt.AlignmentFlag.AlignRight) + font_size_spin_box: SettingsSpinBox = SettingsSpinBox( key='font_size', default=DEFAULT_SETTINGS['font_size'], minimum=1, maximum=99) font_size_spin_box.valueChanged.connect(self.show_restart_warning) - # Images that are too small cause lag, so set a minimum width. - image_list_image_width_spin_box = SettingsSpinBox( + appearance_layout.addWidget(font_size_spin_box, 0, 1, Qt.AlignmentFlag.AlignLeft) + + appearance_layout.addWidget(QLabel('Image width in image list (px)'), 1, 0, Qt.AlignmentFlag.AlignRight) + image_list_image_width_spin_box: SettingsSpinBox = SettingsSpinBox( key='image_list_image_width', default=DEFAULT_SETTINGS['image_list_image_width'], minimum=16, maximum=9999) - image_list_image_width_spin_box.valueChanged.connect( - self.show_restart_warning) - tag_separator_line_edit = QLineEdit() - tag_separator = self.settings.value( + image_list_image_width_spin_box.valueChanged.connect(self.show_restart_warning) + appearance_layout.addWidget(image_list_image_width_spin_box, 1, 1, Qt.AlignmentFlag.AlignLeft) + + # Tagging Tab + tagging_layout: QGridLayout = QGridLayout(tagging_tab) + tagging_layout.addWidget(QLabel('File types to show in image list'), 0, 0, Qt.AlignmentFlag.AlignRight) + file_types_line_edit: SettingsLineEdit = SettingsLineEdit( + key='image_list_file_formats', + default=DEFAULT_SETTINGS['image_list_file_formats']) + file_types_line_edit.setMinimumWidth(400) + file_types_line_edit.textChanged.connect(self.show_restart_warning) + tagging_layout.addWidget(file_types_line_edit, 0, 1, Qt.AlignmentFlag.AlignLeft) + + tagging_layout.addWidget(QLabel('Tag separator'), 1, 0, Qt.AlignmentFlag.AlignRight) + tag_separator_line_edit: QLineEdit = QLineEdit() + tag_separator: str = self.settings.value( 'tag_separator', defaultValue=DEFAULT_SETTINGS['tag_separator'], type=str) tag_separator_line_edit.setMaximumWidth(50) tag_separator_line_edit.setText(tag_separator) - tag_separator_line_edit.textChanged.connect( - self.handle_tag_separator_change) - insert_space_after_tag_separator_check_box = SettingsBigCheckBox( + tag_separator_line_edit.textChanged.connect(self.handle_tag_separator_change) + tagging_layout.addWidget(tag_separator_line_edit, 1, 1, Qt.AlignmentFlag.AlignLeft) + + tagging_layout.addWidget(QLabel('Insert space after tag separator'), 2, 0, Qt.AlignmentFlag.AlignRight) + insert_space_after_tag_separator_check_box: SettingsBigCheckBox = SettingsBigCheckBox( key='insert_space_after_tag_separator', default=DEFAULT_SETTINGS['insert_space_after_tag_separator']) - insert_space_after_tag_separator_check_box.stateChanged.connect( - self.show_restart_warning) - autocomplete_tags_check_box = SettingsBigCheckBox( + insert_space_after_tag_separator_check_box.stateChanged.connect(self.show_restart_warning) + tagging_layout.addWidget(insert_space_after_tag_separator_check_box, 2, 1, Qt.AlignmentFlag.AlignLeft) + + tagging_layout.addWidget(QLabel('Show tag autocomplete suggestions'), 3, 0, Qt.AlignmentFlag.AlignRight) + autocomplete_tags_check_box: SettingsBigCheckBox = SettingsBigCheckBox( key='autocomplete_tags', default=DEFAULT_SETTINGS['autocomplete_tags']) - autocomplete_tags_check_box.stateChanged.connect( - self.show_restart_warning) - self.models_directory_line_edit = SettingsLineEdit( + autocomplete_tags_check_box.stateChanged.connect(self.show_restart_warning) + tagging_layout.addWidget(autocomplete_tags_check_box, 3, 1, Qt.AlignmentFlag.AlignLeft) + + # Directories Tab + directories_layout: QGridLayout = QGridLayout(directories_tab) + directories_layout.addWidget(QLabel('Auto-captioning models directory'), 0, 0, Qt.AlignmentFlag.AlignRight) + self.models_directory_line_edit: SettingsLineEdit = SettingsLineEdit( key='models_directory_path', default=DEFAULT_SETTINGS['models_directory_path']) self.models_directory_line_edit.setMinimumWidth(400) self.models_directory_line_edit.setClearButtonEnabled(True) - self.models_directory_line_edit.textChanged.connect( - self.show_restart_warning) + self.models_directory_line_edit.textChanged.connect(self.show_restart_warning) + directories_layout.addWidget(self.models_directory_line_edit, 0, 1, Qt.AlignmentFlag.AlignLeft) + models_directory_button = QPushButton('Select Directory...') - models_directory_button.setFixedWidth( - int(models_directory_button.sizeHint().width() * 1.3)) + models_directory_button.setFixedWidth(int(models_directory_button.sizeHint().width() * 1.3)) models_directory_button.clicked.connect(self.set_models_directory_path) - file_types_line_edit = SettingsLineEdit( - key='image_list_file_formats', - default=DEFAULT_SETTINGS['image_list_file_formats']) - file_types_line_edit.setMinimumWidth(400) - file_types_line_edit.textChanged.connect(self.show_restart_warning) + directories_layout.addWidget(models_directory_button, 1, 1, Qt.AlignmentFlag.AlignLeft) + + # Comics Tab + comics_layout: QGridLayout = QGridLayout(comics_tab) + comics_layout.addWidget(QLabel('Supported formats:'), 0, 0, Qt.AlignmentFlag.AlignRight) - grid_layout.addWidget(font_size_spin_box, 0, 1, - Qt.AlignmentFlag.AlignLeft) - grid_layout.addWidget(file_types_line_edit, 1, 1, - Qt.AlignmentFlag.AlignLeft) - grid_layout.addWidget(image_list_image_width_spin_box, 2, 1, - Qt.AlignmentFlag.AlignLeft) - grid_layout.addWidget(tag_separator_line_edit, 3, 1, - Qt.AlignmentFlag.AlignLeft) - grid_layout.addWidget(insert_space_after_tag_separator_check_box, 4, 1, - Qt.AlignmentFlag.AlignLeft) - grid_layout.addWidget(autocomplete_tags_check_box, 5, 1, - Qt.AlignmentFlag.AlignLeft) - grid_layout.addWidget(self.models_directory_line_edit, 6, 1, - Qt.AlignmentFlag.AlignLeft) - grid_layout.addWidget(models_directory_button, 7, 1, - Qt.AlignmentFlag.AlignLeft) - layout.addLayout(grid_layout) - - # Prevent the grid layout from moving to the center when the warning - # label is hidden. + self.comics_formats_line_edit: SettingsLineEdit = SettingsLineEdit( + key='comics_formats', + default=DEFAULT_SETTINGS['comics_formats']) + self.comics_formats_line_edit.setMinimumWidth(400) + self.comics_formats_line_edit.setClearButtonEnabled(True) + self.comics_formats_line_edit.textChanged.connect(self.show_restart_warning) + comics_layout.addWidget(self.comics_formats_line_edit, 0, 1, Qt.AlignmentFlag.AlignLeft) + + comics_layout.addWidget(QLabel('Tag options:'), 1, 0, Qt.AlignmentFlag.AlignRight) + self.comic_tag_options_combo: SettingsComboBox = SettingsComboBox( + key = 'comic_tag_type', + default=DEFAULT_SETTINGS['comic_tag_type'] + ) + self.comic_tag_options_combo.addItems(['Tag comic', 'Tag pages', 'Tag both']) + self.comic_tag_options_combo.currentTextChanged.connect(self.show_restart_warning) + comics_layout.addWidget(self.comic_tag_options_combo, 1, 1, Qt.AlignmentFlag.AlignLeft) + + layout.addWidget(tabs) + + # Prevent the grid layout from moving to the center when the warning label is hidden. layout.addStretch() - self.restart_warning = ('Restart the application to apply the new ' - 'settings.') + self.restart_warning = 'Restart the application to apply the new settings.' self.warning_label = QLabel(self.restart_warning) self.warning_label.setAlignment(Qt.AlignmentFlag.AlignCenter) self.warning_label.setStyleSheet('color: red;') layout.addWidget(self.warning_label) - # Fix the size of the dialog to its size when the warning label is - # shown. self.setFixedSize(self.sizeHint()) self.warning_label.hide() @@ -136,8 +155,9 @@ def set_models_directory_path(self): else: initial_directory_path = '' models_directory_path = QFileDialog.getExistingDirectory( - parent=self, caption='Select directory containing auto-captioning ' - 'models', + parent=self, caption='Select directory containing auto-captioning models', dir=initial_directory_path) if models_directory_path: self.models_directory_line_edit.setText(models_directory_path) + + diff --git a/taggui/models/image_list_model.py b/taggui/models/image_list_model.py index 81d0c3c0..0bd509f2 100644 --- a/taggui/models/image_list_model.py +++ b/taggui/models/image_list_model.py @@ -1,5 +1,9 @@ import random import sys +import zipfile +import rarfile +import tarfile +import py7zr from os.path import splitext from collections import Counter, deque from dataclasses import dataclass @@ -111,16 +115,29 @@ def load_directory(self, directory_path: Path): image_suffixes_string = settings.value( 'image_list_file_formats', defaultValue=DEFAULT_SETTINGS['image_list_file_formats'], type=str) + comic_formats_string = settings.value( + 'comics_formats', + defaultValue=DEFAULT_SETTINGS['comics_formats'], type=str) image_suffixes = [] for suffix in image_suffixes_string.split(','): suffix = suffix.strip().lower() if not suffix.startswith('.'): suffix = '.' + suffix image_suffixes.append(suffix) - image_paths = [path for path in file_paths - if str(path).lower().endswith(tuple(image_suffixes))] + comic_suffixes = [] + for suffix in comic_formats_string.split(','): + suffix = suffix.strip().lower() + if not suffix.startswith('.'): + suffix = '.' + suffix + comic_suffixes.append(suffix) + image_paths = [path for path in file_paths + if path.suffix.lower() in image_suffixes] + comic_paths = {path for path in file_paths + if path.suffix.lower() in comic_suffixes} text_file_paths = [path for path in file_paths if path.suffix == '.txt'] + # Comparing paths is slow on some systems, so convert the paths to + # strings. txt_strs = {str(path) for path in text_file_paths} for image_path in image_paths: try: @@ -158,9 +175,60 @@ def load_directory(self, directory_path: Path): tags = [tag for tag in tags if tag] image = Image(image_path, dimensions, tags) self.images.append(image) + + for comic_path in comic_paths: + self.load_comic(comic_path) + self.images.sort(key=lambda image_: image_.path) self.modelReset.emit() + def load_comic(self, comic_path: Path): + settings = get_settings() + comic_tag_type = settings.value('comic_tag_type', defaultValue=1, type=int) + comic_inject_tags = settings.value('comic_inject_tags', defaultValue=True, type=bool) + try: + if comic_path.suffix.lower() == '.cbz': + with zipfile.ZipFile(comic_path, 'r') as archive: + self._process_comic_archive(archive, comic_tag_type, comic_inject_tags) + elif comic_path.suffix.lower() == '.cbr': + with rarfile.RarFile(comic_path, 'r') as archive: + self._process_comic_archive(archive, comic_tag_type, comic_inject_tags) + elif comic_path.suffix.lower() == '.cbt': + with tarfile.open(comic_path, 'r') as archive: + self._process_comic_archive(archive, comic_tag_type, comic_inject_tags) + elif comic_path.suffix.lower() == '.cb7': + with py7zr.SevenZipFile(comic_path, 'r') as archive: + self._process_comic_archive(archive, comic_tag_type, comic_inject_tags) + except Exception as exception: + print(f'Failed to load comic {comic_path}: {exception}', file=sys.stderr) + + def _process_comic_archive(self, archive, comic_tag_type: int, comic_inject_tags: bool): + image_files = [f for f in archive.getnames() if f.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp'))] + text_files = [f for f in archive.getnames() if f.lower().endswith('.txt')] + + for image_file in image_files: + dimensions = imagesize.get(image_file) + tags = [] + + if comic_tag_type in (1, 3): + comic_tags_file = next((f for f in text_files if f.lower() == 'tags.txt'), None) + if comic_tags_file: + with archive.open(comic_tags_file) as file: + caption = file.read().decode('utf-8', errors='replace') + tags.extend(caption.split(self.tag_separator)) + + if comic_tag_type in (2, 3): + image_tag_file = Path(image_file).with_suffix('.txt') + if image_tag_file in text_files: + with archive.open(image_tag_file) as file: + caption = file.read().decode('utf-8', errors='replace') + tags.extend(caption.split(self.tag_separator)) + + tags = [tag.strip() for tag in tags] + tags = [tag for tag in tags if tag] + image = Image(Path(image_file), dimensions, tags) + self.images.append(image) + def add_to_undo_stack(self, action_name: str, should_ask_for_confirmation: bool): """Add the current state of the image tags to the undo stack.""" diff --git a/taggui/utils/settings.py b/taggui/utils/settings.py index 88e1f0f5..25e16fca 100644 --- a/taggui/utils/settings.py +++ b/taggui/utils/settings.py @@ -9,7 +9,10 @@ 'tag_separator': ',', 'insert_space_after_tag_separator': True, 'autocomplete_tags': True, - 'models_directory_path': '' + 'models_directory_path': '', + 'comics_formats': 'cbz, cbr, cb7, cbt', + 'comic_tag_type': 1, + 'comic_inject_tags': True }