diff --git a/README.md b/README.md index c589a5be..19ae0fc1 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 @@ -142,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. @@ -164,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 @@ -253,3 +260,64 @@ 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. + +`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. + +`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. diff --git a/taggui/dialogs/export_dialog.py b/taggui/dialogs/export_dialog.py new file mode 100644 index 00000000..7ecad394 --- /dev/null +++ b/taggui/dialogs/export_dialog.py @@ -0,0 +1,589 @@ +from enum import Enum +from collections import defaultdict +import os +import io +from math import floor +from pathlib import Path +import shutil + +from PySide6.QtCore import Qt, Slot +from PySide6.QtGui import QColorSpace +from PySide6.QtWidgets import (QWidget, QDialog, QFileDialog, QGridLayout, + QLabel, QLineEdit, QPushButton, QTableWidget, + QTableWidgetItem, QProgressBar, QMessageBox, + 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: + 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'), + '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') +} + +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' +} + +class IccProfileList(str, Enum): + SRgb = 'sRGB' + SRgbLinear = 'sRGB (linear gamma)' + AdobeRgb = 'AdobeRGB' + DisplayP3 = 'DisplayP3' + ProPhotoRgb = 'ProPhotoRGB' + Bt2020 = 'BT.2020' + Bt2100Pq = 'BT.2100(PQ)' + Bt2100Hlg = 'BT.2100 (HLG)' + +class BucketStrategy(str, Enum): + CROP = 'crop' + SCALE = 'scale' + CROP_SCALE = 'crop and scale' + +class ExportDialog(QDialog): + def __init__(self, parent, image_list: ImageList): + """ + Main method to create the export dialog. + """ + super().__init__(parent) + self.image_list = image_list + self.settings = get_settings() + self.inhibit_statistics_update = True + self.setWindowTitle('Export') + self.layout = QVBoxLayout(self) + self.layout.setContentsMargins(20, 20, 20, 20) + self.layout.setSpacing(20) + + grid_layout = QGridLayout() + grid_layout.setColumnStretch(0, 0) + 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') + 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', + 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('Bucket resolution size (px)'), grid_row, 0, + Qt.AlignmentFlag.AlignRight) + self.bucket_res_size_spin_box = SettingsSpinBox( + 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.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('Preferred sizes'), grid_row, 0, + Qt.AlignmentFlag.AlignRight) + self.preferred_sizes_line_edit = SettingsLineEdit( + 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.") + 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') + 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) + 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() + format_layout.setContentsMargins(0, 0, 0, 0) + self.format_combo_box = SettingsComboBox(key='export_format') + 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) + format_layout.addWidget(QLabel('Quality'), + Qt.AlignmentFlag.AlignRight) + self.quality_spin_box = SettingsSpinBox( + 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.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) + # 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) + + 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) + self.export_directory_line_edit = SettingsLineEdit( + 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, + 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') + 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(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) + + self.layout.addLayout(grid_layout) + + 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) + self.show_megapixels() + self.inhibit_statistics_update = False + self.show_statistics() + + @Slot() + 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) + self.bucket_res_size_spin_box.setEnabled(True) + else: + self.inhibit_statistics_update = True + 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]) if do_value_change else 0 + self.bucket_res_size_spin_box.setEnabled(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 + + @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 + self.megapixels.setText(f"{megapixels:.3f}") + else: + self.megapixels.setText('-') + + @Slot() + 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) + 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) + elif export_format == ExportFormat.WEBP: + self.quality_spin_box.setValue(80) if do_value_change else 0 + self.quality_spin_box.setEnabled(True) + + @Slot() + 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: + self.quality_spin_box.setStyleSheet('') + + @Slot() + def show_statistics(self): + """ + Update the statistics table content. + """ + if self.inhibit_statistics_update: + return + + 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 = [ + (1, 1, 1), + (2, 1, 2/1), + (3, 2, 3/2), + (4, 3, 4/3), + (16, 9, 16/9), + (21, 9, 21/9), + ] + 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 = 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(), + 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: + 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}%")) + + @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): + """ + 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) + 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): + """ + 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) + no_overwrite = True + only_missing = True + + image_list = self.get_image_list() + self.progress_bar = QProgressBar(self) + self.progress_bar.setMinimum(0) + self.progress_bar.setMaximum(len(image_list)) + self.layout.addWidget(self.progress_bar) + + 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.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_() + if button == QMessageBox.Cancel: + return + if msgBox.clickedButton() == overwrite_button: + no_overwrite = False + only_missing = False + if msgBox.clickedButton() == rename_button: + only_missing = False + else: + QMessageBox.critical( + self, + 'Path error', + 'The export directory path does not exist' + ) + return + + 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)': + color_space = 'sRGB' + save_profile = False + bucket_strategy = self.settings.value('export_bucket_strategy', type=str) + + for image_index, image_entry in enumerate(self.get_image_list()): + self.progress_bar.setValue(image_index) + 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 export_path.exists() and only_missing: + continue + + 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 + + # 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 + 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 = 50, threshold = 0)) + 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)) + 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'), lossless=lossless) + 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(), lossless=lossless) + else: + final_image.save(export_path, format=ExportFormatDict[export_format], quality=quality, icc_profile=None, lossless=lossless) + self.close() + + def get_image_list(self): + image_list_view = self.image_list.list_view + if self.settings.value('export_filter') == ExportFilter.FILTERED: + images = image_list_view.proxy_image_list_model.sourceModel() + image_list = [] + 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 = 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 = 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/dialogs/settings_dialog.py b/taggui/dialogs/settings_dialog.py index 8fb2956b..ee8d6140 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,13 +58,11 @@ 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( - 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( @@ -76,8 +72,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/models/proxy_image_list_model.py b/taggui/models/proxy_image_list_model.py index b730c059..45d6fcbe 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, @@ -37,6 +37,20 @@ 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_[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 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:])) @@ -63,6 +77,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/utils/image.py b/taggui/utils/image.py index f2637da4..17091470 100644 --- a/taggui/utils/image.py +++ b/taggui/utils/image.py @@ -9,4 +9,5 @@ class Image: path: Path dimensions: tuple[int, int] | None tags: list[str] = field(default_factory=list) + target_dimensions: tuple[int, int] | None = None thumbnail: QIcon | None = None diff --git a/taggui/utils/settings.py b/taggui/utils/settings.py index 88e1f0f5..8d4ba547 100644 --- a/taggui/utils/settings.py +++ b/taggui/utils/settings.py @@ -9,7 +9,19 @@ 'tag_separator': ',', 'insert_space_after_tag_separator': True, 'autocomplete_tags': True, - 'models_directory_path': '' + 'models_directory_path': '', + 'export_filter': 'All images', + 'export_preset': 'SDXL, SD3, Flux', + 'export_resolution': 1024, + '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, + 'export_color_space': 'sRGB', + 'export_directory_path': '', + 'export_keep_dir_structure': False } 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())) 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 18af0eb0..e2b3c142 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', 'target'] 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] diff --git a/taggui/widgets/main_window.py b/taggui/widgets/main_window.py index 78a8cf98..35e360ec 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=self.image_list) + 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)