add archive reading support

requirements.txt

@@ -29,3 +29,6 @@ numpy==1.26.4
 # WD Tagger
 huggingface-hub==0.23.2
 onnxruntime==1.18.0
+
+rarfile
+py7zr

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)
+
+

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."""

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
 }