From d01ab042d76646ef0a47f5f7beb78e4435e4a1dc Mon Sep 17 00:00:00 2001 From: StableLlama Date: Mon, 3 Feb 2025 19:47:00 +0100 Subject: [PATCH 01/18] Implement export functionality Still missing: - color space handling - JPEG XL (different PR) - crop editor (different PR) --- taggui/dialogs/export_dialog.py | 482 ++++++++++++++++++++++++++++++++ taggui/utils/image.py | 1 + taggui/utils/settings.py | 11 +- taggui/widgets/main_window.py | 10 + 4 files changed, 503 insertions(+), 1 deletion(-) create mode 100644 taggui/dialogs/export_dialog.py diff --git a/taggui/dialogs/export_dialog.py b/taggui/dialogs/export_dialog.py new file mode 100644 index 00000000..55d297aa --- /dev/null +++ b/taggui/dialogs/export_dialog.py @@ -0,0 +1,482 @@ +from enum import Enum +from collections import defaultdict +from math import floor +import os +from pathlib import Path + +from PySide6.QtCore import Qt, Slot +from PySide6.QtWidgets import (QWidget, QDialog, QFileDialog, QGridLayout, QLabel, + QLineEdit, QPushButton, QVBoxLayout, QHBoxLayout, + QTableWidget, QTableWidgetItem, QSizePolicy, + QMessageBox) +from PIL import Image, ImageFilter #, ImageQt, ImageEnhance, ImageCms + +from utils.settings import DEFAULT_SETTINGS, get_settings +from utils.settings_widgets import (SettingsBigCheckBox, SettingsLineEdit, + SettingsSpinBox, SettingsComboBox) +from models.image_list_model import ImageListModel + +Presets = { + 'manual': (0, 0), + 'Direct feed through': (0, 1), + 'SD1': (512, 64), + 'SDXL, SD3, Flux': (1024, 64) +} + +class ExportFormat(str, Enum): + JPG = '.jpg - JPEG' + PNG = '.png - PNG' + WEBP = '.webp - WEBP' + +ExportFormatDict = { + ExportFormat.JPG: 'jpeg', + ExportFormat.PNG: 'png', + ExportFormat.WEBP: 'webp' +} + +class BucketStrategy(str, Enum): + CROP = 'crop' + SCALE = 'scale' + CROP_SCALE = 'crop and scale' + +class ExportDialog(QDialog): + def __init__(self, parent, image_list_model: ImageListModel): + super().__init__(parent) + self.image_list_model = image_list_model + self.settings = get_settings() + self.inhibit_statistics_update = True + self.resolution_cache: dict[tuple, tuple] = {} + self.setWindowTitle('Export') + layout = QVBoxLayout(self) + layout.setContentsMargins(20, 20, 20, 20) + layout.setSpacing(20) + + grid_layout = QGridLayout() + grid_layout.setColumnStretch(0, 0) + grid_layout.setColumnStretch(1, 1) + + grid_row = 0 + grid_layout.addWidget(QLabel('Preset'), grid_row, 0, + Qt.AlignmentFlag.AlignRight) + preset_combo_box = SettingsComboBox(key='export_preset') + preset_combo_box.addItems(list(Presets)) + preset_combo_box.currentTextChanged.connect(self.apply_preset) + grid_layout.addWidget(preset_combo_box, grid_row, 1, + Qt.AlignmentFlag.AlignLeft) + + grid_row += 1 + grid_layout.addWidget(QLabel('Resolution (px)'), grid_row, 0, + Qt.AlignmentFlag.AlignRight) + self.resolution_spin_box = SettingsSpinBox( + key='export_resolution', default=DEFAULT_SETTINGS['export_resolution'], + minimum=0, maximum=8192) + self.resolution_spin_box.setToolTip('Common values:\n' + '0: disable rescaling\n' + '512: SD1.5\n' + '1024: SDXL, SD3, Flux') + self.resolution_spin_box.textChanged.connect(self.show_megapixels) + self.resolution_spin_box.textChanged.connect(self.show_statistics) + grid_layout.addWidget(self.resolution_spin_box, grid_row, 1, + Qt.AlignmentFlag.AlignLeft) + + grid_row += 1 + grid_layout.addWidget(QLabel('Image size (megapixel)'), grid_row, 0, + Qt.AlignmentFlag.AlignRight) + self.megapixels = QLabel('-') + grid_layout.addWidget(self.megapixels, grid_row, 1, + Qt.AlignmentFlag.AlignLeft) + + grid_row += 1 + grid_layout.addWidget(QLabel('Allow upscaling'), grid_row, 0, + Qt.AlignmentFlag.AlignRight) + self.upscaling_check_box = SettingsBigCheckBox( + key='export_upscaling', + default=DEFAULT_SETTINGS['export_upscaling']) + self.upscaling_check_box.setToolTip('Scale too small images to the requested size.\n' + 'This should be avoided as it lowers the image quality.') + self.upscaling_check_box.stateChanged.connect(self.show_statistics) + grid_layout.addWidget(self.upscaling_check_box, grid_row, 1, + Qt.AlignmentFlag.AlignLeft) + + grid_row += 1 + grid_layout.addWidget(QLabel('Bucket resolution size (px)'), grid_row, 0, + Qt.AlignmentFlag.AlignRight) + self.bucket_res_size_spin_box = SettingsSpinBox( + key='export_bucket_res_size', default=DEFAULT_SETTINGS['export_bucket_res_size'], + minimum=1, maximum=256) + self.bucket_res_size_spin_box.setToolTip('Ensure that the exported image size is divisable by that number.\n' + 'It should match the setting on the training tool.\n' + 'It might cause minor cropping.') + self.bucket_res_size_spin_box.textChanged.connect(self.show_statistics) + grid_layout.addWidget(self.bucket_res_size_spin_box, grid_row, 1, + Qt.AlignmentFlag.AlignLeft) + + grid_row += 1 + grid_layout.addWidget(QLabel('Bucket fitting strategy'), grid_row, 0, + Qt.AlignmentFlag.AlignRight) + bucket_strategy_combo_box = SettingsComboBox(key='export_bucket_strategy') + bucket_strategy_combo_box.addItems(list(BucketStrategy)) + bucket_strategy_combo_box.setToolTip('crop: center crop\n' + 'scale: assymetric scaling\n' + 'crop and scale: use both to minimize each effect') + grid_layout.addWidget(bucket_strategy_combo_box, grid_row, 1, + Qt.AlignmentFlag.AlignLeft) + + grid_row += 1 + grid_layout.addWidget(QLabel('Output format'), grid_row, 0, + Qt.AlignmentFlag.AlignRight) + format_widget = QWidget() + format_layout = QHBoxLayout() + self.format_combo_box = SettingsComboBox(key='export_format') + self.format_combo_box.addItems(list(ExportFormat)) + self.format_combo_box.currentTextChanged.connect(self.format_change) + format_layout.addWidget(self.format_combo_box, + Qt.AlignmentFlag.AlignLeft) + format_layout.addWidget(QLabel('Quality'), + Qt.AlignmentFlag.AlignRight) + self.quality_spin_box = SettingsSpinBox( + key='export_quality', default=DEFAULT_SETTINGS['export_quality'], + minimum=0, maximum=100) + self.quality_spin_box.setToolTip('Only for JPEG and WebP.\n' + '0 is worst and 100 is best.\n' + 'For JPEG numbers above 95 should be avoided') + self.quality_spin_box.textChanged.connect(self.quality_change) + format_layout.addWidget(self.quality_spin_box, + Qt.AlignmentFlag.AlignLeft) + format_widget.setLayout(format_layout) + grid_layout.addWidget(format_widget, grid_row, 1, + Qt.AlignmentFlag.AlignLeft) + + grid_row += 1 + grid_layout.addWidget(QLabel('Export directory'), grid_row, 0, + Qt.AlignmentFlag.AlignRight) + self.export_directory_line_edit = SettingsLineEdit( + key='export_directory_path', + default=DEFAULT_SETTINGS['export_directory_path']) + self.export_directory_line_edit.setMinimumWidth(400) + self.export_directory_line_edit.setClearButtonEnabled(True) + grid_layout.addWidget(self.export_directory_line_edit, grid_row, 1, + Qt.AlignmentFlag.AlignLeft) + + grid_row += 1 + export_directory_button = QPushButton('Select Directory...') + export_directory_button.clicked.connect(self.set_export_directory_path) + grid_layout.addWidget(export_directory_button, grid_row, 1, + Qt.AlignmentFlag.AlignLeft) + + grid_row += 1 + grid_layout.addWidget(QLabel('Keep input directory structure'), grid_row, 0, + Qt.AlignmentFlag.AlignRight) + keep_dir_structure_check_box = SettingsBigCheckBox( + key='export_keep_dir_structure', + default=DEFAULT_SETTINGS['export_keep_dir_structure']) + keep_dir_structure_check_box.setToolTip('Keep the subdirectory structure or export\n' + 'all images in the same export directory') + grid_layout.addWidget(keep_dir_structure_check_box, grid_row, 1, + Qt.AlignmentFlag.AlignLeft) + + grid_row += 1 + grid_layout.addWidget(QLabel('Statistics'), grid_row, 0, + Qt.AlignmentFlag.AlignRight) + self.statistics_table = QTableWidget(0, 5, self) + self.statistics_table.setHorizontalHeaderLabels(['Width', 'Height', 'Count', 'Aspect ratio', 'Size utilization']) + self.statistics_table.setMinimumWidth(400) + grid_layout.addWidget(self.statistics_table, grid_row, 1, + Qt.AlignmentFlag.AlignLeft) + + layout.addLayout(grid_layout) + export_button = QPushButton('Export') + export_button.clicked.connect(self.do_export) + layout.addWidget(export_button) + + # update display + self.apply_preset(preset_combo_box.currentText()) + self.show_megapixels() + self.inhibit_statistics_update = False + self.show_statistics() + + @Slot() + def apply_preset(self, value): + if value == 'manual': + self.resolution_spin_box.setEnabled(True) + self.bucket_res_size_spin_box.setEnabled(True) + else: + preset = Presets[value] + self.inhibit_statistics_update = True + self.resolution_spin_box.setValue(preset[0]) + self.resolution_spin_box.setEnabled(False) + self.bucket_res_size_spin_box.setValue(preset[1]) + self.bucket_res_size_spin_box.setEnabled(False) + self.inhibit_statistics_update = False + self.show_statistics() + + @Slot() + def show_megapixels(self): + resolution = self.resolution_spin_box.value() + if resolution > 0: + megapixels = resolution * resolution / 1024 / 1024 + self.megapixels.setText(f"{megapixels:.3f}") + else: + self.megapixels.setText('-') + + @Slot() + def format_change(self, export_format): + if export_format == ExportFormat.JPG: + self.quality_spin_box.setValue(75) + self.quality_spin_box.setEnabled(True) + elif export_format == ExportFormat.PNG: + self.quality_spin_box.setValue(100) + self.quality_spin_box.setEnabled(False) + elif export_format == ExportFormat.WEBP: + self.quality_spin_box.setValue(80) + self.quality_spin_box.setEnabled(True) + + @Slot() + def quality_change(self, quality): + if (self.format_combo_box.currentText() == ExportFormat.JPG) and int(quality) > 95: + self.quality_spin_box.setStyleSheet('background: orange') + else: + self.quality_spin_box.setStyleSheet('') + + @Slot() + def show_statistics(self): + if self.inhibit_statistics_update: + return + + self.resolution_cache = {} + resolution = self.resolution_spin_box.value() + upscaling = self.upscaling_check_box.isChecked() + bucket_res = self.bucket_res_size_spin_box.value() + + # notable aspect ratios + aspect_ratios = [ + (1, 1, 1), + (2, 1, 2/1), + (3, 2, 3/2), + (4, 3, 4/3), + (16, 9, 16/9), + (21, 9, 21/9), + ] + + image_dimensions = defaultdict(int) + for image_index in range(self.image_list_model.rowCount()): + this_image = self.image_list_model.index(image_index).data(Qt.ItemDataRole.UserRole) + this_image.target_dimensions = self.target_dimensions(this_image.dimensions, resolution, upscaling, bucket_res) + image_dimensions[this_image.target_dimensions] += 1 + + sorted_dimensions = sorted( + image_dimensions.items(), + key=lambda x: x[0][0] / x[0][1] # Sort by width/height ratio + ) + + self.statistics_table.setRowCount(0) # clear old data + for dimensions, count in sorted_dimensions: + width, height = dimensions + aspect_ratio = width / height + rowPosition = self.statistics_table.rowCount() + notable_aspect_ratio = '' + for ar in aspect_ratios: + if abs(ar[2] - aspect_ratio) < 1e-3: + notable_aspect_ratio = f" ({ar[0]}:{ar[1]})" + elif abs(1/ar[2] - aspect_ratio) < 1e-3: + notable_aspect_ratio = f" ({ar[1]}:{ar[0]})" + utilization = (width * height)**0.5 / resolution if resolution > 0 else 1 + + self.statistics_table.insertRow(rowPosition) + self.statistics_table.setItem(rowPosition, 0, QTableWidgetItem(str(width))) + self.statistics_table.setItem(rowPosition, 1, QTableWidgetItem(str(height))) + self.statistics_table.setItem(rowPosition, 2, QTableWidgetItem(str(count))) + self.statistics_table.setItem(rowPosition, 3, QTableWidgetItem(f"{aspect_ratio:.3f}{notable_aspect_ratio}")) + self.statistics_table.setItem(rowPosition, 4, QTableWidgetItem(f"{100*utilization:.1f}%")) + + def target_dimensions(self, dimensions, resolution, upscaling, bucket_res): + """ + Given the original width and height, the bucket resolution step size, + and a maximum allowed area, return new dimensions (width, height) + where both dimensions are multiples of `bucket_res`, their product + does not exceed resolution**2, and the new aspect ratio (width/height) + is as close as possible to the original aspect ratio. + + Note: this gives the optimal answer and thus can be slower than the Kohya bucket + algorithm + """ + if resolution == 0: + # no rescale in this case, only cropping + return ((dimensions[0] // bucket_res)*bucket_res, (dimensions[1] // bucket_res)*bucket_res) + + if dimensions in self.resolution_cache: + return self.resolution_cache[dimensions] + + max_area = resolution**2 + + # Compute the original aspect ratio. + target_ratio = dimensions[0] / dimensions[1] + + # The maximum allowed product of multipliers. + T = max_area // (bucket_res * bucket_res) + + best_candidate = None # will hold (new_width, new_height, error, area) + + # Loop over possible values for b (the vertical multiplier). + # We choose b from 1 up to T (although many values will be skipped + # because the corresponding a then makes a * b > T). + for b in range(1, T + 1): + # Choose a so that a / b is as close as possible to target_ratio. + # (We round the ideal value a = target_ratio * b to the nearest integer.) + a = round(target_ratio * b) + if a < 1: + a = 1 # ensure at least bucket_res pixels + + # Check that the candidate image area (in multiplier units) does not exceed T. + if a * b > T: + # If a*b is too big, skip the candidate. + continue + + candidate_width = a * bucket_res + candidate_height = b * bucket_res + candidate_area = candidate_width * candidate_height + + if not upscaling and (candidate_width > dimensions[0] or candidate_height > dimensions[1]): + continue + + # Compute the aspect ratio error. + candidate_ratio = a / b + error = abs(candidate_ratio - target_ratio) + # compute the mean squared error of ratio and normalized maximum size + error = (candidate_ratio - target_ratio)**2 + ((max_area-candidate_area)/max_area)**2 + + # We choose the candidate with the lowest error. In case of a tie, we choose + # the one that uses the largest area (i.e. as close as possible to resolution**2). + if best_candidate is None: + best_candidate = (candidate_width, candidate_height, error, candidate_area) + else: + _, _, best_error, best_area = best_candidate + if (error < best_error) or (abs(error - best_error) < 1e-9 and candidate_area > best_area): + best_candidate = (candidate_width, candidate_height, error, candidate_area) + + # Fallback: if no candidate is found (this shouldn't happen for reasonable values), + # simply return the smallest possible image. + if best_candidate is None: + return bucket_res, bucket_res + else: + new_width, new_height, _, _ = best_candidate + self.resolution_cache[dimensions] = (new_width, new_height) + return new_width, new_height + + @Slot() + def set_export_directory_path(self): + export_directory_path = self.settings.value( + 'export_directory_path', + defaultValue=DEFAULT_SETTINGS['export_directory_path'], type=str) + if export_directory_path: + initial_directory_path = export_directory_path + elif self.settings.contains('directory_path'): + initial_directory_path = self.settings.value('directory_path') + else: + initial_directory_path = '' + export_directory_path = QFileDialog.getExistingDirectory( + parent=self, caption='Select directory for image export', + dir=initial_directory_path) + if export_directory_path: + self.export_directory_line_edit.setText(export_directory_path) + + @Slot() + def do_export(self): + directory_path = self.settings.value('directory_path', type=str) + export_directory_path = Path(self.settings.value('export_directory_path', type=str)) + export_keep_dir_structure = self.settings.value('export_keep_dir_structure', type=bool) + no_overwrite = True + if os.path.exists(export_directory_path): + if os.path.isfile(export_directory_path): + QMessageBox.critical( + self, + 'Path error', + 'The export directory path points to a file and not to a directory' + ) + return + if os.listdir(export_directory_path): + msgBox = QMessageBox() + msgBox.setIcon(QMessageBox.Warning) + msgBox.setWindowTitle('Path warning') + msgBox.setText('The export directory path is not empty') + overwrite_button = msgBox.addButton('Overwrite', QMessageBox.YesRole) + rename_button = msgBox.addButton('Rename', QMessageBox.NoRole) + msgBox.addButton(QMessageBox.Cancel) + msgBox.setDefaultButton(QMessageBox.Cancel) + button = msgBox.exec_() + if button == QMessageBox.Cancel: + return + if msgBox.clickedButton() == overwrite_button: + no_overwrite = False + else: + QMessageBox.critical( + self, + 'Path error', + 'The export directory path does not exist' + ) + return + + resolution = self.resolution_spin_box.value() + upscaling = self.upscaling_check_box.isChecked() + bucket_res = self.bucket_res_size_spin_box.value() + export_format = self.format_combo_box.currentText() + quality = self.quality_spin_box.value() + bucket_strategy = self.settings.value('export_bucket_strategy', type=str) + + for image_index in range(self.image_list_model.rowCount()): + image_entry = self.image_list_model.index(image_index).data(Qt.ItemDataRole.UserRole) + if export_keep_dir_structure: + relative_path = image_entry.path.relative_to(directory_path) + export_path = export_directory_path / relative_path + export_path.parent.mkdir(parents=True, exist_ok=True) + else: + export_path = export_directory_path / image_entry.path.name + export_path = export_path.with_suffix(export_format.split(' ', 1)[0]) + + if no_overwrite: + stem = export_path.stem + counter = 0 + while export_path.exists(): + export_path = export_path.parent / f"{stem}_{counter}{export_path.suffix}" + counter += 1 + + image_file = Image.open(image_entry.path) + # Preserve alpha if present: + if image_file.mode in ("RGBA", "LA", "PA") and not export_format == ExportFormat.JPG: # Check for alpha channels + image_file = image_file.convert("RGBA") + else: + image_file = image_file.convert("RGB") # Otherwise, convert to RGB + + new_width, new_height = image_entry.target_dimensions + current_width, current_height = image_file.size + if bucket_strategy == BucketStrategy.CROP or bucket_strategy == BucketStrategy.CROP_SCALE: + if current_height * new_width / current_width < new_height: # too wide + new_width = floor(current_width * new_height / current_height) + else: # too high + new_height = floor(current_height * new_width / current_width) + if bucket_strategy == BucketStrategy.CROP_SCALE: + new_width = floor((image_entry.target_dimensions[0] + new_width)/2) + new_height = floor((image_entry.target_dimensions[1] + new_height)/2) + if image_file.size[0] != new_width or image_file.size[1] != new_height: + # resize with the best method available + resized_image = image_file.resize((new_width, new_height), Image.LANCZOS) + # followed by a slight sharpening as it should be done + sharpend_image = resized_image.filter(ImageFilter.UnsharpMask(radius = 0.5, percent = 100, threshold = 3)) + else: + sharpend_image = image_file + + # crop to the desired size + current_width, current_height = sharpend_image.size + crop_width = floor((current_width - image_entry.target_dimensions[0]) / 2) + crop_height = floor((current_height - image_entry.target_dimensions[1]) / 2) + cropped_image = sharpend_image.crop((crop_width, crop_height, current_width - crop_width, current_height - crop_height)) + + # TODO: apply color management as the input images might have an + # arbitrary color space that doesn't fit to the expectation + # of the trainer (most likely sRGB, but probably something + # like Rec. 2100 with HDR in the future?) + #final_image = ImageCms.profileToProfile(cropped_image, srgb_profile, profile) + final_image = cropped_image + + final_image.save(export_path, format=ExportFormatDict[export_format], quality=quality, icc_profile=final_image.info.get('icc_profile')) + self.close() diff --git a/taggui/utils/image.py b/taggui/utils/image.py index f2637da4..96554486 100644 --- a/taggui/utils/image.py +++ b/taggui/utils/image.py @@ -8,5 +8,6 @@ class Image: path: Path dimensions: tuple[int, int] | None + target_dimensions: tuple[int, int] | None tags: list[str] = field(default_factory=list) thumbnail: QIcon | None = None diff --git a/taggui/utils/settings.py b/taggui/utils/settings.py index 88e1f0f5..efe94440 100644 --- a/taggui/utils/settings.py +++ b/taggui/utils/settings.py @@ -9,7 +9,16 @@ 'tag_separator': ',', 'insert_space_after_tag_separator': True, 'autocomplete_tags': True, - 'models_directory_path': '' + 'models_directory_path': '', + 'export_preset': 'SDXL, SD3, Flux', + 'export_resolution': 1024, + 'export_upscaling': False, + 'export_bucket_res_size': 64, + 'export_bucket_strategy': 'crop', + 'export_format': '.jpg - JPEG', + 'export_quality': 75, + 'export_directory_path': '', + 'export_keep_dir_structure': False } diff --git a/taggui/widgets/main_window.py b/taggui/widgets/main_window.py index 78a8cf98..87533a8f 100644 --- a/taggui/widgets/main_window.py +++ b/taggui/widgets/main_window.py @@ -10,6 +10,7 @@ from dialogs.batch_reorder_tags_dialog import BatchReorderTagsDialog from dialogs.find_and_replace_dialog import FindAndReplaceDialog +from dialogs.export_dialog import ExportDialog from dialogs.settings_dialog import SettingsDialog from models.image_list_model import ImageListModel from models.image_tag_list_model import ImageTagListModel @@ -255,6 +256,12 @@ def reload_directory(self): self.image_list.list_view.setCurrentIndex( self.proxy_image_list_model.index(select_index, 0)) + @Slot() + def export_images_dialog(self): + export_dialog = ExportDialog(parent=self, image_list_model=self.image_list_model) + export_dialog.exec() + return + @Slot() def show_settings_dialog(self): settings_dialog = SettingsDialog(parent=self) @@ -313,6 +320,9 @@ def create_menus(self): [QKeySequence('Ctrl+Shift+L'), QKeySequence('F5')]) self.reload_directory_action.triggered.connect(self.reload_directory) file_menu.addAction(self.reload_directory_action) + export_action = QAction('Export...', parent=self) + export_action.triggered.connect(self.export_images_dialog) + file_menu.addAction(export_action) settings_action = QAction('Settings...', parent=self) settings_action.setShortcut(QKeySequence('Ctrl+Alt+S')) settings_action.triggered.connect(self.show_settings_dialog) From 8259d06b4a02b6b274ff09ba4569f3b9645f0c68 Mon Sep 17 00:00:00 2001 From: StableLlama Date: Tue, 4 Feb 2025 21:24:35 +0100 Subject: [PATCH 02/18] Add color space conversion --- taggui/dialogs/export_dialog.py | 57 +++++++++++++++++++++++++++------ taggui/utils/settings.py | 1 + 2 files changed, 49 insertions(+), 9 deletions(-) diff --git a/taggui/dialogs/export_dialog.py b/taggui/dialogs/export_dialog.py index 55d297aa..ff54b327 100644 --- a/taggui/dialogs/export_dialog.py +++ b/taggui/dialogs/export_dialog.py @@ -2,14 +2,16 @@ from collections import defaultdict from math import floor import os +import io from pathlib import Path from PySide6.QtCore import Qt, Slot +from PySide6.QtGui import QColorSpace from PySide6.QtWidgets import (QWidget, QDialog, QFileDialog, QGridLayout, QLabel, QLineEdit, QPushButton, QVBoxLayout, QHBoxLayout, QTableWidget, QTableWidgetItem, QSizePolicy, QMessageBox) -from PIL import Image, ImageFilter #, ImageQt, ImageEnhance, ImageCms +from PIL import Image, ImageFilter, ImageCms from utils.settings import DEFAULT_SETTINGS, get_settings from utils.settings_widgets import (SettingsBigCheckBox, SettingsLineEdit, @@ -34,6 +36,17 @@ class ExportFormat(str, Enum): ExportFormat.WEBP: 'webp' } +class IccProfileList(str, Enum): + SRgb = 'sRGB' + SRgbLinear = 'sRGB (linear gamma)' + AdobeRgb = 'AdobeRGB' + DisplayP3 = 'DisplayP3' + ProPhotoRgb = 'ProPhotoRGB' + # since PySide6.8: + Bt2020 = 'BT.2020' + Bt2100Pq = 'BT.2100(PQ)' + Bt2100Hlg = 'BT.2100 (HLG)' + class BucketStrategy(str, Enum): CROP = 'crop' SCALE = 'scale' @@ -147,6 +160,21 @@ def __init__(self, parent, image_list_model: ImageListModel): grid_layout.addWidget(format_widget, grid_row, 1, Qt.AlignmentFlag.AlignLeft) + grid_row += 1 + grid_layout.addWidget(QLabel('Output color space'), grid_row, 0, + Qt.AlignmentFlag.AlignRight) + color_space_combo_box = SettingsComboBox(key='export_color_space') + color_space_combo_box.addItem("feed through (don't touch)") + color_space_combo_box.addItem('sRGB (implicit, without profile)') + color_space_combo_box.addItems([IccProfileList[e.name] for e in QColorSpace.NamedColorSpace]) + color_space_combo_box.setToolTip('Color space of the exported images.\n' + 'Most likely the trainer expects sRGB!\n' + '\n' + 'Use "feed through" to keep the color space as it is.\n' + 'Use "sRGB (implicit, without profile)" to save in sRGB but don\'t embed the ICC profile to save 8k file size.') + grid_layout.addWidget(color_space_combo_box, grid_row, 1, + Qt.AlignmentFlag.AlignLeft) + grid_row += 1 grid_layout.addWidget(QLabel('Export directory'), grid_row, 0, Qt.AlignmentFlag.AlignRight) @@ -421,6 +449,11 @@ def do_export(self): bucket_res = self.bucket_res_size_spin_box.value() export_format = self.format_combo_box.currentText() quality = self.quality_spin_box.value() + color_space = self.settings.value('export_color_space', type=str) + save_profile = True + if color_space == 'sRGB (implicit, without profile)': + color_space = 'sRGB' + save_profile = False bucket_strategy = self.settings.value('export_bucket_strategy', type=str) for image_index in range(self.image_list_model.rowCount()): @@ -471,12 +504,18 @@ def do_export(self): crop_height = floor((current_height - image_entry.target_dimensions[1]) / 2) cropped_image = sharpend_image.crop((crop_width, crop_height, current_width - crop_width, current_height - crop_height)) - # TODO: apply color management as the input images might have an - # arbitrary color space that doesn't fit to the expectation - # of the trainer (most likely sRGB, but probably something - # like Rec. 2100 with HDR in the future?) - #final_image = ImageCms.profileToProfile(cropped_image, srgb_profile, profile) - final_image = cropped_image - - final_image.save(export_path, format=ExportFormatDict[export_format], quality=quality, icc_profile=final_image.info.get('icc_profile')) + if color_space == "feed through (don't touch)": + cropped_image.save(export_path, format=ExportFormatDict[export_format], quality=quality, icc_profile=cropped_image.info.get('icc_profile')) + else: + source_profile_raw = image_file.info.get('icc_profile') + if source_profile_raw is None: # assume sRGB + source_profile_raw = QColorSpace(QColorSpace.SRgb).iccProfile() + source_profile = ImageCms.ImageCmsProfile(io.BytesIO(source_profile_raw)) + target_profile_raw = QColorSpace(getattr(QColorSpace, IccProfileList(color_space).name)).iccProfile() + target_profile = ImageCms.ImageCmsProfile(io.BytesIO(target_profile_raw)) + final_image = ImageCms.profileToProfile(cropped_image, source_profile, target_profile) + if save_profile: + final_image.save(export_path, format=ExportFormatDict[export_format], quality=quality, icc_profile=target_profile.tobytes()) + else: + final_image.save(export_path, format=ExportFormatDict[export_format], quality=quality, icc_profile=None) self.close() diff --git a/taggui/utils/settings.py b/taggui/utils/settings.py index efe94440..518e8c46 100644 --- a/taggui/utils/settings.py +++ b/taggui/utils/settings.py @@ -17,6 +17,7 @@ 'export_bucket_strategy': 'crop', 'export_format': '.jpg - JPEG', 'export_quality': 75, + 'export_color_space': 'sRGB', 'export_directory_path': '', 'export_keep_dir_structure': False } From ab5be1af280b8b3c9c2feb1dec8083b539140ad8 Mon Sep 17 00:00:00 2001 From: StableLlama Date: Tue, 4 Feb 2025 23:47:22 +0100 Subject: [PATCH 03/18] Little display fixes and add infrastructure for preferred sizes. --- taggui/dialogs/export_dialog.py | 75 +++++++++++++++++++++------------ taggui/utils/settings.py | 3 +- 2 files changed, 49 insertions(+), 29 deletions(-) diff --git a/taggui/dialogs/export_dialog.py b/taggui/dialogs/export_dialog.py index ff54b327..8a12f6ff 100644 --- a/taggui/dialogs/export_dialog.py +++ b/taggui/dialogs/export_dialog.py @@ -1,8 +1,8 @@ from enum import Enum from collections import defaultdict -from math import floor import os import io +from math import floor from pathlib import Path from PySide6.QtCore import Qt, Slot @@ -19,10 +19,10 @@ from models.image_list_model import ImageListModel Presets = { - 'manual': (0, 0), - 'Direct feed through': (0, 1), - 'SD1': (512, 64), - 'SDXL, SD3, Flux': (1024, 64) + 'manual': (0, 0, '1:1, 2:1, 3:2, 4:3, 16:9, 21:9'), + 'Direct feed through': (0, 1, '1:1, 2:1, 3:2, 4:3, 16:9, 21:9'), + 'SD1': (512, 64, '512:512, 640:320, 576:384, 512:384, 640:384, 704:320'), + 'SDXL, SD3, Flux': (1024, 64, '1024:1024, 1408:704, 1216:832, 1152:896, 1344:768, 1536:640') } class ExportFormat(str, Enum): @@ -99,18 +99,6 @@ def __init__(self, parent, image_list_model: ImageListModel): grid_layout.addWidget(self.megapixels, grid_row, 1, Qt.AlignmentFlag.AlignLeft) - grid_row += 1 - grid_layout.addWidget(QLabel('Allow upscaling'), grid_row, 0, - Qt.AlignmentFlag.AlignRight) - self.upscaling_check_box = SettingsBigCheckBox( - key='export_upscaling', - default=DEFAULT_SETTINGS['export_upscaling']) - self.upscaling_check_box.setToolTip('Scale too small images to the requested size.\n' - 'This should be avoided as it lowers the image quality.') - self.upscaling_check_box.stateChanged.connect(self.show_statistics) - grid_layout.addWidget(self.upscaling_check_box, grid_row, 1, - Qt.AlignmentFlag.AlignLeft) - grid_row += 1 grid_layout.addWidget(QLabel('Bucket resolution size (px)'), grid_row, 0, Qt.AlignmentFlag.AlignRight) @@ -124,6 +112,30 @@ def __init__(self, parent, image_list_model: ImageListModel): grid_layout.addWidget(self.bucket_res_size_spin_box, grid_row, 1, Qt.AlignmentFlag.AlignLeft) + grid_row += 1 + grid_layout.addWidget(QLabel('Prefered sizes'), grid_row, 0, + Qt.AlignmentFlag.AlignRight) + self.preferred_sizes_line_edit = SettingsLineEdit( + key='export_preferred_sizes', + default=DEFAULT_SETTINGS['export_preferred_sizes']) + self.preferred_sizes_line_edit.setMinimumWidth(500) + self.preferred_sizes_line_edit.setToolTip('Comma separated list of preferred sizes and aspect ratios.\n' + "The inverse aspect ratio is automatically derived and doesn't need to be included.") + grid_layout.addWidget(self.preferred_sizes_line_edit, grid_row, 1, + Qt.AlignmentFlag.AlignLeft) + + grid_row += 1 + grid_layout.addWidget(QLabel('Allow upscaling'), grid_row, 0, + Qt.AlignmentFlag.AlignRight) + self.upscaling_check_box = SettingsBigCheckBox( + key='export_upscaling', + default=DEFAULT_SETTINGS['export_upscaling']) + self.upscaling_check_box.setToolTip('Scale too small images to the requested size.\n' + 'This should be avoided as it lowers the image quality.') + self.upscaling_check_box.stateChanged.connect(self.show_statistics) + grid_layout.addWidget(self.upscaling_check_box, grid_row, 1, + Qt.AlignmentFlag.AlignLeft) + grid_row += 1 grid_layout.addWidget(QLabel('Bucket fitting strategy'), grid_row, 0, Qt.AlignmentFlag.AlignRight) @@ -140,6 +152,9 @@ def __init__(self, parent, image_list_model: ImageListModel): Qt.AlignmentFlag.AlignRight) format_widget = QWidget() format_layout = QHBoxLayout() + format_layout.setContentsMargins(0, 0, 0, 0) + current_format = self.settings.value('export_format', type=str) + current_quality = self.settings.value('export_quality', type=str) self.format_combo_box = SettingsComboBox(key='export_format') self.format_combo_box.addItems(list(ExportFormat)) self.format_combo_box.currentTextChanged.connect(self.format_change) @@ -159,6 +174,9 @@ def __init__(self, parent, image_list_model: ImageListModel): format_widget.setLayout(format_layout) grid_layout.addWidget(format_widget, grid_row, 1, Qt.AlignmentFlag.AlignLeft) + # ensure correct enable/disable and background color of the quality + self.format_change(current_format, False) + self.quality_change(current_quality) grid_row += 1 grid_layout.addWidget(QLabel('Output color space'), grid_row, 0, @@ -181,7 +199,7 @@ def __init__(self, parent, image_list_model: ImageListModel): self.export_directory_line_edit = SettingsLineEdit( key='export_directory_path', default=DEFAULT_SETTINGS['export_directory_path']) - self.export_directory_line_edit.setMinimumWidth(400) + self.export_directory_line_edit.setMinimumWidth(500) self.export_directory_line_edit.setClearButtonEnabled(True) grid_layout.addWidget(self.export_directory_line_edit, grid_row, 1, Qt.AlignmentFlag.AlignLeft) @@ -208,7 +226,7 @@ def __init__(self, parent, image_list_model: ImageListModel): Qt.AlignmentFlag.AlignRight) self.statistics_table = QTableWidget(0, 5, self) self.statistics_table.setHorizontalHeaderLabels(['Width', 'Height', 'Count', 'Aspect ratio', 'Size utilization']) - self.statistics_table.setMinimumWidth(400) + self.statistics_table.setMinimumWidth(500) grid_layout.addWidget(self.statistics_table, grid_row, 1, Qt.AlignmentFlag.AlignLeft) @@ -218,25 +236,26 @@ def __init__(self, parent, image_list_model: ImageListModel): layout.addWidget(export_button) # update display - self.apply_preset(preset_combo_box.currentText()) + self.apply_preset(preset_combo_box.currentText(), False) self.show_megapixels() self.inhibit_statistics_update = False self.show_statistics() @Slot() - def apply_preset(self, value): + def apply_preset(self, value, do_value_change = True): + preset = Presets[value] if value == 'manual': self.resolution_spin_box.setEnabled(True) self.bucket_res_size_spin_box.setEnabled(True) else: - preset = Presets[value] self.inhibit_statistics_update = True - self.resolution_spin_box.setValue(preset[0]) + self.resolution_spin_box.setValue(preset[0]) if do_value_change else 0 self.resolution_spin_box.setEnabled(False) - self.bucket_res_size_spin_box.setValue(preset[1]) + self.bucket_res_size_spin_box.setValue(preset[1]) if do_value_change else 0 self.bucket_res_size_spin_box.setEnabled(False) self.inhibit_statistics_update = False self.show_statistics() + self.preferred_sizes_line_edit.setText(preset[2]) if do_value_change else 0 @Slot() def show_megapixels(self): @@ -248,15 +267,15 @@ def show_megapixels(self): self.megapixels.setText('-') @Slot() - def format_change(self, export_format): + def format_change(self, export_format, do_value_change = True): if export_format == ExportFormat.JPG: - self.quality_spin_box.setValue(75) + self.quality_spin_box.setValue(75) if do_value_change else 0 self.quality_spin_box.setEnabled(True) elif export_format == ExportFormat.PNG: - self.quality_spin_box.setValue(100) + self.quality_spin_box.setValue(100) if do_value_change else 0 self.quality_spin_box.setEnabled(False) elif export_format == ExportFormat.WEBP: - self.quality_spin_box.setValue(80) + self.quality_spin_box.setValue(80) if do_value_change else 0 self.quality_spin_box.setEnabled(True) @Slot() diff --git a/taggui/utils/settings.py b/taggui/utils/settings.py index 518e8c46..d256d2c4 100644 --- a/taggui/utils/settings.py +++ b/taggui/utils/settings.py @@ -12,8 +12,9 @@ 'models_directory_path': '', 'export_preset': 'SDXL, SD3, Flux', 'export_resolution': 1024, - 'export_upscaling': False, 'export_bucket_res_size': 64, + 'export_preferred_sizes' : '1024:1024, 1408:704, 1216:832, 1152:896, 1344:768, 1536:640', + 'export_upscaling': False, 'export_bucket_strategy': 'crop', 'export_format': '.jpg - JPEG', 'export_quality': 75, From 476e218cef61e86f53e18c5ae446b1345617b339 Mon Sep 17 00:00:00 2001 From: StableLlama Date: Wed, 5 Feb 2025 22:44:05 +0100 Subject: [PATCH 04/18] Change algorithm for bucketing --- taggui/dialogs/export_dialog.py | 149 +++++++++++++++++++------------- 1 file changed, 90 insertions(+), 59 deletions(-) diff --git a/taggui/dialogs/export_dialog.py b/taggui/dialogs/export_dialog.py index 8a12f6ff..63ee2e41 100644 --- a/taggui/dialogs/export_dialog.py +++ b/taggui/dialogs/export_dialog.py @@ -2,7 +2,8 @@ from collections import defaultdict import os import io -from math import floor +import re +from math import floor, sqrt from pathlib import Path from PySide6.QtCore import Qt, Slot @@ -21,7 +22,7 @@ Presets = { 'manual': (0, 0, '1:1, 2:1, 3:2, 4:3, 16:9, 21:9'), 'Direct feed through': (0, 1, '1:1, 2:1, 3:2, 4:3, 16:9, 21:9'), - 'SD1': (512, 64, '512:512, 640:320, 576:384, 512:384, 640:384, 704:320'), + 'SD1': (512, 64, '512:512, 640:320, 576:384, 512:384, 640:384, 768:320'), 'SDXL, SD3, Flux': (1024, 64, '1024:1024, 1408:704, 1216:832, 1152:896, 1344:768, 1536:640') } @@ -113,7 +114,7 @@ def __init__(self, parent, image_list_model: ImageListModel): Qt.AlignmentFlag.AlignLeft) grid_row += 1 - grid_layout.addWidget(QLabel('Prefered sizes'), grid_row, 0, + grid_layout.addWidget(QLabel('Preferred sizes'), grid_row, 0, Qt.AlignmentFlag.AlignRight) self.preferred_sizes_line_edit = SettingsLineEdit( key='export_preferred_sizes', @@ -304,6 +305,25 @@ def show_statistics(self): (16, 9, 16/9), (21, 9, 21/9), ] + self.preferred_sizes = [] + for res_str in re.split(r'\s*,\s*', self.settings.value('export_preferred_sizes')): + try: + size_str = res_str.split(':') + width = max(int(size_str[0]), int(size_str[1])) + height = min(int(size_str[0]), int(size_str[1])) + self.preferred_sizes.append((width, height)) + if not width == height: + self.preferred_sizes.append((height, width)) + # add exact aspect ratio of the preferred size to label it similar to the perfect one + aspect_ratio = width / height + for ar in aspect_ratios: + if abs(ar[2] - aspect_ratio) < 0.15: + aspect_ratios.append((ar[0], ar[1], aspect_ratio)) + break + except ValueError: + # Handle cases where the resolution string is not in the correct format + print(f"Warning: Invalid resolution format: {res_str}. Skipping.") + continue # Skip to the next resolution if there's an error image_dimensions = defaultdict(int) for image_index in range(self.image_list_model.rowCount()): @@ -347,68 +367,79 @@ def target_dimensions(self, dimensions, resolution, upscaling, bucket_res): Note: this gives the optimal answer and thus can be slower than the Kohya bucket algorithm """ + width, height = dimensions if resolution == 0: # no rescale in this case, only cropping - return ((dimensions[0] // bucket_res)*bucket_res, (dimensions[1] // bucket_res)*bucket_res) + return ((width // bucket_res)*bucket_res, (height // bucket_res)*bucket_res) + + if width < bucket_res or height < bucket_res: + # it doesn't make sense to use such a small image. But we shouldn't + # patronize the user + return dimensions if dimensions in self.resolution_cache: return self.resolution_cache[dimensions] - max_area = resolution**2 - - # Compute the original aspect ratio. - target_ratio = dimensions[0] / dimensions[1] - - # The maximum allowed product of multipliers. - T = max_area // (bucket_res * bucket_res) - - best_candidate = None # will hold (new_width, new_height, error, area) - - # Loop over possible values for b (the vertical multiplier). - # We choose b from 1 up to T (although many values will be skipped - # because the corresponding a then makes a * b > T). - for b in range(1, T + 1): - # Choose a so that a / b is as close as possible to target_ratio. - # (We round the ideal value a = target_ratio * b to the nearest integer.) - a = round(target_ratio * b) - if a < 1: - a = 1 # ensure at least bucket_res pixels - - # Check that the candidate image area (in multiplier units) does not exceed T. - if a * b > T: - # If a*b is too big, skip the candidate. - continue - - candidate_width = a * bucket_res - candidate_height = b * bucket_res - candidate_area = candidate_width * candidate_height - - if not upscaling and (candidate_width > dimensions[0] or candidate_height > dimensions[1]): - continue - - # Compute the aspect ratio error. - candidate_ratio = a / b - error = abs(candidate_ratio - target_ratio) - # compute the mean squared error of ratio and normalized maximum size - error = (candidate_ratio - target_ratio)**2 + ((max_area-candidate_area)/max_area)**2 - - # We choose the candidate with the lowest error. In case of a tie, we choose - # the one that uses the largest area (i.e. as close as possible to resolution**2). - if best_candidate is None: - best_candidate = (candidate_width, candidate_height, error, candidate_area) - else: - _, _, best_error, best_area = best_candidate - if (error < best_error) or (abs(error - best_error) < 1e-9 and candidate_area > best_area): - best_candidate = (candidate_width, candidate_height, error, candidate_area) - - # Fallback: if no candidate is found (this shouldn't happen for reasonable values), - # simply return the smallest possible image. - if best_candidate is None: - return bucket_res, bucket_res - else: - new_width, new_height, _, _ = best_candidate - self.resolution_cache[dimensions] = (new_width, new_height) - return new_width, new_height + preferred_sizes_bonus = 0.4 # reduce the loss by this factor + + max_pixels = resolution * resolution + opt_width = resolution * sqrt(width/height) + opt_height = resolution * sqrt(height/width) + if not upscaling: + opt_width = min(width, opt_width) + opt_height = min(height, opt_height) + + # test 1, guaranteed to find a solution: shrink and crop + # 1.1: exact width + candidate_width = (opt_width // bucket_res) * bucket_res + candidate_height = ((height * candidate_width / width) // bucket_res) * bucket_res + loss = ((height * candidate_width / width) - candidate_height) * candidate_width + if (candidate_width, candidate_height) in self.preferred_sizes: + loss *= preferred_sizes_bonus + # 1.2: exact height + test_height = (opt_height // bucket_res) * bucket_res + test_width = ((width * test_height / height) // bucket_res) * bucket_res + test_loss = ((width * test_height / height) - test_width) * test_height + if (test_height, test_width) in self.preferred_sizes: + test_loss *= preferred_sizes_bonus + if test_loss < loss: + candidate_width = test_width + candidate_height = test_height + loss = test_loss + + # test 2, going bigger might still fit in the size budget due to cropping + # 2.1: exact width + for delta in range(1, 10): + test_width = (opt_width // bucket_res + delta) * bucket_res + test_height = ((height * test_width / width) // bucket_res) * bucket_res + if test_width * test_height > max_pixels: + break + if (test_width > width or test_height > height) and not upscaling: + break + test_loss = ((height * test_width / width) - test_height) * test_width + if (test_height, test_width) in self.preferred_sizes: + test_loss *= preferred_sizes_bonus + if test_loss < loss: + candidate_width = test_width + candidate_height = test_height + loss = test_loss + # 2.2: exact height + for delta in range(1, 10): + test_height = (opt_height // bucket_res + delta) * bucket_res + test_width = ((width * test_height / height) // bucket_res) * bucket_res + if test_width * test_height > max_pixels: + break + if (test_width > width or test_height > height) and not upscaling: + break + test_loss = ((width * test_height / height) - test_width) * test_height + if (test_height, test_width) in self.preferred_sizes: + test_loss *= preferred_sizes_bonus + if test_loss < loss: + candidate_width = test_width + candidate_height = test_height + loss = test_loss + + return int(candidate_width), int(candidate_height) @Slot() def set_export_directory_path(self): From 465e2fe6b7042788697724363d894d0ba0bd1280 Mon Sep 17 00:00:00 2001 From: StableLlama Date: Wed, 5 Feb 2025 23:18:33 +0100 Subject: [PATCH 05/18] Code documentation --- taggui/dialogs/export_dialog.py | 58 ++++++++++++++++++++++++++------- 1 file changed, 46 insertions(+), 12 deletions(-) diff --git a/taggui/dialogs/export_dialog.py b/taggui/dialogs/export_dialog.py index 63ee2e41..42e793f8 100644 --- a/taggui/dialogs/export_dialog.py +++ b/taggui/dialogs/export_dialog.py @@ -55,6 +55,9 @@ class BucketStrategy(str, Enum): class ExportDialog(QDialog): def __init__(self, parent, image_list_model: ImageListModel): + """ + Main method to create the export dialog. + """ super().__init__(parent) self.image_list_model = image_list_model self.settings = get_settings() @@ -243,7 +246,11 @@ def __init__(self, parent, image_list_model: ImageListModel): self.show_statistics() @Slot() - def apply_preset(self, value, do_value_change = True): + def apply_preset(self, value: str, do_value_change: bool = True): + """ + Slot to call when a new preset was selected to help the user to set + important settings to a consistent state. + """ preset = Presets[value] if value == 'manual': self.resolution_spin_box.setEnabled(True) @@ -260,6 +267,10 @@ def apply_preset(self, value, do_value_change = True): @Slot() def show_megapixels(self): + """ + Slot to call when the resolution was changes to update the megapixel + display. + """ resolution = self.resolution_spin_box.value() if resolution > 0: megapixels = resolution * resolution / 1024 / 1024 @@ -268,7 +279,10 @@ def show_megapixels(self): self.megapixels.setText('-') @Slot() - def format_change(self, export_format, do_value_change = True): + def format_change(self, export_format: ExportFormat, do_value_change: bool = True): + """ + Slot to call when the export format was changed. + """ if export_format == ExportFormat.JPG: self.quality_spin_box.setValue(75) if do_value_change else 0 self.quality_spin_box.setEnabled(True) @@ -280,7 +294,10 @@ def format_change(self, export_format, do_value_change = True): self.quality_spin_box.setEnabled(True) @Slot() - def quality_change(self, quality): + def quality_change(self, quality: str): + """ + Slot to call when the export quality setting was changed. + """ if (self.format_combo_box.currentText() == ExportFormat.JPG) and int(quality) > 95: self.quality_spin_box.setStyleSheet('background: orange') else: @@ -288,6 +305,9 @@ def quality_change(self, quality): @Slot() def show_statistics(self): + """ + Update the statistics table content. + """ if self.inhibit_statistics_update: return @@ -356,16 +376,24 @@ def show_statistics(self): self.statistics_table.setItem(rowPosition, 3, QTableWidgetItem(f"{aspect_ratio:.3f}{notable_aspect_ratio}")) self.statistics_table.setItem(rowPosition, 4, QTableWidgetItem(f"{100*utilization:.1f}%")) - def target_dimensions(self, dimensions, resolution, upscaling, bucket_res): + def target_dimensions(self, dimensions: tuple[int, int], resolution: int, upscaling: bool, bucket_res: int): """ - Given the original width and height, the bucket resolution step size, - and a maximum allowed area, return new dimensions (width, height) - where both dimensions are multiples of `bucket_res`, their product - does not exceed resolution**2, and the new aspect ratio (width/height) - is as close as possible to the original aspect ratio. - - Note: this gives the optimal answer and thus can be slower than the Kohya bucket - algorithm + Determine the dimensions of an image it should have when it is exported. + + Note: this gives the optimal answer and thus can be slower than the Kohya + bucket algorithm. + + Parameters + ---------- + dimensions : tuple[int, int] + The width and height of the image + resolution : int + The target resolution of the AI model. The target image pixels + will not exceed the square of this number + upscaling : bool + Is upscaling of images allowed? + bucket_res : int + The resolution of the buckets """ width, height = dimensions if resolution == 0: @@ -443,6 +471,9 @@ def target_dimensions(self, dimensions, resolution, upscaling, bucket_res): @Slot() def set_export_directory_path(self): + """ + Set the path of the directory to export to. + """ export_directory_path = self.settings.value( 'export_directory_path', defaultValue=DEFAULT_SETTINGS['export_directory_path'], type=str) @@ -460,6 +491,9 @@ def set_export_directory_path(self): @Slot() def do_export(self): + """ + Export all images with the configured settings. + """ directory_path = self.settings.value('directory_path', type=str) export_directory_path = Path(self.settings.value('export_directory_path', type=str)) export_keep_dir_structure = self.settings.value('export_keep_dir_structure', type=bool) From 0651748067f6dcd74f373503bdb960e34d0f37a9 Mon Sep 17 00:00:00 2001 From: StableLlama Date: Wed, 5 Feb 2025 23:39:10 +0100 Subject: [PATCH 06/18] Add documentation --- README.md | 55 +++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 55 insertions(+) diff --git a/README.md b/README.md index c589a5be..961f6bdb 100644 --- a/README.md +++ b/README.md @@ -17,6 +17,7 @@ like Stable Diffusion. Tagger, and many more - Batch tag operations for renaming, deleting, and sorting tags - Advanced image list filtering +- Export images ready to be used for training ## Installation @@ -253,3 +254,57 @@ You can nest parentheses and operators to create arbitrarily complex filters. The `Edit` menu contains additional features for batch tag operations, such as `Find and Replace` (`Ctrl`+`R`) and `Batch Reorder Tags` (`Ctrl`+`B`). + +## Export + +Exporting the images to a directory allows different options. By choosing the +preset for the target AI model many important settings are automatically set. + +Resolution +: The native resolution of the model, like 1024 for SDXL or Flux. + +Image size +: A hint showing the megapixels. The exported images will not exceed this +: number. + +Bucket resolution size +: The bucket size the training tool is using. + +Preferres sizes +: A comma separated list of target sizes that should be preferred for the +: exported images. + +Allow upscaling +: Do upscale images when set. This is bad for the quality but might reduce the +: number of buckets that must be used for training. + +Bucket fitting strategy +: The method to make sure an image fits into a bucket. It can be a direct crop +: that removes information from the side of an image. Or a scaling that changes +: the aspect ratio of an image and can create slight distortions. Or a +: combination of both that reduces each effect. + +Output format +: The file type and quality setting for formats that have a lossy compression. +: Note: for JPEG a number above 95 should be avoided. + +Output color space +: Most models will expect the images in sRGB format and don't contain any +: color management. So it is important that the exporter handles this as +: the images used for the training might use a different color space. +: To save 8 kB for each image you might want to select "sRGB implicit" as that +: converts the image to sRGB but doesn't store the ICC information. +: When no color space convertation should happen you can choose "feed through". +: +: The simple "sRGB" is most likely the setting you want to choose here unless +: you are an expert and have special requirements. + +Export directory +: The place to export the images to. + +Keep input directory structure +: When the source images are organized in subdirectories this structure will +: be used for the exported images as well when selected. + +Statistics +: Preview of the generated image sizes from the export funtion. From c13689b990e862dd33ae721a8840069b616a37bb Mon Sep 17 00:00:00 2001 From: StableLlama Date: Wed, 5 Feb 2025 23:43:15 +0100 Subject: [PATCH 07/18] Fix markdown style --- README.md | 96 +++++++++++++++++++++++++++---------------------------- 1 file changed, 48 insertions(+), 48 deletions(-) diff --git a/README.md b/README.md index 961f6bdb..ec25985b 100644 --- a/README.md +++ b/README.md @@ -260,51 +260,51 @@ The `Edit` menu contains additional features for batch tag operations, such as Exporting the images to a directory allows different options. By choosing the preset for the target AI model many important settings are automatically set. -Resolution -: The native resolution of the model, like 1024 for SDXL or Flux. - -Image size -: A hint showing the megapixels. The exported images will not exceed this -: number. - -Bucket resolution size -: The bucket size the training tool is using. - -Preferres sizes -: A comma separated list of target sizes that should be preferred for the -: exported images. - -Allow upscaling -: Do upscale images when set. This is bad for the quality but might reduce the -: number of buckets that must be used for training. - -Bucket fitting strategy -: The method to make sure an image fits into a bucket. It can be a direct crop -: that removes information from the side of an image. Or a scaling that changes -: the aspect ratio of an image and can create slight distortions. Or a -: combination of both that reduces each effect. - -Output format -: The file type and quality setting for formats that have a lossy compression. -: Note: for JPEG a number above 95 should be avoided. - -Output color space -: Most models will expect the images in sRGB format and don't contain any -: color management. So it is important that the exporter handles this as -: the images used for the training might use a different color space. -: To save 8 kB for each image you might want to select "sRGB implicit" as that -: converts the image to sRGB but doesn't store the ICC information. -: When no color space convertation should happen you can choose "feed through". -: -: The simple "sRGB" is most likely the setting you want to choose here unless -: you are an expert and have special requirements. - -Export directory -: The place to export the images to. - -Keep input directory structure -: When the source images are organized in subdirectories this structure will -: be used for the exported images as well when selected. - -Statistics -: Preview of the generated image sizes from the export funtion. +`Resolution`: +The native resolution of the model, like 1024 for SDXL or Flux. + +`Image size`: +A hint showing the megapixels. The exported images will not exceed this +number. + +`Bucket resolution size`: +The bucket size the training tool is using. + +`Preferres sizes`: +A comma separated list of target sizes that should be preferred for the +exported images. + +`Allow upscaling`: +Do upscale images when set. This is bad for the quality but might reduce the +number of buckets that must be used for training. + +`Bucket fitting strategy`: +The method to make sure an image fits into a bucket. It can be a direct crop +that removes information from the side of an image. Or a scaling that changes +the aspect ratio of an image and can create slight distortions. Or a +combination of both that reduces each effect. + +`Output format`: +The file type and quality setting for formats that have a lossy compression. +Note: for JPEG a number above 95 should be avoided. + +`Output color space`: +Most models will expect the images in sRGB format and don't contain any +color management. So it is important that the exporter handles this as +the images used for the training might use a different color space. +To save 8 kB for each image you might want to select "sRGB implicit" as that +converts the image to sRGB but doesn't store the ICC information. +When no color space convertation should happen you can choose "feed through". + +The simple "sRGB" is most likely the setting you want to choose here unless +you are an expert and have special requirements. + +`Export directory`: +The place to export the images to. + +`Keep input directory structure`: +When the source images are organized in subdirectories this structure will +be used for the exported images as well when selected. + +`Statistics`: +Preview of the generated image sizes from the export funtion. From f397e98d11d7944d9bb29391bc34d520c5587797 Mon Sep 17 00:00:00 2001 From: StableLlama Date: Thu, 6 Feb 2025 00:18:54 +0100 Subject: [PATCH 08/18] Make sure to export the caption files as well --- taggui/dialogs/export_dialog.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/taggui/dialogs/export_dialog.py b/taggui/dialogs/export_dialog.py index 42e793f8..515c9f69 100644 --- a/taggui/dialogs/export_dialog.py +++ b/taggui/dialogs/export_dialog.py @@ -5,6 +5,7 @@ import re from math import floor, sqrt from pathlib import Path +import shutil from PySide6.QtCore import Qt, Slot from PySide6.QtGui import QColorSpace @@ -557,6 +558,11 @@ def do_export(self): export_path = export_path.parent / f"{stem}_{counter}{export_path.suffix}" counter += 1 + # copy the tag file first + if image_entry.path.with_suffix('.txt').exists(): + shutil.copyfile(str(image_entry.path.with_suffix('.txt')), str(export_path.with_suffix('.txt'))) + + # then handle the image image_file = Image.open(image_entry.path) # Preserve alpha if present: if image_file.mode in ("RGBA", "LA", "PA") and not export_format == ExportFormat.JPG: # Check for alpha channels From d3ac289c7079e905ba74aa372139bb535a2c631d Mon Sep 17 00:00:00 2001 From: StableLlama Date: Thu, 6 Feb 2025 23:51:54 +0100 Subject: [PATCH 09/18] Refactor DEFAULT_SETTINGS to ease initial value access Code refactor to respect mostly a width of 80 --- taggui/dialogs/export_dialog.py | 135 +++++++++++++++++------------- taggui/dialogs/settings_dialog.py | 12 +-- taggui/utils/settings_widgets.py | 33 +++++--- 3 files changed, 103 insertions(+), 77 deletions(-) diff --git a/taggui/dialogs/export_dialog.py b/taggui/dialogs/export_dialog.py index 515c9f69..b3997bb3 100644 --- a/taggui/dialogs/export_dialog.py +++ b/taggui/dialogs/export_dialog.py @@ -9,10 +9,10 @@ from PySide6.QtCore import Qt, Slot from PySide6.QtGui import QColorSpace -from PySide6.QtWidgets import (QWidget, QDialog, QFileDialog, QGridLayout, QLabel, - QLineEdit, QPushButton, QVBoxLayout, QHBoxLayout, - QTableWidget, QTableWidgetItem, QSizePolicy, - QMessageBox) +from PySide6.QtWidgets import (QWidget, QDialog, QFileDialog, QGridLayout, + QLabel, QLineEdit, QPushButton, QTableWidget, + QTableWidgetItem, QProgressBar, QMessageBox, + QVBoxLayout, QHBoxLayout, QSizePolicy) from PIL import Image, ImageFilter, ImageCms from utils.settings import DEFAULT_SETTINGS, get_settings @@ -44,7 +44,6 @@ class IccProfileList(str, Enum): AdobeRgb = 'AdobeRGB' DisplayP3 = 'DisplayP3' ProPhotoRgb = 'ProPhotoRGB' - # since PySide6.8: Bt2020 = 'BT.2020' Bt2100Pq = 'BT.2100(PQ)' Bt2100Hlg = 'BT.2100 (HLG)' @@ -65,9 +64,9 @@ def __init__(self, parent, image_list_model: ImageListModel): self.inhibit_statistics_update = True self.resolution_cache: dict[tuple, tuple] = {} self.setWindowTitle('Export') - layout = QVBoxLayout(self) - layout.setContentsMargins(20, 20, 20, 20) - layout.setSpacing(20) + self.layout = QVBoxLayout(self) + self.layout.setContentsMargins(20, 20, 20, 20) + self.layout.setSpacing(20) grid_layout = QGridLayout() grid_layout.setColumnStretch(0, 0) @@ -86,12 +85,13 @@ def __init__(self, parent, image_list_model: ImageListModel): grid_layout.addWidget(QLabel('Resolution (px)'), grid_row, 0, Qt.AlignmentFlag.AlignRight) self.resolution_spin_box = SettingsSpinBox( - key='export_resolution', default=DEFAULT_SETTINGS['export_resolution'], + key='export_resolution', minimum=0, maximum=8192) - self.resolution_spin_box.setToolTip('Common values:\n' - '0: disable rescaling\n' - '512: SD1.5\n' - '1024: SDXL, SD3, Flux') + self.resolution_spin_box.setToolTip( + 'Common values:\n' + '0: disable rescaling\n' + '512: SD1.5\n' + '1024: SDXL, SD3, Flux') self.resolution_spin_box.textChanged.connect(self.show_megapixels) self.resolution_spin_box.textChanged.connect(self.show_statistics) grid_layout.addWidget(self.resolution_spin_box, grid_row, 1, @@ -108,11 +108,12 @@ def __init__(self, parent, image_list_model: ImageListModel): grid_layout.addWidget(QLabel('Bucket resolution size (px)'), grid_row, 0, Qt.AlignmentFlag.AlignRight) self.bucket_res_size_spin_box = SettingsSpinBox( - key='export_bucket_res_size', default=DEFAULT_SETTINGS['export_bucket_res_size'], + key='export_bucket_res_size', minimum=1, maximum=256) - self.bucket_res_size_spin_box.setToolTip('Ensure that the exported image size is divisable by that number.\n' - 'It should match the setting on the training tool.\n' - 'It might cause minor cropping.') + self.bucket_res_size_spin_box.setToolTip( + 'Ensure that the exported image size is divisable by that number.\n' + 'It should match the setting on the training tool.\n' + 'It might cause minor cropping.') self.bucket_res_size_spin_box.textChanged.connect(self.show_statistics) grid_layout.addWidget(self.bucket_res_size_spin_box, grid_row, 1, Qt.AlignmentFlag.AlignLeft) @@ -121,22 +122,21 @@ def __init__(self, parent, image_list_model: ImageListModel): grid_layout.addWidget(QLabel('Preferred sizes'), grid_row, 0, Qt.AlignmentFlag.AlignRight) self.preferred_sizes_line_edit = SettingsLineEdit( - key='export_preferred_sizes', - default=DEFAULT_SETTINGS['export_preferred_sizes']) + key='export_preferred_sizes') self.preferred_sizes_line_edit.setMinimumWidth(500) - self.preferred_sizes_line_edit.setToolTip('Comma separated list of preferred sizes and aspect ratios.\n' - "The inverse aspect ratio is automatically derived and doesn't need to be included.") + self.preferred_sizes_line_edit.setToolTip( + 'Comma separated list of preferred sizes and aspect ratios.\n' + "The inverse aspect ratio is automatically derived and doesn't need to be included.") grid_layout.addWidget(self.preferred_sizes_line_edit, grid_row, 1, Qt.AlignmentFlag.AlignLeft) grid_row += 1 grid_layout.addWidget(QLabel('Allow upscaling'), grid_row, 0, Qt.AlignmentFlag.AlignRight) - self.upscaling_check_box = SettingsBigCheckBox( - key='export_upscaling', - default=DEFAULT_SETTINGS['export_upscaling']) - self.upscaling_check_box.setToolTip('Scale too small images to the requested size.\n' - 'This should be avoided as it lowers the image quality.') + self.upscaling_check_box = SettingsBigCheckBox(key='export_upscaling') + self.upscaling_check_box.setToolTip( + 'Scale too small images to the requested size.\n' + 'This should be avoided as it lowers the image quality.') self.upscaling_check_box.stateChanged.connect(self.show_statistics) grid_layout.addWidget(self.upscaling_check_box, grid_row, 1, Qt.AlignmentFlag.AlignLeft) @@ -144,11 +144,13 @@ def __init__(self, parent, image_list_model: ImageListModel): grid_row += 1 grid_layout.addWidget(QLabel('Bucket fitting strategy'), grid_row, 0, Qt.AlignmentFlag.AlignRight) - bucket_strategy_combo_box = SettingsComboBox(key='export_bucket_strategy') + bucket_strategy_combo_box = SettingsComboBox( + key='export_bucket_strategy') bucket_strategy_combo_box.addItems(list(BucketStrategy)) - bucket_strategy_combo_box.setToolTip('crop: center crop\n' - 'scale: assymetric scaling\n' - 'crop and scale: use both to minimize each effect') + bucket_strategy_combo_box.setToolTip( + 'crop: center crop\n' + 'scale: assymetric scaling\n' + 'crop and scale: use both to minimize each effect') grid_layout.addWidget(bucket_strategy_combo_box, grid_row, 1, Qt.AlignmentFlag.AlignLeft) @@ -158,8 +160,6 @@ def __init__(self, parent, image_list_model: ImageListModel): format_widget = QWidget() format_layout = QHBoxLayout() format_layout.setContentsMargins(0, 0, 0, 0) - current_format = self.settings.value('export_format', type=str) - current_quality = self.settings.value('export_quality', type=str) self.format_combo_box = SettingsComboBox(key='export_format') self.format_combo_box.addItems(list(ExportFormat)) self.format_combo_box.currentTextChanged.connect(self.format_change) @@ -168,11 +168,12 @@ def __init__(self, parent, image_list_model: ImageListModel): format_layout.addWidget(QLabel('Quality'), Qt.AlignmentFlag.AlignRight) self.quality_spin_box = SettingsSpinBox( - key='export_quality', default=DEFAULT_SETTINGS['export_quality'], + key='export_quality', minimum=0, maximum=100) - self.quality_spin_box.setToolTip('Only for JPEG and WebP.\n' - '0 is worst and 100 is best.\n' - 'For JPEG numbers above 95 should be avoided') + self.quality_spin_box.setToolTip( + 'Only for JPEG and WebP.\n' + '0 is worst and 100 is best.\n' + 'For JPEG numbers above 95 should be avoided') self.quality_spin_box.textChanged.connect(self.quality_change) format_layout.addWidget(self.quality_spin_box, Qt.AlignmentFlag.AlignLeft) @@ -180,6 +181,8 @@ def __init__(self, parent, image_list_model: ImageListModel): grid_layout.addWidget(format_widget, grid_row, 1, Qt.AlignmentFlag.AlignLeft) # ensure correct enable/disable and background color of the quality + current_format = self.settings.value('export_format', type=str) + current_quality = self.settings.value('export_quality', type=int) self.format_change(current_format, False) self.quality_change(current_quality) @@ -190,11 +193,12 @@ def __init__(self, parent, image_list_model: ImageListModel): color_space_combo_box.addItem("feed through (don't touch)") color_space_combo_box.addItem('sRGB (implicit, without profile)') color_space_combo_box.addItems([IccProfileList[e.name] for e in QColorSpace.NamedColorSpace]) - color_space_combo_box.setToolTip('Color space of the exported images.\n' - 'Most likely the trainer expects sRGB!\n' - '\n' - 'Use "feed through" to keep the color space as it is.\n' - 'Use "sRGB (implicit, without profile)" to save in sRGB but don\'t embed the ICC profile to save 8k file size.') + color_space_combo_box.setToolTip( + 'Color space of the exported images.\n' + 'Most likely the trainer expects sRGB!\n' + '\n' + 'Use "feed through" to keep the color space as it is.\n' + 'Use "sRGB (implicit, without profile)" to save in sRGB but don\'t embed the ICC profile to save 8k file size.') grid_layout.addWidget(color_space_combo_box, grid_row, 1, Qt.AlignmentFlag.AlignLeft) @@ -202,8 +206,7 @@ def __init__(self, parent, image_list_model: ImageListModel): grid_layout.addWidget(QLabel('Export directory'), grid_row, 0, Qt.AlignmentFlag.AlignRight) self.export_directory_line_edit = SettingsLineEdit( - key='export_directory_path', - default=DEFAULT_SETTINGS['export_directory_path']) + key='export_directory_path') self.export_directory_line_edit.setMinimumWidth(500) self.export_directory_line_edit.setClearButtonEnabled(True) grid_layout.addWidget(self.export_directory_line_edit, grid_row, 1, @@ -219,10 +222,10 @@ def __init__(self, parent, image_list_model: ImageListModel): grid_layout.addWidget(QLabel('Keep input directory structure'), grid_row, 0, Qt.AlignmentFlag.AlignRight) keep_dir_structure_check_box = SettingsBigCheckBox( - key='export_keep_dir_structure', - default=DEFAULT_SETTINGS['export_keep_dir_structure']) - keep_dir_structure_check_box.setToolTip('Keep the subdirectory structure or export\n' - 'all images in the same export directory') + key='export_keep_dir_structure') + keep_dir_structure_check_box.setToolTip( + 'Keep the subdirectory structure or export\n' + 'all images in the same export directory') grid_layout.addWidget(keep_dir_structure_check_box, grid_row, 1, Qt.AlignmentFlag.AlignLeft) @@ -230,15 +233,20 @@ def __init__(self, parent, image_list_model: ImageListModel): grid_layout.addWidget(QLabel('Statistics'), grid_row, 0, Qt.AlignmentFlag.AlignRight) self.statistics_table = QTableWidget(0, 5, self) - self.statistics_table.setHorizontalHeaderLabels(['Width', 'Height', 'Count', 'Aspect ratio', 'Size utilization']) + self.statistics_table.setHorizontalHeaderLabels( + ['Width', 'Height', 'Count', 'Aspect ratio', 'Size utilization']) self.statistics_table.setMinimumWidth(500) grid_layout.addWidget(self.statistics_table, grid_row, 1, Qt.AlignmentFlag.AlignLeft) - layout.addLayout(grid_layout) + self.layout.addLayout(grid_layout) + export_button = QPushButton('Export') - export_button.clicked.connect(self.do_export) - layout.addWidget(export_button) + if self.image_list_model.rowCount() > 0: + export_button.clicked.connect(self.do_export) + else: + export_button.setEnabled(False) + self.layout.addWidget(export_button) # update display self.apply_preset(preset_combo_box.currentText(), False) @@ -252,6 +260,7 @@ def apply_preset(self, value: str, do_value_change: bool = True): Slot to call when a new preset was selected to help the user to set important settings to a consistent state. """ + inhibit_statistics_update_current = self.inhibit_statistics_update preset = Presets[value] if value == 'manual': self.resolution_spin_box.setEnabled(True) @@ -262,7 +271,7 @@ def apply_preset(self, value: str, do_value_change: bool = True): self.resolution_spin_box.setEnabled(False) self.bucket_res_size_spin_box.setValue(preset[1]) if do_value_change else 0 self.bucket_res_size_spin_box.setEnabled(False) - self.inhibit_statistics_update = False + self.inhibit_statistics_update = inhibit_statistics_update_current self.show_statistics() self.preferred_sizes_line_edit.setText(preset[2]) if do_value_change else 0 @@ -327,15 +336,18 @@ def show_statistics(self): (21, 9, 21/9), ] self.preferred_sizes = [] - for res_str in re.split(r'\s*,\s*', self.settings.value('export_preferred_sizes')): + for res_str in re.split(r'\s*,\s*', self.settings.value('export_preferred_sizes') or ''): try: + if res_str == '': + continue size_str = res_str.split(':') width = max(int(size_str[0]), int(size_str[1])) height = min(int(size_str[0]), int(size_str[1])) self.preferred_sizes.append((width, height)) if not width == height: self.preferred_sizes.append((height, width)) - # add exact aspect ratio of the preferred size to label it similar to the perfect one + # add exact aspect ratio of the preferred size to label it + # similar to the perfect one aspect_ratio = width / height for ar in aspect_ratios: if abs(ar[2] - aspect_ratio) < 0.15: @@ -343,13 +355,14 @@ def show_statistics(self): break except ValueError: # Handle cases where the resolution string is not in the correct format - print(f"Warning: Invalid resolution format: {res_str}. Skipping.") + print(f'Warning: Invalid resolution format: {res_str}. Skipping.') continue # Skip to the next resolution if there's an error image_dimensions = defaultdict(int) for image_index in range(self.image_list_model.rowCount()): this_image = self.image_list_model.index(image_index).data(Qt.ItemDataRole.UserRole) - this_image.target_dimensions = self.target_dimensions(this_image.dimensions, resolution, upscaling, bucket_res) + this_image.target_dimensions = self.target_dimensions( + this_image.dimensions, resolution, upscaling, bucket_res) image_dimensions[this_image.target_dimensions] += 1 sorted_dimensions = sorted( @@ -499,6 +512,13 @@ def do_export(self): export_directory_path = Path(self.settings.value('export_directory_path', type=str)) export_keep_dir_structure = self.settings.value('export_keep_dir_structure', type=bool) no_overwrite = True + + image_count = self.image_list_model.rowCount() + self.progress_bar = QProgressBar(self) + self.progress_bar.setMinimum(0) + self.progress_bar.setMaximum(image_count) + self.layout.addWidget(self.progress_bar) + if os.path.exists(export_directory_path): if os.path.isfile(export_directory_path): QMessageBox.critical( @@ -541,7 +561,8 @@ def do_export(self): save_profile = False bucket_strategy = self.settings.value('export_bucket_strategy', type=str) - for image_index in range(self.image_list_model.rowCount()): + for image_index in range(image_count): + self.progress_bar.setValue(image_index) image_entry = self.image_list_model.index(image_index).data(Qt.ItemDataRole.UserRole) if export_keep_dir_structure: relative_path = image_entry.path.relative_to(directory_path) diff --git a/taggui/dialogs/settings_dialog.py b/taggui/dialogs/settings_dialog.py index 8fb2956b..49d27f70 100644 --- a/taggui/dialogs/settings_dialog.py +++ b/taggui/dialogs/settings_dialog.py @@ -33,19 +33,17 @@ def __init__(self, parent): Qt.AlignmentFlag.AlignRight) font_size_spin_box = SettingsSpinBox( - key='font_size', default=DEFAULT_SETTINGS['font_size'], + key='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( 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) self.insert_space_after_tag_separator_check_box = SettingsBigCheckBox( - key='insert_space_after_tag_separator', - default=DEFAULT_SETTINGS['insert_space_after_tag_separator']) + key='insert_space_after_tag_separator') self.insert_space_after_tag_separator_check_box.stateChanged.connect( self.show_restart_warning) tag_separator_line_edit = QLineEdit() @@ -60,8 +58,7 @@ def __init__(self, parent): tag_separator_line_edit.textChanged.connect( self.handle_tag_separator_change) autocomplete_tags_check_box = SettingsBigCheckBox( - key='autocomplete_tags', - default=DEFAULT_SETTINGS['autocomplete_tags']) + key='autocomplete_tags') autocomplete_tags_check_box.stateChanged.connect( self.show_restart_warning) self.models_directory_line_edit = SettingsLineEdit( @@ -76,8 +73,7 @@ def __init__(self, parent): 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']) + key='image_list_file_formats') file_types_line_edit.setMinimumWidth(400) file_types_line_edit.textChanged.connect(self.show_restart_warning) diff --git a/taggui/utils/settings_widgets.py b/taggui/utils/settings_widgets.py index 43b0b5ec..3369b618 100644 --- a/taggui/utils/settings_widgets.py +++ b/taggui/utils/settings_widgets.py @@ -4,14 +4,16 @@ from utils.big_widgets import BigCheckBox from utils.focused_scroll_mixin import FocusedScrollMixin -from utils.settings import get_settings +from utils.settings import DEFAULT_SETTINGS, get_settings class SettingsBigCheckBox(BigCheckBox): - def __init__(self, key: str, default: bool, text: str | None = None): + def __init__(self, key: str, default: bool | None = None, text: str | None = None): super().__init__(text) settings = get_settings() - self.setChecked(settings.value(key, default, type=bool)) + if not settings.contains(key): + settings.setValue(key, default or DEFAULT_SETTINGS.get(key)) + self.setChecked(settings.value(key, type=bool)) self.stateChanged.connect( lambda state: settings.setValue( key, state == Qt.CheckState.Checked.value)) @@ -20,12 +22,13 @@ def __init__(self, key: str, default: bool, text: str | None = None): class SettingsComboBox(QComboBox): def __init__(self, key: str, default: str | None = None): super().__init__() - self.key = key - self.default = default self.settings = get_settings() + self.key = key + if not self.settings.contains(key): + self.settings.setValue(key, default or DEFAULT_SETTINGS.get(key)) def addItems(self, texts: list[str]): - setting: str = self.settings.value(self.key, self.default, type=str) + setting: str = self.settings.value(self.key, type=str) super().addItems(texts) self.currentTextChanged.connect( lambda text: self.settings.setValue(self.key, text)) @@ -50,11 +53,13 @@ def __init__(self, key: str, default: float, minimum: float, class SettingsSpinBox(QSpinBox): - def __init__(self, key: str, default: int, minimum: int, maximum: int): + def __init__(self, key: str, minimum: int, maximum: int, default: int | None = None): super().__init__() self.setRange(minimum, maximum) settings = get_settings() - self.setValue(settings.value(key, default, type=int)) + if not settings.contains(key): + settings.setValue(key, default or DEFAULT_SETTINGS.get(key)) + self.setValue(settings.value(key, type=int)) self.valueChanged.connect(lambda value: settings.setValue(key, value)) @@ -63,17 +68,21 @@ class FocusedScrollSettingsSpinBox(FocusedScrollMixin, SettingsSpinBox): class SettingsLineEdit(QLineEdit): - def __init__(self, key: str, default: str = ''): + def __init__(self, key: str, default: str | None = None): super().__init__() settings = get_settings() - self.setText(settings.value(key, default, type=str)) + if not settings.contains(key): + settings.setValue(key, default or DEFAULT_SETTINGS.get(key, '')) + self.setText(settings.value(key, type=str)) self.textChanged.connect(lambda text: settings.setValue(key, text)) class SettingsPlainTextEdit(QPlainTextEdit): - def __init__(self, key: str, default: str = ''): + def __init__(self, key: str, default: str | None = None): super().__init__() settings = get_settings() - self.setPlainText(settings.value(key, default, type=str)) + if not settings.contains(key): + settings.setValue(key, default or DEFAULT_SETTINGS.get(key, '')) + self.setPlainText(settings.value(key, type=str)) self.textChanged.connect(lambda: settings.setValue(key, self.toPlainText())) From 15e384200a43181ce007de1b42fedd9f37a46175 Mon Sep 17 00:00:00 2001 From: StableLlama Date: Fri, 7 Feb 2025 00:08:55 +0100 Subject: [PATCH 10/18] Use new default infrastructure --- taggui/dialogs/export_dialog.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/taggui/dialogs/export_dialog.py b/taggui/dialogs/export_dialog.py index b3997bb3..13c045d8 100644 --- a/taggui/dialogs/export_dialog.py +++ b/taggui/dialogs/export_dialog.py @@ -322,9 +322,9 @@ def show_statistics(self): return self.resolution_cache = {} - resolution = self.resolution_spin_box.value() - upscaling = self.upscaling_check_box.isChecked() - bucket_res = self.bucket_res_size_spin_box.value() + resolution = self.settings.value('export_resolution', type=int) + upscaling = self.settings.value('export_upscaling', type=int) + bucket_res = self.settings.value('export_bucket_res_size', type=int) # notable aspect ratios aspect_ratios = [ @@ -549,11 +549,11 @@ def do_export(self): ) return - resolution = self.resolution_spin_box.value() - upscaling = self.upscaling_check_box.isChecked() - bucket_res = self.bucket_res_size_spin_box.value() - export_format = self.format_combo_box.currentText() - quality = self.quality_spin_box.value() + resolution = self.settings.value('export_resolution', type=int) + upscaling = self.settings.value('export_upscaling', type=int) + bucket_res = self.settings.value('export_bucket_res_size', type=int) + export_format = self.settings.value('export_format', type=str) + quality = self.settings.value('export_quality', type=int) color_space = self.settings.value('export_color_space', type=str) save_profile = True if color_space == 'sRGB (implicit, without profile)': From 5beaa26cf3b47dce516241ed8a9af9488f7377ef Mon Sep 17 00:00:00 2001 From: StableLlama Date: Fri, 7 Feb 2025 00:40:27 +0100 Subject: [PATCH 11/18] Fix broken tagging --- taggui/dialogs/export_dialog.py | 3 ++- taggui/dialogs/settings_dialog.py | 3 +-- taggui/utils/image.py | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/taggui/dialogs/export_dialog.py b/taggui/dialogs/export_dialog.py index 13c045d8..7c5d3ed3 100644 --- a/taggui/dialogs/export_dialog.py +++ b/taggui/dialogs/export_dialog.py @@ -355,7 +355,8 @@ def show_statistics(self): break except ValueError: # Handle cases where the resolution string is not in the correct format - print(f'Warning: Invalid resolution format: {res_str}. Skipping.') + print(f'Warning: Invalid resolution format: {res_str}. Skipping.', + file=sys.stderr) continue # Skip to the next resolution if there's an error image_dimensions = defaultdict(int) diff --git a/taggui/dialogs/settings_dialog.py b/taggui/dialogs/settings_dialog.py index 49d27f70..ee8d6140 100644 --- a/taggui/dialogs/settings_dialog.py +++ b/taggui/dialogs/settings_dialog.py @@ -62,8 +62,7 @@ def __init__(self, parent): autocomplete_tags_check_box.stateChanged.connect( self.show_restart_warning) self.models_directory_line_edit = SettingsLineEdit( - key='models_directory_path', - default=DEFAULT_SETTINGS['models_directory_path']) + key='models_directory_path') self.models_directory_line_edit.setMinimumWidth(400) self.models_directory_line_edit.setClearButtonEnabled(True) self.models_directory_line_edit.textChanged.connect( diff --git a/taggui/utils/image.py b/taggui/utils/image.py index 96554486..17091470 100644 --- a/taggui/utils/image.py +++ b/taggui/utils/image.py @@ -8,6 +8,6 @@ class Image: path: Path dimensions: tuple[int, int] | None - target_dimensions: tuple[int, int] | None tags: list[str] = field(default_factory=list) + target_dimensions: tuple[int, int] | None = None thumbnail: QIcon | None = None From 1de5c5705c5aec0f7dbbc34876508c3ec298f576 Mon Sep 17 00:00:00 2001 From: StableLlama Date: Fri, 7 Feb 2025 10:39:40 +0100 Subject: [PATCH 12/18] Finetune sharpening --- taggui/dialogs/export_dialog.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/taggui/dialogs/export_dialog.py b/taggui/dialogs/export_dialog.py index 7c5d3ed3..ba9298d7 100644 --- a/taggui/dialogs/export_dialog.py +++ b/taggui/dialogs/export_dialog.py @@ -606,7 +606,7 @@ def do_export(self): # resize with the best method available resized_image = image_file.resize((new_width, new_height), Image.LANCZOS) # followed by a slight sharpening as it should be done - sharpend_image = resized_image.filter(ImageFilter.UnsharpMask(radius = 0.5, percent = 100, threshold = 3)) + sharpend_image = resized_image.filter(ImageFilter.UnsharpMask(radius = 0.5, percent = 50, threshold = 0)) else: sharpend_image = image_file From deb790e5f29ae0261922ac9c01142a7041a1e20b Mon Sep 17 00:00:00 2001 From: StableLlama Date: Wed, 12 Feb 2025 22:30:39 +0100 Subject: [PATCH 13/18] Add export with the JPEG XL format Advantages are smaller file sizes when compression is acceptable (quality < 100) or lossless compression with quality = 100. Also the alpha channel is supported and kept in the images. Note 1: this only adds support for the export function, and is not general JPEG XL support for taggui. There is the fork https://github.com/yggdrasil75/taggui that does exactly this. Note 2: You might need `pip install pillow-jxl-plugin` beforehand to be able to export into JPEG XL. --- taggui/dialogs/export_dialog.py | 24 ++++++++++++++++++++---- 1 file changed, 20 insertions(+), 4 deletions(-) diff --git a/taggui/dialogs/export_dialog.py b/taggui/dialogs/export_dialog.py index ba9298d7..4b6c0a91 100644 --- a/taggui/dialogs/export_dialog.py +++ b/taggui/dialogs/export_dialog.py @@ -20,6 +20,11 @@ SettingsSpinBox, SettingsComboBox) from models.image_list_model import ImageListModel +try: + import pillow_jxl +except ModuleNotFoundError: + pass + Presets = { 'manual': (0, 0, '1:1, 2:1, 3:2, 4:3, 16:9, 21:9'), 'Direct feed through': (0, 1, '1:1, 2:1, 3:2, 4:3, 16:9, 21:9'), @@ -29,11 +34,13 @@ class ExportFormat(str, Enum): JPG = '.jpg - JPEG' + JPGXL = '.jxl - JPEG XL' PNG = '.png - PNG' WEBP = '.webp - WEBP' ExportFormatDict = { ExportFormat.JPG: 'jpeg', + ExportFormat.JPGXL: 'jxl', ExportFormat.PNG: 'png', ExportFormat.WEBP: 'webp' } @@ -161,7 +168,12 @@ def __init__(self, parent, image_list_model: ImageListModel): format_layout = QHBoxLayout() format_layout.setContentsMargins(0, 0, 0, 0) self.format_combo_box = SettingsComboBox(key='export_format') - self.format_combo_box.addItems(list(ExportFormat)) + supported_extensions = set(Image.registered_extensions().keys()) + supported_formats = [ + format for format in ExportFormat + if any(ext in supported_extensions for ext in format.value.split(' - ')[0].split(',')) + ] + self.format_combo_box.addItems(supported_formats) self.format_combo_box.currentTextChanged.connect(self.format_change) format_layout.addWidget(self.format_combo_box, Qt.AlignmentFlag.AlignLeft) @@ -296,6 +308,9 @@ def format_change(self, export_format: ExportFormat, do_value_change: bool = Tru if export_format == ExportFormat.JPG: self.quality_spin_box.setValue(75) if do_value_change else 0 self.quality_spin_box.setEnabled(True) + if export_format == ExportFormat.JPGXL: + self.quality_spin_box.setValue(100) if do_value_change else 0 + self.quality_spin_box.setEnabled(True) elif export_format == ExportFormat.PNG: self.quality_spin_box.setValue(100) if do_value_change else 0 self.quality_spin_box.setEnabled(False) @@ -615,9 +630,10 @@ def do_export(self): crop_width = floor((current_width - image_entry.target_dimensions[0]) / 2) crop_height = floor((current_height - image_entry.target_dimensions[1]) / 2) cropped_image = sharpend_image.crop((crop_width, crop_height, current_width - crop_width, current_height - crop_height)) + lossless = quality > 99 if color_space == "feed through (don't touch)": - cropped_image.save(export_path, format=ExportFormatDict[export_format], quality=quality, icc_profile=cropped_image.info.get('icc_profile')) + cropped_image.save(export_path, format=ExportFormatDict[export_format], quality=quality, icc_profile=cropped_image.info.get('icc_profile'), lossless=lossless) else: source_profile_raw = image_file.info.get('icc_profile') if source_profile_raw is None: # assume sRGB @@ -627,7 +643,7 @@ def do_export(self): target_profile = ImageCms.ImageCmsProfile(io.BytesIO(target_profile_raw)) final_image = ImageCms.profileToProfile(cropped_image, source_profile, target_profile) if save_profile: - final_image.save(export_path, format=ExportFormatDict[export_format], quality=quality, icc_profile=target_profile.tobytes()) + final_image.save(export_path, format=ExportFormatDict[export_format], quality=quality, icc_profile=target_profile.tobytes(), lossless=lossless) else: - final_image.save(export_path, format=ExportFormatDict[export_format], quality=quality, icc_profile=None) + final_image.save(export_path, format=ExportFormatDict[export_format], quality=quality, icc_profile=None, lossless=lossless) self.close() From f9321a6b5b2ecd9d03756c7f8337c5d56551cf74 Mon Sep 17 00:00:00 2001 From: StableLlama Date: Fri, 14 Feb 2025 00:04:30 +0100 Subject: [PATCH 14/18] Allow export to respect filter and selection of images --- taggui/dialogs/export_dialog.py | 61 +++++++++++++++++++++++++-------- taggui/utils/settings.py | 1 + taggui/widgets/main_window.py | 2 +- 3 files changed, 48 insertions(+), 16 deletions(-) diff --git a/taggui/dialogs/export_dialog.py b/taggui/dialogs/export_dialog.py index 4b6c0a91..c06a0056 100644 --- a/taggui/dialogs/export_dialog.py +++ b/taggui/dialogs/export_dialog.py @@ -18,13 +18,18 @@ from utils.settings import DEFAULT_SETTINGS, get_settings from utils.settings_widgets import (SettingsBigCheckBox, SettingsLineEdit, SettingsSpinBox, SettingsComboBox) -from models.image_list_model import ImageListModel +from widgets.image_list import ImageList try: import pillow_jxl except ModuleNotFoundError: pass +class ExportFilter(str, Enum): + NONE = 'All images' + FILTERED = 'Filtered images' + SELECTED = 'Selected images' + Presets = { 'manual': (0, 0, '1:1, 2:1, 3:2, 4:3, 16:9, 21:9'), 'Direct feed through': (0, 1, '1:1, 2:1, 3:2, 4:3, 16:9, 21:9'), @@ -61,12 +66,12 @@ class BucketStrategy(str, Enum): CROP_SCALE = 'crop and scale' class ExportDialog(QDialog): - def __init__(self, parent, image_list_model: ImageListModel): + def __init__(self, parent, image_list: ImageList): """ Main method to create the export dialog. """ super().__init__(parent) - self.image_list_model = image_list_model + self.image_list_view = image_list.list_view self.settings = get_settings() self.inhibit_statistics_update = True self.resolution_cache: dict[tuple, tuple] = {} @@ -80,6 +85,15 @@ def __init__(self, parent, image_list_model: ImageListModel): grid_layout.setColumnStretch(1, 1) grid_row = 0 + grid_layout.addWidget(QLabel('Image selection'), grid_row, 0, + Qt.AlignmentFlag.AlignRight) + preset_combo_box = SettingsComboBox(key='export_filter') + preset_combo_box.addItems(list(ExportFilter)) + preset_combo_box.currentTextChanged.connect(self.show_statistics) + grid_layout.addWidget(preset_combo_box, grid_row, 1, + Qt.AlignmentFlag.AlignLeft) + + grid_row += 1 grid_layout.addWidget(QLabel('Preset'), grid_row, 0, Qt.AlignmentFlag.AlignRight) preset_combo_box = SettingsComboBox(key='export_preset') @@ -253,12 +267,10 @@ def __init__(self, parent, image_list_model: ImageListModel): self.layout.addLayout(grid_layout) - export_button = QPushButton('Export') - if self.image_list_model.rowCount() > 0: - export_button.clicked.connect(self.do_export) - else: - export_button.setEnabled(False) - self.layout.addWidget(export_button) + self.export_button = QPushButton('Export') + self.export_button.clicked.connect(self.do_export) + self.export_button.setEnabled(False) + self.layout.addWidget(self.export_button) # update display self.apply_preset(preset_combo_box.currentText(), False) @@ -374,9 +386,9 @@ def show_statistics(self): file=sys.stderr) continue # Skip to the next resolution if there's an error + image_list = self.get_image_list() image_dimensions = defaultdict(int) - for image_index in range(self.image_list_model.rowCount()): - this_image = self.image_list_model.index(image_index).data(Qt.ItemDataRole.UserRole) + for this_image in image_list: this_image.target_dimensions = self.target_dimensions( this_image.dimensions, resolution, upscaling, bucket_res) image_dimensions[this_image.target_dimensions] += 1 @@ -385,6 +397,8 @@ def show_statistics(self): image_dimensions.items(), key=lambda x: x[0][0] / x[0][1] # Sort by width/height ratio ) + self.export_button.setEnabled(len(image_list) > 0) + self.statistics_table.setRowCount(0) # clear old data for dimensions, count in sorted_dimensions: @@ -529,10 +543,10 @@ def do_export(self): export_keep_dir_structure = self.settings.value('export_keep_dir_structure', type=bool) no_overwrite = True - image_count = self.image_list_model.rowCount() + image_list = self.get_image_list() self.progress_bar = QProgressBar(self) self.progress_bar.setMinimum(0) - self.progress_bar.setMaximum(image_count) + self.progress_bar.setMaximum(len(image_list)) self.layout.addWidget(self.progress_bar) if os.path.exists(export_directory_path): @@ -577,9 +591,8 @@ def do_export(self): save_profile = False bucket_strategy = self.settings.value('export_bucket_strategy', type=str) - for image_index in range(image_count): + for image_index, image_entry in enumerate(self.get_image_list()): self.progress_bar.setValue(image_index) - image_entry = self.image_list_model.index(image_index).data(Qt.ItemDataRole.UserRole) if export_keep_dir_structure: relative_path = image_entry.path.relative_to(directory_path) export_path = export_directory_path / relative_path @@ -647,3 +660,21 @@ def do_export(self): else: final_image.save(export_path, format=ExportFormatDict[export_format], quality=quality, icc_profile=None, lossless=lossless) self.close() + + def get_image_list(self): + if self.settings.value('export_filter') == ExportFilter.FILTERED: + images = self.image_list_view.proxy_image_list_model.sourceModel() + image_list = [] + for row in range(self.image_list_view.proxy_image_list_model.sourceModel().rowCount()): + source_index = self.image_list_view.proxy_image_list_model.sourceModel().index(row, 0) + proxy_index = self.image_list_view.proxy_image_list_model.mapFromSource(source_index) + if proxy_index.isValid(): + image_list.append(source_index.data(Qt.ItemDataRole.UserRole)) + elif self.settings.value('export_filter') == ExportFilter.SELECTED: + images = self.image_list_view.proxy_image_list_model.sourceModel() + image_list = [image_index.data(Qt.ItemDataRole.UserRole) for image_index in self.image_list_view.get_selected_image_indices()] + else: # ExportFilter.NONE + images = self.image_list_view.proxy_image_list_model.sourceModel() + image_list = [images.index(image_index).data(Qt.ItemDataRole.UserRole) for image_index in range(images.rowCount())] + + return image_list diff --git a/taggui/utils/settings.py b/taggui/utils/settings.py index d256d2c4..8d4ba547 100644 --- a/taggui/utils/settings.py +++ b/taggui/utils/settings.py @@ -10,6 +10,7 @@ 'insert_space_after_tag_separator': True, 'autocomplete_tags': True, 'models_directory_path': '', + 'export_filter': 'All images', 'export_preset': 'SDXL, SD3, Flux', 'export_resolution': 1024, 'export_bucket_res_size': 64, diff --git a/taggui/widgets/main_window.py b/taggui/widgets/main_window.py index 87533a8f..35e360ec 100644 --- a/taggui/widgets/main_window.py +++ b/taggui/widgets/main_window.py @@ -258,7 +258,7 @@ def reload_directory(self): @Slot() def export_images_dialog(self): - export_dialog = ExportDialog(parent=self, image_list_model=self.image_list_model) + export_dialog = ExportDialog(parent=self, image_list=self.image_list) export_dialog.exec() return From 4fd486814d240c46516a464f858b42ea3c65485d Mon Sep 17 00:00:00 2001 From: StableLlama Date: Fri, 14 Feb 2025 17:24:26 +0100 Subject: [PATCH 15/18] Add possibility to filter for image size --- README.md | 13 +++++++++++++ taggui/models/proxy_image_list_model.py | 10 ++++++++++ taggui/widgets/image_list.py | 4 ++-- 3 files changed, 25 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index ec25985b..19ae0fc1 100644 --- a/README.md +++ b/README.md @@ -143,6 +143,9 @@ apply: - `path`: Images that contain the filter term in the full file path - `path:cat` will match images such as `C:\Users\cats\dog.jpg` or `/home/dogs/cat.jpg`. +- `size`: Images that have the given size, stated in double colon separated + numbers. + - `size:512:512` will match images of the dimension 512x512 pixels. - You can also use a filter term with no prefix to filter for images that contain the term in either the caption or the file path. - `cat` will match images containing `cat` in the caption or file path. @@ -165,6 +168,9 @@ comparison. caption. - `tokens:<=50` will match images that have 50 or fewer tokens in the caption. +- `x` and `y`: will match images with the specified x or y dimension. + - `x:>512` will match images where the width is greater than 512 pixels. + - `y:=1024` will match images where the height is exactly 1024 pixels. ### Spaces and quotes @@ -260,6 +266,13 @@ The `Edit` menu contains additional features for batch tag operations, such as Exporting the images to a directory allows different options. By choosing the preset for the target AI model many important settings are automatically set. +`Image selection`: +Select whether all images, or those with the current filter or only the +currently selected images should be exported. + +`Preset`: +Choose a given preset or `manual` to set your own values. + `Resolution`: The native resolution of the model, like 1024 for SDXL or Flux. diff --git a/taggui/models/proxy_image_list_model.py b/taggui/models/proxy_image_list_model.py index b730c059..734ed6c6 100644 --- a/taggui/models/proxy_image_list_model.py +++ b/taggui/models/proxy_image_list_model.py @@ -37,6 +37,12 @@ def does_image_match_filter(self, image: Image, return fnmatchcase(image.path.name, f'*{filter_[1]}*') if filter_[0] == 'path': return fnmatchcase(str(image.path), f'*{filter_[1]}*') + if filter_[0] == 'size': + # accept any dimension separator of [x:] + dimension = (filter_[1]).replace(':', 'x').split('x') + return (len(dimension) == 2 + and dimension[0] == str(image.dimensions[0]) + and dimension[1] == str(image.dimensions[1])) if filter_[1] == 'AND': return (self.does_image_match_filter(image, filter_[0]) and self.does_image_match_filter(image, filter_[2:])) @@ -63,6 +69,10 @@ def does_image_match_filter(self, image: Image, caption = self.tag_separator.join(image.tags) # Subtract 2 for the `<|startoftext|>` and `<|endoftext|>` tokens. number_to_compare = len(self.tokenizer(caption).input_ids) - 2 + elif filter_[0] == 'x': + number_to_compare = image.dimensions[0] + elif filter_[0] == 'y': + number_to_compare = image.dimensions[1] return comparison_operator(number_to_compare, int(filter_[2])) def filterAcceptsRow(self, source_row: int, diff --git a/taggui/widgets/image_list.py b/taggui/widgets/image_list.py index 18af0eb0..c8dab6ad 100644 --- a/taggui/widgets/image_list.py +++ b/taggui/widgets/image_list.py @@ -48,12 +48,12 @@ def __init__(self): | QuotedString(quote_char="'", esc_char='\\') | Word(printables, exclude_chars='()')) - string_filter_keys = ['tag', 'caption', 'name', 'path'] + string_filter_keys = ['tag', 'caption', 'name', 'path', 'size'] string_filter_expressions = [Group(CaselessLiteral(key) + Suppress(':') + optionally_quoted_string) for key in string_filter_keys] comparison_operator = one_of('= == != < > <= >=') - number_filter_keys = ['tags', 'chars', 'tokens'] + number_filter_keys = ['tags', 'chars', 'tokens', 'x', 'y'] number_filter_expressions = [Group(CaselessLiteral(key) + Suppress(':') + comparison_operator + Word(nums)) for key in number_filter_keys] From a846f7a3849c94fa6751ae7ffa0451bb5f5686d4 Mon Sep 17 00:00:00 2001 From: StableLlama Date: Sat, 15 Feb 2025 01:10:36 +0100 Subject: [PATCH 16/18] Add filtering for target size. Also refactor the code a bit to be able to use the target size calculation more easily in (future) other modules as well. --- taggui/dialogs/export_dialog.py | 161 +++++------------------- taggui/models/proxy_image_list_model.py | 10 +- taggui/utils/target_dimension.py | 143 +++++++++++++++++++++ taggui/widgets/image_list.py | 2 +- 4 files changed, 184 insertions(+), 132 deletions(-) create mode 100644 taggui/utils/target_dimension.py diff --git a/taggui/dialogs/export_dialog.py b/taggui/dialogs/export_dialog.py index c06a0056..5420a2b5 100644 --- a/taggui/dialogs/export_dialog.py +++ b/taggui/dialogs/export_dialog.py @@ -2,8 +2,7 @@ from collections import defaultdict import os import io -import re -from math import floor, sqrt +from math import floor from pathlib import Path import shutil @@ -12,12 +11,14 @@ from PySide6.QtWidgets import (QWidget, QDialog, QFileDialog, QGridLayout, QLabel, QLineEdit, QPushButton, QTableWidget, QTableWidgetItem, QProgressBar, QMessageBox, - QVBoxLayout, QHBoxLayout, QSizePolicy) + QVBoxLayout, QHBoxLayout, QSizePolicy, + QAbstractItemView) from PIL import Image, ImageFilter, ImageCms from utils.settings import DEFAULT_SETTINGS, get_settings from utils.settings_widgets import (SettingsBigCheckBox, SettingsLineEdit, SettingsSpinBox, SettingsComboBox) +import utils.target_dimension as target_dimension from widgets.image_list import ImageList try: @@ -71,10 +72,9 @@ def __init__(self, parent, image_list: ImageList): Main method to create the export dialog. """ super().__init__(parent) - self.image_list_view = image_list.list_view + self.image_list = image_list self.settings = get_settings() self.inhibit_statistics_update = True - self.resolution_cache: dict[tuple, tuple] = {} self.setWindowTitle('Export') self.layout = QVBoxLayout(self) self.layout.setContentsMargins(20, 20, 20, 20) @@ -262,6 +262,8 @@ def __init__(self, parent, image_list: ImageList): self.statistics_table.setHorizontalHeaderLabels( ['Width', 'Height', 'Count', 'Aspect ratio', 'Size utilization']) self.statistics_table.setMinimumWidth(500) + self.statistics_table.setEditTriggers(QAbstractItemView.NoEditTriggers) + self.statistics_table.itemDoubleClicked.connect(self.set_filter) grid_layout.addWidget(self.statistics_table, grid_row, 1, Qt.AlignmentFlag.AlignLeft) @@ -348,7 +350,6 @@ def show_statistics(self): if self.inhibit_statistics_update: return - self.resolution_cache = {} resolution = self.settings.value('export_resolution', type=int) upscaling = self.settings.value('export_upscaling', type=int) bucket_res = self.settings.value('export_bucket_res_size', type=int) @@ -362,36 +363,15 @@ def show_statistics(self): (16, 9, 16/9), (21, 9, 21/9), ] - self.preferred_sizes = [] - for res_str in re.split(r'\s*,\s*', self.settings.value('export_preferred_sizes') or ''): - try: - if res_str == '': - continue - size_str = res_str.split(':') - width = max(int(size_str[0]), int(size_str[1])) - height = min(int(size_str[0]), int(size_str[1])) - self.preferred_sizes.append((width, height)) - if not width == height: - self.preferred_sizes.append((height, width)) - # add exact aspect ratio of the preferred size to label it - # similar to the perfect one - aspect_ratio = width / height - for ar in aspect_ratios: - if abs(ar[2] - aspect_ratio) < 0.15: - aspect_ratios.append((ar[0], ar[1], aspect_ratio)) - break - except ValueError: - # Handle cases where the resolution string is not in the correct format - print(f'Warning: Invalid resolution format: {res_str}. Skipping.', - file=sys.stderr) - continue # Skip to the next resolution if there's an error + aspect_ratios = target_dimension.prepare(aspect_ratios) image_list = self.get_image_list() image_dimensions = defaultdict(int) for this_image in image_list: - this_image.target_dimensions = self.target_dimensions( - this_image.dimensions, resolution, upscaling, bucket_res) + this_image.target_dimensions = target_dimension.get( + this_image.dimensions) image_dimensions[this_image.target_dimensions] += 1 + self.image_list.proxy_image_list_model.invalidate() sorted_dimensions = sorted( image_dimensions.items(), @@ -420,98 +400,18 @@ def show_statistics(self): self.statistics_table.setItem(rowPosition, 3, QTableWidgetItem(f"{aspect_ratio:.3f}{notable_aspect_ratio}")) self.statistics_table.setItem(rowPosition, 4, QTableWidgetItem(f"{100*utilization:.1f}%")) - def target_dimensions(self, dimensions: tuple[int, int], resolution: int, upscaling: bool, bucket_res: int): - """ - Determine the dimensions of an image it should have when it is exported. - - Note: this gives the optimal answer and thus can be slower than the Kohya - bucket algorithm. - - Parameters - ---------- - dimensions : tuple[int, int] - The width and height of the image - resolution : int - The target resolution of the AI model. The target image pixels - will not exceed the square of this number - upscaling : bool - Is upscaling of images allowed? - bucket_res : int - The resolution of the buckets - """ - width, height = dimensions - if resolution == 0: - # no rescale in this case, only cropping - return ((width // bucket_res)*bucket_res, (height // bucket_res)*bucket_res) - - if width < bucket_res or height < bucket_res: - # it doesn't make sense to use such a small image. But we shouldn't - # patronize the user - return dimensions - - if dimensions in self.resolution_cache: - return self.resolution_cache[dimensions] - - preferred_sizes_bonus = 0.4 # reduce the loss by this factor - - max_pixels = resolution * resolution - opt_width = resolution * sqrt(width/height) - opt_height = resolution * sqrt(height/width) - if not upscaling: - opt_width = min(width, opt_width) - opt_height = min(height, opt_height) - - # test 1, guaranteed to find a solution: shrink and crop - # 1.1: exact width - candidate_width = (opt_width // bucket_res) * bucket_res - candidate_height = ((height * candidate_width / width) // bucket_res) * bucket_res - loss = ((height * candidate_width / width) - candidate_height) * candidate_width - if (candidate_width, candidate_height) in self.preferred_sizes: - loss *= preferred_sizes_bonus - # 1.2: exact height - test_height = (opt_height // bucket_res) * bucket_res - test_width = ((width * test_height / height) // bucket_res) * bucket_res - test_loss = ((width * test_height / height) - test_width) * test_height - if (test_height, test_width) in self.preferred_sizes: - test_loss *= preferred_sizes_bonus - if test_loss < loss: - candidate_width = test_width - candidate_height = test_height - loss = test_loss - - # test 2, going bigger might still fit in the size budget due to cropping - # 2.1: exact width - for delta in range(1, 10): - test_width = (opt_width // bucket_res + delta) * bucket_res - test_height = ((height * test_width / width) // bucket_res) * bucket_res - if test_width * test_height > max_pixels: - break - if (test_width > width or test_height > height) and not upscaling: - break - test_loss = ((height * test_width / width) - test_height) * test_width - if (test_height, test_width) in self.preferred_sizes: - test_loss *= preferred_sizes_bonus - if test_loss < loss: - candidate_width = test_width - candidate_height = test_height - loss = test_loss - # 2.2: exact height - for delta in range(1, 10): - test_height = (opt_height // bucket_res + delta) * bucket_res - test_width = ((width * test_height / height) // bucket_res) * bucket_res - if test_width * test_height > max_pixels: - break - if (test_width > width or test_height > height) and not upscaling: - break - test_loss = ((width * test_height / height) - test_width) * test_height - if (test_height, test_width) in self.preferred_sizes: - test_loss *= preferred_sizes_bonus - if test_loss < loss: - candidate_width = test_width - candidate_height = test_height - loss = test_loss - - return int(candidate_width), int(candidate_height) + @Slot() + def set_filter(self, selected_table_item): + row = selected_table_item.row() + width = self.statistics_table.model().index(row, 0).data() + height = self.statistics_table.model().index(row, 1).data() + filter = f'target:{width}:{height}' + text = self.image_list.filter_line_edit.text() + if text != '': + self.image_list.filter_line_edit.setText(f'{filter} AND ({text})') + else: + self.image_list.filter_line_edit.setText(filter) + self.close() @Slot() def set_export_directory_path(self): @@ -662,19 +562,20 @@ def do_export(self): self.close() def get_image_list(self): + image_list_view = self.image_list.list_view if self.settings.value('export_filter') == ExportFilter.FILTERED: - images = self.image_list_view.proxy_image_list_model.sourceModel() + images = image_list_view.proxy_image_list_model.sourceModel() image_list = [] - for row in range(self.image_list_view.proxy_image_list_model.sourceModel().rowCount()): - source_index = self.image_list_view.proxy_image_list_model.sourceModel().index(row, 0) - proxy_index = self.image_list_view.proxy_image_list_model.mapFromSource(source_index) + for row in range(image_list_view.proxy_image_list_model.sourceModel().rowCount()): + source_index = image_list_view.proxy_image_list_model.sourceModel().index(row, 0) + proxy_index = image_list_view.proxy_image_list_model.mapFromSource(source_index) if proxy_index.isValid(): image_list.append(source_index.data(Qt.ItemDataRole.UserRole)) elif self.settings.value('export_filter') == ExportFilter.SELECTED: - images = self.image_list_view.proxy_image_list_model.sourceModel() - image_list = [image_index.data(Qt.ItemDataRole.UserRole) for image_index in self.image_list_view.get_selected_image_indices()] + images = image_list_view.proxy_image_list_model.sourceModel() + image_list = [image_index.data(Qt.ItemDataRole.UserRole) for image_index in image_list_view.get_selected_image_indices()] else: # ExportFilter.NONE - images = self.image_list_view.proxy_image_list_model.sourceModel() + images = image_list_view.proxy_image_list_model.sourceModel() image_list = [images.index(image_index).data(Qt.ItemDataRole.UserRole) for image_index in range(images.rowCount())] return image_list diff --git a/taggui/models/proxy_image_list_model.py b/taggui/models/proxy_image_list_model.py index 734ed6c6..e848a90a 100644 --- a/taggui/models/proxy_image_list_model.py +++ b/taggui/models/proxy_image_list_model.py @@ -6,7 +6,7 @@ from models.image_list_model import ImageListModel from utils.image import Image - +import utils.target_dimension as target_dimension class ProxyImageListModel(QSortFilterProxyModel): def __init__(self, image_list_model: ImageListModel, @@ -43,6 +43,14 @@ def does_image_match_filter(self, image: Image, return (len(dimension) == 2 and dimension[0] == str(image.dimensions[0]) and dimension[1] == str(image.dimensions[1])) + if filter_[0] == 'target': + # accept any dimension separator of [x:] + dimension = (filter_[1]).replace(':', 'x').split('x') + if image.target_dimensions == None: + image.target_dimensions = target_dimension.get(image.dimensions) + return (len(dimension) == 2 #and image.target_dimensions != None + and dimension[0] == str(image.target_dimensions[0]) + and dimension[1] == str(image.target_dimensions[1])) if filter_[1] == 'AND': return (self.does_image_match_filter(image, filter_[0]) and self.does_image_match_filter(image, filter_[2:])) diff --git a/taggui/utils/target_dimension.py b/taggui/utils/target_dimension.py new file mode 100644 index 00000000..b7cac27b --- /dev/null +++ b/taggui/utils/target_dimension.py @@ -0,0 +1,143 @@ +from math import sqrt +import re + +from utils.settings import DEFAULT_SETTINGS, get_settings + +# singleton data store +_preferred_sizes : list[tuple[int, int]] | None = None + +def prepare(aspect_ratios : list[tuple[int, int, int]] | None = None) -> list[tuple[int, int, int]] | None: + """ + Prepare by parsing the user supplied preferred sizes. + + Parameters + ---------- + aspect_ratios : list(tuple[int, int, int]) | None + A list of typical aspect ratios to take care of + + Return + ------ + The same list of aspect ratios (when supplied) but extrended by the real + aspect ratios of the preferred sizes. + """ + global _preferred_sizes + _preferred_sizes = [] + for res_str in re.split(r'\s*,\s*', get_settings().value('export_preferred_sizes') or ''): + try: + if res_str == '': + continue + size_str = res_str.split(':') + width = max(int(size_str[0]), int(size_str[1])) + height = min(int(size_str[0]), int(size_str[1])) + _preferred_sizes.append((width, height)) + if not width == height: + _preferred_sizes.append((height, width)) + if aspect_ratios != None: + # add exact aspect ratio of the preferred size to label it + # similar to the perfect one + aspect_ratio = width / height + for ar in aspect_ratios: + if abs(ar[2] - aspect_ratio) < 0.15: + aspect_ratios.append((ar[0], ar[1], aspect_ratio)) + break + except ValueError: + # Handle cases where the resolution string is not in the correct format + print(f'Warning: Invalid resolution format: {res_str}. Skipping.', + file=sys.stderr) + continue # Skip to the next resolution if there's an error + return aspect_ratios + +def get(dimensions: tuple[int, int]): + """ + Determine the dimensions of an image it should have when it is exported. + + Note: this gives the optimal answer and thus can be slower than the Kohya + bucket algorithm. + + Parameters + ---------- + dimensions : tuple[int, int] + The width and height of the image + """ + global _preferred_sizes + width, height = dimensions + # The target resolution of the AI model. The target image pixels + # will not exceed the square of this number + resolution = get_settings().value('export_resolution', defaultValue=DEFAULT_SETTINGS['export_resolution'], type=int) + # Is upscaling of images allowed? + upscaling = get_settings().value('export_upscaling', defaultValue=DEFAULT_SETTINGS['export_upscaling'], type=bool) + # The resolution of the buckets + bucket_res = get_settings().value('export_bucket_res_size', defaultValue=DEFAULT_SETTINGS['export_bucket_res_size'], type=int) + + if not _preferred_sizes: + prepare() + + if resolution == 0: + # no rescale in this case, only cropping + return ((width // bucket_res)*bucket_res, (height // bucket_res)*bucket_res) + + if width < bucket_res or height < bucket_res: + # it doesn't make sense to use such a small image. But we shouldn't + # patronize the user + return dimensions + + preferred_sizes_bonus = 0.4 # reduce the loss by this factor + + max_pixels = resolution * resolution + opt_width = resolution * sqrt(width/height) + opt_height = resolution * sqrt(height/width) + if not upscaling: + opt_width = min(width, opt_width) + opt_height = min(height, opt_height) + + # test 1, guaranteed to find a solution: shrink and crop + # 1.1: exact width + candidate_width = (opt_width // bucket_res) * bucket_res + candidate_height = ((height * candidate_width / width) // bucket_res) * bucket_res + loss = ((height * candidate_width / width) - candidate_height) * candidate_width + if (candidate_width, candidate_height) in _preferred_sizes: + loss *= preferred_sizes_bonus + # 1.2: exact height + test_height = (opt_height // bucket_res) * bucket_res + test_width = ((width * test_height / height) // bucket_res) * bucket_res + test_loss = ((width * test_height / height) - test_width) * test_height + if (test_height, test_width) in _preferred_sizes: + test_loss *= preferred_sizes_bonus + if test_loss < loss: + candidate_width = test_width + candidate_height = test_height + loss = test_loss + + # test 2, going bigger might still fit in the size budget due to cropping + # 2.1: exact width + for delta in range(1, 10): + test_width = (opt_width // bucket_res + delta) * bucket_res + test_height = ((height * test_width / width) // bucket_res) * bucket_res + if test_width * test_height > max_pixels: + break + if (test_width > width or test_height > height) and not upscaling: + break + test_loss = ((height * test_width / width) - test_height) * test_width + if (test_height, test_width) in _preferred_sizes: + test_loss *= preferred_sizes_bonus + if test_loss < loss: + candidate_width = test_width + candidate_height = test_height + loss = test_loss + # 2.2: exact height + for delta in range(1, 10): + test_height = (opt_height // bucket_res + delta) * bucket_res + test_width = ((width * test_height / height) // bucket_res) * bucket_res + if test_width * test_height > max_pixels: + break + if (test_width > width or test_height > height) and not upscaling: + break + test_loss = ((width * test_height / height) - test_width) * test_height + if (test_height, test_width) in _preferred_sizes: + test_loss *= preferred_sizes_bonus + if test_loss < loss: + candidate_width = test_width + candidate_height = test_height + loss = test_loss + + return int(candidate_width), int(candidate_height) diff --git a/taggui/widgets/image_list.py b/taggui/widgets/image_list.py index c8dab6ad..e2b3c142 100644 --- a/taggui/widgets/image_list.py +++ b/taggui/widgets/image_list.py @@ -48,7 +48,7 @@ def __init__(self): | QuotedString(quote_char="'", esc_char='\\') | Word(printables, exclude_chars='()')) - string_filter_keys = ['tag', 'caption', 'name', 'path', 'size'] + string_filter_keys = ['tag', 'caption', 'name', 'path', 'size', 'target'] string_filter_expressions = [Group(CaselessLiteral(key) + Suppress(':') + optionally_quoted_string) for key in string_filter_keys] From 8040003666317f73a0fbc44c321a6864771bab43 Mon Sep 17 00:00:00 2001 From: StableLlama Date: Sat, 15 Feb 2025 01:15:30 +0100 Subject: [PATCH 17/18] Remove little left over --- taggui/models/proxy_image_list_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/taggui/models/proxy_image_list_model.py b/taggui/models/proxy_image_list_model.py index e848a90a..45d6fcbe 100644 --- a/taggui/models/proxy_image_list_model.py +++ b/taggui/models/proxy_image_list_model.py @@ -48,7 +48,7 @@ def does_image_match_filter(self, image: Image, dimension = (filter_[1]).replace(':', 'x').split('x') if image.target_dimensions == None: image.target_dimensions = target_dimension.get(image.dimensions) - return (len(dimension) == 2 #and image.target_dimensions != None + return (len(dimension) == 2 and dimension[0] == str(image.target_dimensions[0]) and dimension[1] == str(image.target_dimensions[1])) if filter_[1] == 'AND': From 012b7ed9c8394690631cbd61ba7c8eb79617a4ce Mon Sep 17 00:00:00 2001 From: StableLlama Date: Sat, 15 Feb 2025 11:37:26 +0100 Subject: [PATCH 18/18] Add option to only export missing images --- taggui/dialogs/export_dialog.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/taggui/dialogs/export_dialog.py b/taggui/dialogs/export_dialog.py index 5420a2b5..7ecad394 100644 --- a/taggui/dialogs/export_dialog.py +++ b/taggui/dialogs/export_dialog.py @@ -442,6 +442,7 @@ def do_export(self): export_directory_path = Path(self.settings.value('export_directory_path', type=str)) export_keep_dir_structure = self.settings.value('export_keep_dir_structure', type=bool) no_overwrite = True + only_missing = True image_list = self.get_image_list() self.progress_bar = QProgressBar(self) @@ -462,8 +463,9 @@ def do_export(self): msgBox.setIcon(QMessageBox.Warning) msgBox.setWindowTitle('Path warning') msgBox.setText('The export directory path is not empty') - overwrite_button = msgBox.addButton('Overwrite', QMessageBox.YesRole) - rename_button = msgBox.addButton('Rename', QMessageBox.NoRole) + overwrite_button = msgBox.addButton('Overwrite', QMessageBox.DestructiveRole) + rename_button = msgBox.addButton('Rename', QMessageBox.YesRole) + only_missing_button = msgBox.addButton('Only missing', QMessageBox.AcceptRole) msgBox.addButton(QMessageBox.Cancel) msgBox.setDefaultButton(QMessageBox.Cancel) button = msgBox.exec_() @@ -471,6 +473,9 @@ def do_export(self): return if msgBox.clickedButton() == overwrite_button: no_overwrite = False + only_missing = False + if msgBox.clickedButton() == rename_button: + only_missing = False else: QMessageBox.critical( self, @@ -501,6 +506,9 @@ def do_export(self): export_path = export_directory_path / image_entry.path.name export_path = export_path.with_suffix(export_format.split(' ', 1)[0]) + if export_path.exists() and only_missing: + continue + if no_overwrite: stem = export_path.stem counter = 0