diff --git a/README.md b/README.md index f6021fcd..880ebc39 100644 --- a/README.md +++ b/README.md @@ -17,6 +17,11 @@ models. 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 +- Mark images manually or with the help of YOLO models to create masks for + inclusion and exclusion of image parts for masked training +- Crop images with advanced hints that respect relevant aspect ratios and + bucket sizes of the training scripts ## Installation @@ -137,11 +142,30 @@ apply: - `caption:cat` will match images that have `cat` anywhere in the caption. For example, images with the tag `orange cat` or the tag `catastrophe`. +- `marking`: Images that contain at least one marking with this label. It + doesn't matter whether it is a _hint_, _include_ or _exclude_ marking. Is also + doesn't matter whether the marking is on the exported image, i.e., cropped + away or not. + - `marking:cat:>0.789` when the label is followed by a colon with a + relational operator and a number then only those markings are matched + where the confidence fits to the specified number, in this example case + a `cat` marking must have a confidence number higher than 0.789. +- `crops`: Images that contain at least one marking that will be cropped on the + resulting image + - `crops:hand` at least one `hand` marking is not completely visible on the + exported image. +- `visible`: Images that contain the marking in the exported area, completely + or partly + - `visible:face`: At least a part of the face must be visible when the image + gets exported. - `name`: Images that contain the filter term in the file name - `name:cat` will match images such as `cat-1.jpg` or `large_cat.png`. - `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 +188,13 @@ comparison. caption. - `tokens:<=50` will match images that have 50 or fewer tokens in the caption. +- `stars`: Images with this rating in stars + - `stars:>=4` will match all images with 4 or 5 stars. +- `width` and `height`: will match images with the specified width or height. + - `width:>512` will match images where the width is greater than 512 pixels. + - `height:=1024` will match images where the height is exactly 1024 pixels. +- `area`: will match images with the specified amount of pixels. + - `area:<1048576` will match images with less than 1 mega pixels (1024²). ### Spaces and quotes @@ -253,3 +284,226 @@ 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`). + +## Cropping and masking with markings + +Next to tagging images with words and text, taggui supports visual tagging, +called _marking_. There are different types of marking and except the `crop` +they can be changed into each other. + +All markings are marking the pixels inside the border 🞑, not any pixels below +the border. + +TagGUI cropping and masking feature + +### Crop + +The _crop_, shown with a blue border, defines +the part of the image that will be exported. Depending on the export settings +a bucketing is likely configured. When the cropped area doesn't exactly fit +into a bucket as defined by the _Bucket resolution size_ and the _Bucket fitting +strategy_, it might be necessary to crop even more. This additional cropped +area is shown by a semitransparent red overlay. + +During the editing of the crop, you get hints to help with this task: + +In the image list, an overlay shows the size of the original images as well as +the size of the cropped area. The small number in the brackets shows how many +pixels the crop is extended to fit into a bucket. The target size is the image +size when exported, a checkmark is shown when this size is one of the preferred +sizes. And when the crop has a well-known aspect ratio, it is also shown. + +In the main image, lines are shown to quickly be able to select the best size. +The straight lines follow well-known aspect ratios. And a green line shows where +the size of the crop is big enough to fully use the native _Resolution_ of the +model. +These lines are intended for a quick orientation and thus placed at the +theoretical optimal position, taking the discrete nature of image pixels into +account the real optimal position might be in a slightly different place. +Especially when the _Bucket resolution size_ is changed it might be necessary +to optimize the crop when a pixel-perfect result is required. + +Inside the cropping area lines are shown to help with aesthetical alignment +of the content. These lines are drawn in the middle (stroked), following the +1/3rd rule (dashed) and in the golden ratio (dotted). +These lines can temporarily be hidden by pressing the `alt` key. + +### Hint + +A _hint_, shown with a gray border, is just a hint and has no effect on +exporting the image. +A _hint_ has a label where you can give it a name and which you can use for +filtering images which contain the given marking. +A _hint_ can be changed in an _exclude_ or an _include_. + +### Exclude + +An _exclude_, shown with a red border, is an area guaranteed to be masked +(made transparent) when the image is exported. +When _Latent size_ and _Quantize alpha channel_ are set and the _exclude_ area +doesn't fit, the mask will be grown to make sure that no excluded pixel will +stay unmasked. +An _exclude_ can be changed in an _include_ or a _hint_. + +### Include + +An _include_, shown with a green border, is an area included when the image is +exported. +When no _include_ is set, the full image (of course respecting the _crop_) is +included. +When an _include_ and an _exclude_ are overlapping, the _exclude_ takes +precedence. +And when _Latent size_ and _Quantize alpha channel_ are set and the _include_ +area doesn't fit, the mask will be shrunken to make sure that only included +pixels will stay unmasked. +An _include_ can be changed in a _hint_ or an _include_. + +### Working with markings + +Markings can be created by the use of the toolbar or by holding the `ctrl` key +to create a _hint_ or with `ctrl + alt` to create an _exclude_. +The position and size can be changed by dragging them to the desired place. +When the `shift` key is pressed during dragging the current part is snapped +to the next position that fits the current _export_ settings. +The marking label can be edited by clicking on it. +And the type can be changed in the toolbar or with a right mouse button click on +the marking. + +### Automatic marking detection + +When you have the path to the YOLO models configured in the _Settings..._ +dialog, you can use them to automatically detect features in the image that +can be used for marking. After selecting the relevant model in the drop-down +list, you see a table with all classes it can detect. In each line, you can +decide whether it should be ignored or added as a hint, exclude or include. +Next to the label, you can see the confidence the model had during detection. + +More detailed control about the minimal required confidence, the IoU and +the maximal number of detected markings per image can be set in the +advanced settings. + +Note: When you are already using image generation tools like ADetailer you +will most likely have already relevant YOLO models. There are many models +widely available when you are searching for "YOLO" or "ADetailer". + +## 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. + +`Latent size`: +The size of one latent space pixel in image pixels. + +`Quantize alpha channel`: +When exporting with _include_ or _exclude_ markings in an image format that +supports alpha masks (all, but not the classic JPEG) you can make sure that +the masks are aligned to the latent pixels that the trainer is using for +masked training. + +`Masking strategy`: +Select how the _include_ and _exclude_ masks are applied. They can be completely +ignored, replace the image content, or, when the image format supports +transparency, make that part of the image transparent. You can also opt to +get separate masking files. +The replacement content, or with _replace_ the invisible content hidden by +the transparency, is defined by _Masked content_. + +`Masked content`: +When exporting with _include_ or _exclude_ markings in an image format that +supports alpha masks (all, but not the classic JPEG), you can change the +content that is invisible due to the mask. +It is known that some masked content can slightly bleed through during the +masked training, so it can be beneficial when this (supposed to be invisible) +content is replaced. In simple cases (e.g., masking the face for cloth LoRA +training) the default "blur + noise" is a good choice. In hard cases, to really +hide the original data, "grey + noise" can be a good choice. "original" doesn't +modify the data behind the mask. + +`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 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 conversation 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. + +`Caption`: +Define how the tags are exported into a caption. +It can be a simple tag list, using the same _Tag separator_ as used for the +image tag files. +_Only first tag_ and _Only last tag_ are exporting exactly one tag. This can +be a good option when you are using tagging for image organizing and then +derive a complete captioning sentence out of those for modern models like SD3 +or Flux. +_Enumeration_ creates a natural language list out of the tags. I.e., it is a +comma separated list with the word "and" in the last place. Also, the +Oxford comma is used (a comma directly before the "and"). +_Prefixed enumeration_ is similar to _enumeration_, but the first tag is used +as a prefix and the remaining tags are translated to the enumeration. The +use case is also using tags to organize the images and then transfer those +to a caption by prepending it with something like "Photography of" which you'd +place then as the first tag. + +`Separate by #newline`: +Create a new line in the captioning file for every `#newline` tag. +When _Prefixed enumeration_ is selected the first tag is repeated for +every new line. + +This option is relevant when you want to use multi captioning, which is +supported by many trainers. + +`Fiter (other) hashtag (#) tags`: +When set, all tags that start with a hashtag (i.e., the "#" character) are not +included in the exported captions. + +This allows you to use tags for internal image organization like marking +images that you don't want to export or adding notes about an image into the +tag list. + +`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 function. diff --git a/images/doc/cropping.jpg b/images/doc/cropping.jpg new file mode 100644 index 00000000..a89db7fb Binary files /dev/null and b/images/doc/cropping.jpg differ diff --git a/images/show_label.png b/images/show_label.png new file mode 100644 index 00000000..7adc8603 Binary files /dev/null and b/images/show_label.png differ diff --git a/images/show_label.svg b/images/show_label.svg new file mode 100644 index 00000000..79cf035d --- /dev/null +++ b/images/show_label.svg @@ -0,0 +1,56 @@ + + + + + + + + + diff --git a/images/show_marking.png b/images/show_marking.png new file mode 100644 index 00000000..e5c8f366 Binary files /dev/null and b/images/show_marking.png differ diff --git a/images/show_marking.svg b/images/show_marking.svg new file mode 100644 index 00000000..2b55e279 --- /dev/null +++ b/images/show_marking.svg @@ -0,0 +1,66 @@ + + + + + + + + + diff --git a/images/show_marking_latent.png b/images/show_marking_latent.png new file mode 100644 index 00000000..0d17b322 Binary files /dev/null and b/images/show_marking_latent.png differ diff --git a/images/show_marking_latent.svg b/images/show_marking_latent.svg new file mode 100644 index 00000000..63efa827 --- /dev/null +++ b/images/show_marking_latent.svg @@ -0,0 +1,66 @@ + + + + + + + + + diff --git a/images/toggle_marking.png b/images/toggle_marking.png new file mode 100644 index 00000000..db6e1d73 Binary files /dev/null and b/images/toggle_marking.png differ diff --git a/images/toggle_marking.svg b/images/toggle_marking.svg new file mode 100644 index 00000000..29af277a --- /dev/null +++ b/images/toggle_marking.svg @@ -0,0 +1,156 @@ + + + + + + + + + + + + + + + + + + + + + diff --git a/requirements.txt b/requirements.txt index fc39e660..e1cb5b42 100644 --- a/requirements.txt +++ b/requirements.txt @@ -34,3 +34,6 @@ https://download.pytorch.org/whl/cu124/torch-2.6.0%2Bcu124-cp311-cp311-win_amd64 flash-attn==2.7.4.post1; platform_system == "Linux" https://github.com/kingbri1/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu124torch2.6.0cxx11abiFALSE-cp312-cp312-win_amd64.whl; platform_system == "Windows" and python_version == "3.12" https://github.com/kingbri1/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu124torch2.6.0cxx11abiFALSE-cp311-cp311-win_amd64.whl; platform_system == "Windows" and python_version == "3.11" + +# YOLO +ultralytics>=8.3.100 diff --git a/taggui/auto_captioning/auto_captioning_model.py b/taggui/auto_captioning/auto_captioning_model.py index 6ca7ce5f..63d2a2a4 100644 --- a/taggui/auto_captioning/auto_captioning_model.py +++ b/taggui/auto_captioning/auto_captioning_model.py @@ -16,20 +16,23 @@ from utils.image import Image -def replace_template_variable(match: re.Match, image: Image) -> str: +def replace_template_variable(match: re.Match, image: Image, skip_hash: bool) -> str: template_variable = match.group(0)[1:-1].lower() if template_variable == 'tags': - return ', '.join(image.tags) + if skip_hash: + return ', '.join([t for t in image.tags if not t.startswith('#')]) + else: + return ', '.join(image.tags) if template_variable == 'name': return image.path.stem if template_variable in ('directory', 'folder'): return image.path.parent.name -def replace_template_variables(text: str, image: Image) -> str: +def replace_template_variables(text: str, image: Image, skip_hash: bool) -> str: # Replace template variables inside curly braces that are not escaped. text = re.sub(r'(? str: - return captioning_start_datetime.strftime('%Y-%m-%d %H:%M:%S') - - def get_captioning_message(self, are_multiple_images_selected: bool, - captioning_start_datetime: datetime) -> str: - if are_multiple_images_selected: - captioning_start_datetime_string = ( - self.get_captioning_start_datetime_string( - captioning_start_datetime)) - return (f'Captioning... (device: {self.device}, start time: ' - f'{captioning_start_datetime_string})') - return f'Captioning... (device: {self.device})' + def get_generation_text() -> str: + return 'Captioning' @staticmethod def get_default_prompt() -> str: @@ -197,7 +190,8 @@ def format_prompt(prompt: str) -> str: def get_image_prompt(self, image: Image) -> str | None: if self.prompt: - image_prompt = replace_template_variables(self.prompt, image) + image_prompt = replace_template_variables(self.prompt, image, + self.skip_hash) else: self.prompt = self.get_default_prompt() image_prompt = self.prompt @@ -211,17 +205,19 @@ def get_input_text(self, image_prompt: str) -> str: text = image_prompt or self.caption_start return text - def load_image(self, image: Image) -> PilImage: + def load_image(self, image: Image, crop: bool) -> PilImage: pil_image = PilImage.open(image.path) # Rotate the image according to the orientation tag. pil_image = exif_transpose(pil_image) pil_image = pil_image.convert(self.image_mode) + if crop and image.crop is not None: + pil_image = pil_image.crop(image.crop.getCoords()) return pil_image - def get_model_inputs(self, image_prompt: str, - image: Image) -> BatchFeature | dict | np.ndarray: + def get_model_inputs(self, image_prompt: str, image: Image, + crop: bool) -> BatchFeature | dict | np.ndarray: text = self.get_input_text(image_prompt) - pil_image = self.load_image(image) + pil_image = self.load_image(image, crop) model_inputs = (self.processor(text=text, images=pil_image, return_tensors='pt') .to(self.device, **self.dtype_argument)) diff --git a/taggui/auto_captioning/captioning_thread.py b/taggui/auto_captioning/captioning_thread.py index 59e6daef..a35e799e 100644 --- a/taggui/auto_captioning/captioning_thread.py +++ b/taggui/auto_captioning/captioning_thread.py @@ -1,9 +1,6 @@ -from datetime import datetime from pathlib import Path -from time import perf_counter -from PIL import UnidentifiedImageError -from PySide6.QtCore import QModelIndex, QThread, Qt, Signal +from PySide6.QtCore import QModelIndex, Signal from auto_captioning.auto_captioning_model import AutoCaptioningModel from auto_captioning.models_list import get_model_class @@ -11,6 +8,7 @@ from utils.enums import CaptionPosition from utils.image import Image from utils.settings import get_tag_separator +from utils.ModelThread import ModelThread def add_caption_to_tags(tags: list[str], caption: str, @@ -36,113 +34,50 @@ def add_caption_to_tags(tags: list[str], caption: str, return tags -def format_duration(seconds: float) -> str: - seconds_per_minute = 60 - seconds_per_hour = 60 * seconds_per_minute - seconds_per_day = 24 * seconds_per_hour - if seconds < seconds_per_minute: - return f'{seconds:.1f} seconds' - if seconds < seconds_per_hour: - minutes = seconds / seconds_per_minute - return f'{minutes:.1f} minutes' - if seconds < seconds_per_day: - hours = seconds / seconds_per_hour - return f'{hours:.1f} hours' - days = seconds / seconds_per_day - return f'{days:.1f} days' - - -class CaptioningThread(QThread): - text_outputted = Signal(str) - clear_console_text_edit_requested = Signal() +class CaptioningThread(ModelThread): # The image index, the caption, and the tags with the caption added. The # third parameter must be declared as `list` instead of `list[str]` for it # to work. caption_generated = Signal(QModelIndex, str, list) - progress_bar_update_requested = Signal(int) def __init__(self, parent, image_list_model: ImageListModel, selected_image_indices: list[QModelIndex], caption_settings: dict, tag_separator: str, models_directory_path: Path | None): - super().__init__(parent) - self.image_list_model = image_list_model - self.selected_image_indices = selected_image_indices + super().__init__(parent, image_list_model, selected_image_indices) self.caption_settings = caption_settings self.tag_separator = tag_separator self.models_directory_path = models_directory_path - self.is_error = False - self.is_canceled = False + self.model: AutoCaptioningModel | None = None - def run_captioning(self): + def load_model(self): model_id = self.caption_settings['model_id'] model_class = get_model_class(model_id) - model: AutoCaptioningModel = model_class( + self.model = model_class( captioning_thread_=self, caption_settings=self.caption_settings) - error_message = model.get_error_message() - if error_message: + self.error_message = self.model.get_error_message() + if self.error_message: self.is_error = True - self.clear_console_text_edit_requested.emit() - print(error_message) - return - model.load_processor_and_model() - model.monkey_patch_after_loading() - if self.is_canceled: - print('Canceled captioning.') return - self.clear_console_text_edit_requested.emit() - selected_image_count = len(self.selected_image_indices) - are_multiple_images_selected = selected_image_count > 1 - captioning_start_datetime = datetime.now() - captioning_message = model.get_captioning_message( - are_multiple_images_selected, captioning_start_datetime) - print(captioning_message) - caption_position = self.caption_settings['caption_position'] - for i, image_index in enumerate(self.selected_image_indices): - start_time = perf_counter() - if self.is_canceled: - print('Canceled captioning.') - return - image: Image = self.image_list_model.data(image_index, - Qt.ItemDataRole.UserRole) - image_prompt = model.get_image_prompt(image) - try: - model_inputs = model.get_model_inputs(image_prompt, image) - except UnidentifiedImageError: - print(f'Skipping {image.path.name} because its file format is ' - 'not supported or it is a corrupted image.') - continue - caption, console_output_caption = model.generate_caption( - model_inputs, image_prompt) - tags = add_caption_to_tags(image.tags, caption, caption_position) - self.caption_generated.emit(image_index, caption, tags) - if are_multiple_images_selected: - self.progress_bar_update_requested.emit(i + 1) - if i == 0 and not are_multiple_images_selected: - self.clear_console_text_edit_requested.emit() - if console_output_caption is None: - console_output_caption = caption - print(f'{image.path.name} ({perf_counter() - start_time:.1f} s):\n' - f'{console_output_caption}') - if are_multiple_images_selected: - captioning_end_datetime = datetime.now() - total_captioning_duration = ((captioning_end_datetime - - captioning_start_datetime) - .total_seconds()) - average_captioning_duration = (total_captioning_duration / - selected_image_count) - print(f'Finished captioning {selected_image_count} images in ' - f'{format_duration(total_captioning_duration)} ' - f'({average_captioning_duration:.1f} s/image) at ' - f'{captioning_end_datetime.strftime("%Y-%m-%d %H:%M:%S")}.') + self.model.load_processor_and_model() + self.model.monkey_patch_after_loading() + self.device = self.model.device + self.text = { + 'Generating': self.model.get_generation_text(), + 'generating': 'captioning' + } - def run(self): - try: - self.run_captioning() - except Exception as exception: - self.is_error = True - # Show the error message in the console text edit. - raise exception + def get_model_inputs(self, image: Image): + image_prompt = self.model.get_image_prompt(image) + crop = self.caption_settings['limit_to_crop'] + return image_prompt, self.model.get_model_inputs(image_prompt, + image, + crop) - def write(self, text: str): - self.text_outputted.emit(text) + def generate_output(self, image_index, image: Image, image_prompt: str | None, model_inputs) -> str: + caption_position = self.caption_settings['caption_position'] + caption, console_output_caption = self.model.generate_caption( + model_inputs, image_prompt) + tags = add_caption_to_tags(image.tags, caption, caption_position) + self.caption_generated.emit(image_index, caption, tags) + return console_output_caption diff --git a/taggui/auto_captioning/models/cogvlm.py b/taggui/auto_captioning/models/cogvlm.py index 1130c624..e468b6ad 100644 --- a/taggui/auto_captioning/models/cogvlm.py +++ b/taggui/auto_captioning/models/cogvlm.py @@ -65,9 +65,9 @@ def get_input_text(self, image_prompt: str) -> str: # `caption_start` is added later. return image_prompt - def get_model_inputs(self, image_prompt: str, image: Image) -> dict: + def get_model_inputs(self, image_prompt: str, image: Image, crop: bool) -> dict: text = self.get_input_text(image_prompt) - pil_image = self.load_image(image) + pil_image = self.load_image(image, crop) model_inputs = self.model.build_conversation_input_ids( self.processor, query=text, images=[pil_image], template_version=None) diff --git a/taggui/auto_captioning/models/cogvlm2.py b/taggui/auto_captioning/models/cogvlm2.py index e5b6ac80..52584ce7 100644 --- a/taggui/auto_captioning/models/cogvlm2.py +++ b/taggui/auto_captioning/models/cogvlm2.py @@ -65,9 +65,9 @@ def get_default_prompt() -> str: def format_prompt(prompt: str) -> str: return f'Question: {prompt} Answer:' - def get_model_inputs(self, image_prompt: str, image: Image) -> dict: + def get_model_inputs(self, image_prompt: str, image: Image, crop: bool) -> dict: text = self.get_input_text(image_prompt) - pil_image = self.load_image(image) + pil_image = self.load_image(image, crop) image_size = self.model.config.vision_config['image_size'] patch_size = self.model.config.vision_config['patch_size'] vision_tokens_count = ((image_size // patch_size // 2) diff --git a/taggui/auto_captioning/models/moondream.py b/taggui/auto_captioning/models/moondream.py index 448d44ae..d569a1c3 100644 --- a/taggui/auto_captioning/models/moondream.py +++ b/taggui/auto_captioning/models/moondream.py @@ -63,9 +63,9 @@ def patch_source_code(self) -> bool: del sys.modules[phi_module.__name__] return True - def get_model_inputs(self, image_prompt: str, image: Image) -> dict: + def get_model_inputs(self, image_prompt: str, image: Image, crop: bool) -> dict: text = self.get_input_text(image_prompt) - pil_image = self.load_image(image) + pil_image = self.load_image(image, crop) encoded_image = self.model.encode_image(pil_image) eos_tokens_ids = self.processor('').input_ids inputs_embeds = self.model.input_embeds(text, encoded_image, @@ -102,9 +102,9 @@ def get_model_load_arguments(self) -> dict: arguments['revision'] = MOONDREAM2_REVISION return arguments - def get_model_inputs(self, image_prompt: str, image: Image) -> dict: + def get_model_inputs(self, image_prompt: str, image: Image, crop: bool) -> dict: text = self.get_input_text(image_prompt) - pil_image = self.load_image(image) + pil_image = self.load_image(image, crop) encoded_image = self.model.encode_image(pil_image) inputs_embeds = self.model.input_embeds(text, encoded_image, self.processor) diff --git a/taggui/auto_captioning/models/phi_3_vision.py b/taggui/auto_captioning/models/phi_3_vision.py index fc3429d9..f77a3420 100644 --- a/taggui/auto_captioning/models/phi_3_vision.py +++ b/taggui/auto_captioning/models/phi_3_vision.py @@ -27,8 +27,8 @@ def get_input_text(self, image_prompt: str) -> str: return image_prompt + self.caption_start def get_model_inputs(self, image_prompt: str, - image: Image) -> BatchFeature: - model_inputs = super().get_model_inputs(image_prompt, image) + image: Image, crop: bool) -> BatchFeature: + model_inputs = super().get_model_inputs(image_prompt, image, crop) self.input_length = model_inputs['input_ids'].shape[1] return model_inputs diff --git a/taggui/auto_captioning/models/wd_tagger.py b/taggui/auto_captioning/models/wd_tagger.py index 2b5f3619..db05e27b 100644 --- a/taggui/auto_captioning/models/wd_tagger.py +++ b/taggui/auto_captioning/models/wd_tagger.py @@ -108,18 +108,12 @@ def get_processor(self): def get_model(self): return WdTaggerModel(self.model_id) - def get_captioning_message(self, are_multiple_images_selected: bool, - captioning_start_datetime: datetime) -> str: - if are_multiple_images_selected: - captioning_start_datetime_string = ( - self.get_captioning_start_datetime_string( - captioning_start_datetime)) - return (f'Generating tags... (start time: ' - f'{captioning_start_datetime_string})') - return 'Generating tags...' - - def get_model_inputs(self, image_prompt: str, image: Image) -> np.ndarray: - pil_image = self.load_image(image) + @staticmethod + def get_generation_text() -> str: + return 'Generating tags' + + def get_model_inputs(self, image_prompt: str, image: Image, crop: bool) -> np.ndarray: + pil_image = self.load_image(image, crop) # Add a white background to the image in case it has transparent areas. canvas = PilImage.new('RGBA', pil_image.size, (255, 255, 255)) canvas.alpha_composite(pil_image) diff --git a/taggui/auto_captioning/models/xcomposer2.py b/taggui/auto_captioning/models/xcomposer2.py index aa7308eb..283e9c64 100644 --- a/taggui/auto_captioning/models/xcomposer2.py +++ b/taggui/auto_captioning/models/xcomposer2.py @@ -107,9 +107,9 @@ def format_prompt(prompt: str) -> str: def get_input_text(self, image_prompt: str) -> str: return image_prompt + self.caption_start - def get_model_inputs(self, image_prompt: str, image: Image) -> dict: + def get_model_inputs(self, image_prompt: str, image: Image, crop: bool) -> dict: text = self.get_input_text(image_prompt) - pil_image = self.load_image(image) + pil_image = self.load_image(image, crop) input_embeddings_parts = [] image_mask_parts = [] processed_image = self.model.vis_processor(pil_image).unsqueeze(0).to( @@ -189,7 +189,7 @@ class Xcomposer2_4khd(Xcomposer2): def monkey_patch_after_loading(self): return - def load_image(self, image: Image) -> PilImage: - pil_image = super().load_image(image) + def load_image(self, image: Image, crop: bool) -> PilImage: + pil_image = super().load_image(image, crop) pil_image = hd_transform(pil_image) return pil_image diff --git a/taggui/auto_marking/marking_thread.py b/taggui/auto_marking/marking_thread.py new file mode 100644 index 00000000..3f3663f2 --- /dev/null +++ b/taggui/auto_marking/marking_thread.py @@ -0,0 +1,66 @@ +from PySide6.QtCore import QModelIndex, QThread, Signal, Qt + +from ultralytics import YOLO + +from models.image_list_model import ImageListModel +from utils.image import Image +from utils.ModelThread import ModelThread + + +class MarkingThread(ModelThread): + # The image index, the caption, and the tags with the caption added. The + # third parameter must be declared as `list` instead of `list[str]` for it + # to work. + marking_generated = Signal(QModelIndex, list) + + def __init__(self, parent, image_list_model: ImageListModel, + selected_image_indices: list[QModelIndex], + marking_settings: dict): + super().__init__(parent, image_list_model, selected_image_indices) + self.marking_settings = marking_settings + self.model: YOLO | None = None + self.text = { + 'Generating': 'Marking', + 'generating': 'marking' + } + + def load_model(self): + if not self.model: + self.error_message = 'Model not preloaded.' + self.is_error = True + pass + + def preload_model(self): + if self.marking_settings['model_path'] is None: + self.error_message = 'Model path not set' + self.is_error = True + self.model = None + return + self.model = YOLO(self.marking_settings['model_path']) + + def get_model_inputs(self, image: Image): + return '', {} + + def generate_output(self, image_index, image: Image, image_prompt, model_inputs) -> str: + if len(self.marking_settings['classes']) == 0: + return 'No classes to mark selected.' + classes = list(self.marking_settings['classes'].keys()) + results = self.model.predict(source=image.path, + conf=self.marking_settings['conf'], + iou=self.marking_settings['iou'], + max_det=self.marking_settings['max_det'], + classes=classes, + retina_masks=True) + markings = [] + for r in results: + for box, class_id, confidence in zip(r.boxes.xyxy.to('cpu').tolist(), + r.boxes.cls.to('cpu').tolist(), + r.boxes.conf.to('cpu').tolist()): + marking = self.marking_settings['classes'].get(class_id) + if marking is not None: + markings.append({'box': box, + 'label': marking[0], + 'type': marking[1], + 'confidence': round(confidence, 3)}) + self.marking_generated.emit(image_index, markings) + return f'Found {len(markings)} marking(s).' diff --git a/taggui/dialogs/batch_reorder_tags_dialog.py b/taggui/dialogs/batch_reorder_tags_dialog.py index c6f041f1..cd6bf51a 100644 --- a/taggui/dialogs/batch_reorder_tags_dialog.py +++ b/taggui/dialogs/batch_reorder_tags_dialog.py @@ -48,6 +48,19 @@ def __init__(self, parent, image_list_model: ImageListModel, top_buttons_layout.addWidget(shuffle_button) top_layout.addLayout(top_buttons_layout) horizontal_line = HorizontalLine() + middle_layout = QHBoxLayout() + middle_layout.setContentsMargins(20, 20, 20, 20) + middle_layout.setSpacing(20) + separate_newline_check_box = SettingsBigCheckBox( + key='reorder_tags_separate_newline', default=True) + separate_newline_check_box.setText('Separate by #newline') + middle_layout.addWidget(separate_newline_check_box) + sort_sentences_button = QPushButton('Sort Sentence Tags to Bottom') + sort_sentences_button.clicked.connect( + lambda: self.image_list_model.sort_sentences_down( + separate_newline_check_box.isChecked())) + middle_layout.addWidget(sort_sentences_button) + horizontal_line2 = HorizontalLine() bottom_layout = QVBoxLayout() bottom_layout.setContentsMargins(20, 20, 20, 20) bottom_layout.setSpacing(20) @@ -65,6 +78,8 @@ def __init__(self, parent, image_list_model: ImageListModel, bottom_layout.addWidget(self.move_tags_button) layout.addLayout(top_layout) layout.addWidget(horizontal_line) + layout.addLayout(middle_layout) + layout.addWidget(horizontal_line2) layout.addLayout(bottom_layout) self.move_tags_line_edit.textChanged.emit( diff --git a/taggui/dialogs/caption_multiple_images_dialog.py b/taggui/dialogs/caption_multiple_images_dialog.py index 13abbe0e..0b3e4fee 100644 --- a/taggui/dialogs/caption_multiple_images_dialog.py +++ b/taggui/dialogs/caption_multiple_images_dialog.py @@ -3,9 +3,10 @@ class CaptionMultipleImagesDialog(ConfirmationDialog): - def __init__(self, selected_image_count: int): - title = 'Generate Captions' - question = f'Caption {selected_image_count} selected images?' + def __init__(self, selected_image_count: int, caption_singular = 'Caption', + caption_plural = 'Captions'): + title = f'Generate {caption_plural}' + question = f'{caption_singular} {selected_image_count} selected images?' super().__init__(title=title, question=question) self.show_alert_check_box = SettingsBigCheckBox( key='show_alert_when_captioning_finished', default=True, diff --git a/taggui/dialogs/export_dialog.py b/taggui/dialogs/export_dialog.py new file mode 100644 index 00000000..2464eb32 --- /dev/null +++ b/taggui/dialogs/export_dialog.py @@ -0,0 +1,865 @@ +from collections import defaultdict +import os +import io +from math import ceil, floor +from pathlib import Path +import shutil +import numpy as np + +from PySide6.QtCore import QRect, QSize, Qt, Slot +from PySide6.QtGui import QColorSpace +from PySide6.QtWidgets import (QWidget, QDialog, QFileDialog, QGridLayout, + QHeaderView, QLabel, QPushButton, QTableWidget, + QTableWidgetItem, QProgressBar, QMessageBox, + QVBoxLayout, QHBoxLayout, QAbstractItemView) +from PIL import Image, ImageFilter, ImageCms + +from utils.enums import (ExportFilter, Presets, MaskingStrategy, MaskedContent, + ExportFormat, ExportFormatDict, IccProfileList, + BucketStrategy, CaptionStrategy, HashNewlineHandling) +from utils.settings import DEFAULT_SETTINGS, settings +from utils.image import ImageMarking +from utils.settings_widgets import (SettingsBigCheckBox, SettingsLineEdit, + SettingsSpinBox, SettingsComboBox) +import utils.target_dimension as target_dimension +from utils.grid import Grid +from widgets.image_list import ImageList + +try: + import pillow_jxl +except ModuleNotFoundError: + pass + + +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.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) + resolution_widget = QWidget() + resolution_layout = QHBoxLayout() + resolution_layout.setContentsMargins(0, 0, 0, 0) + 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.valueChanged.connect(self.show_megapixels) + self.resolution_spin_box.valueChanged.connect(self.show_statistics) + resolution_layout.addWidget(self.resolution_spin_box, + Qt.AlignmentFlag.AlignLeft) + + resolution_layout.addWidget(QLabel('Image size (megapixel)'), + Qt.AlignmentFlag.AlignRight) + self.megapixels = QLabel('-') + resolution_layout.addWidget(self.megapixels, + Qt.AlignmentFlag.AlignLeft) + resolution_widget.setLayout(resolution_layout) + grid_layout.addWidget(resolution_widget, 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.valueChanged.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('Latent size (px)'), grid_row, 0, + Qt.AlignmentFlag.AlignRight) + latent_widget = QWidget() + latent_layout = QHBoxLayout() + latent_layout.setContentsMargins(0, 0, 0, 0) + self.latent_size_spin_box = SettingsSpinBox( + key='export_latent_size', + minimum=1, maximum=256) + self.latent_size_spin_box.setToolTip( + 'Size of one latent space pixel in image space pixels') + latent_layout.addWidget(self.latent_size_spin_box, + Qt.AlignmentFlag.AlignLeft) + latent_layout.addWidget(QLabel('Quantize alpha channel'), + Qt.AlignmentFlag.AlignRight) + self.quantize_alpha_check_box = SettingsBigCheckBox(key='export_quantize_alpha') + self.quantize_alpha_check_box.setToolTip( + 'Align the masks due to include and exclude marking with the\n' + 'latent space pixels.\n' + 'Only available when the output format supports an alpha channel.') + latent_layout.addWidget(self.quantize_alpha_check_box, + Qt.AlignmentFlag.AlignLeft) + latent_widget.setLayout(latent_layout) + grid_layout.addWidget(latent_widget, grid_row, 1, + Qt.AlignmentFlag.AlignLeft) + + grid_row += 1 + grid_layout.addWidget(QLabel('Masking strategy'), grid_row, 0, + Qt.AlignmentFlag.AlignRight) + masking_widget = QWidget() + masking_layout = QHBoxLayout() + masking_layout.setContentsMargins(0, 0, 0, 0) + self.masking_strategy_combo_box = SettingsComboBox(key='export_masking_strategy') + self.masking_strategy_combo_box.addItems(list(MaskingStrategy)) + self.masking_strategy_combo_box.setToolTip( + 'Ignore the exclude masks, replace the content of the masks or\n' + 'remove it (make it transparent) when supported by the image format.') + masking_layout.addWidget(self.masking_strategy_combo_box, + Qt.AlignmentFlag.AlignLeft) + masking_layout.addWidget(QLabel('Masked content'), + Qt.AlignmentFlag.AlignRight) + self.masked_content_combo_box = SettingsComboBox(key='export_masked_content') + self.masked_content_combo_box.addItems(list(MaskedContent)) + masking_layout.addWidget(self.masked_content_combo_box, + Qt.AlignmentFlag.AlignLeft) + masking_widget.setLayout(masking_layout) + grid_layout.addWidget(masking_widget, 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(600) + 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: asymmetric 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.valueChanged.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 = settings.value('export_format', type=str) + current_quality = 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('Caption'), grid_row, 0, + Qt.AlignmentFlag.AlignRight) + caption_algorithm_combo_box = SettingsComboBox(key='export_caption_algorithm') + caption_algorithm_combo_box.addItems(list(CaptionStrategy)) + caption_algorithm_combo_box.setToolTip( + 'Define how the tags should be exported:\n' + 'tag list - just like the tag text files\n' + 'only first tag, only last tag - only this one tag\n' + 'enumeration - natural language list with commas\n' + 'prefixed enumeration - first tag directly followed by enumeration') + grid_layout.addWidget(caption_algorithm_combo_box, grid_row, 1, + Qt.AlignmentFlag.AlignLeft) + + grid_row += 1 + grid_layout.addWidget(QLabel('Fiter hashtag (#) tags'), grid_row, 0, + Qt.AlignmentFlag.AlignRight) + caption_hashtag_widget = QWidget() + caption_hashtag_layout = QHBoxLayout() + caption_hashtag_layout.setContentsMargins(0, 0, 0, 0) + self.filter_hashtag_check_box = SettingsBigCheckBox(key='export_filter_hashtag') + self.filter_hashtag_check_box.setToolTip( + 'Do not export tags that start with a hashtag (#)') + caption_hashtag_layout.addWidget(self.filter_hashtag_check_box, + Qt.AlignmentFlag.AlignLeft) + self.separate_newline_combo_box = SettingsComboBox(key='export_separate_newline') + self.separate_newline_combo_box.addItems(list(HashNewlineHandling)) + self.separate_newline_combo_box.setToolTip( + 'Create a multi-caption file where each line contains a caption for\n' + 'the image. The tags are split by the tag "#newline" and the\n' + 'captioning algorith is used for each group. Only for prefixed\n' + 'enumeration the first tag is used repeatedly for each group.') + caption_hashtag_layout.addWidget(QLabel('Handle #newline'), + Qt.AlignmentFlag.AlignRight) + caption_hashtag_layout.addWidget(self.separate_newline_combo_box, + Qt.AlignmentFlag.AlignLeft) + caption_hashtag_widget.setLayout(caption_hashtag_layout) + grid_layout.addWidget(caption_hashtag_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') + self.export_directory_line_edit.setMinimumWidth(600) + 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.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch) + self.statistics_table.setMinimumWidth(600) + self.statistics_table.setMinimumHeight(100) + 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) + + image_list = self.get_image_list() + self.progress_bar = QProgressBar(self) + self.progress_bar.setMinimum(0) + self.progress_bar.setMaximum(len(image_list)) + self.progress_bar.hide() + self.layout.addWidget(self.progress_bar) + + # 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) + self.latent_size_spin_box.setEnabled(True) + else: + self.inhibit_statistics_update = True + if do_value_change: + self.resolution_spin_box.setValue(preset[0]) + self.bucket_res_size_spin_box.setValue(preset[1]) + self.latent_size_spin_box.setValue(preset[2]) + self.resolution_spin_box.setEnabled(False) + self.bucket_res_size_spin_box.setEnabled(False) + self.latent_size_spin_box.setEnabled(False) + self.preferred_sizes_line_edit.setText(preset[3]) if do_value_change else 0 + self.inhibit_statistics_update = inhibit_statistics_update_current + self.show_statistics() + + @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. + """ + replace_mask_item = self.masking_strategy_combo_box.model().item(2) + if export_format == ExportFormat.JPG: + self.quality_spin_box.setValue(75) if do_value_change else 0 + self.quality_spin_box.setEnabled(True) + self.quantize_alpha_check_box.setEnabled(False) + if self.masking_strategy_combo_box.currentIndex() == 2: + self.masking_strategy_combo_box.setCurrentIndex(1) + replace_mask_item.setFlags(replace_mask_item.flags() & ~Qt.ItemIsEnabled) + if export_format == ExportFormat.JPGXL: + self.quality_spin_box.setValue(100) if do_value_change else 0 + self.quality_spin_box.setEnabled(True) + self.quantize_alpha_check_box.setEnabled(True) + replace_mask_item.setFlags(replace_mask_item.flags() | Qt.ItemIsEnabled) + elif export_format == ExportFormat.PNG: + self.quality_spin_box.setValue(100) if do_value_change else 0 + self.quality_spin_box.setEnabled(False) + self.quantize_alpha_check_box.setEnabled(True) + replace_mask_item.setFlags(replace_mask_item.flags() | Qt.ItemIsEnabled) + elif export_format == ExportFormat.WEBP: + self.quality_spin_box.setValue(80) if do_value_change else 0 + self.quality_spin_box.setEnabled(True) + self.quantize_alpha_check_box.setEnabled(True) + replace_mask_item.setFlags(replace_mask_item.flags() | Qt.ItemIsEnabled) + + @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 = settings.value('export_resolution', type=int) + + image_list = self.get_image_list() + image_dimensions = defaultdict(int) + for this_image in image_list: + if this_image.crop is not None: + this_image.target_dimension = target_dimension.get( + this_image.crop.size()) + else: + this_image.target_dimension = target_dimension.get( + QSize(*this_image.dimensions)) + image_dimensions[this_image.target_dimension.toTuple()] += 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() + ar = target_dimension.get_noteable_aspect_ratio(width, height) + notable_aspect_ratio = f' ({ar[0]}:{ar[1]})' if ar is not None else '' + 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: QTableWidgetItem): + row = selected_table_item.row() + width = self.statistics_table.model().index(row, 0).data() + height = self.statistics_table.model().index(row, 1).data() + add_filter = f'target:{width}:{height}' + text = self.image_list.filter_line_edit.text().strip() + if text != '': + text = text if text.startswith('(') else f'({text})' + self.image_list.filter_line_edit.setText(f'{add_filter} AND {text}') + else: + self.image_list.filter_line_edit.setText(add_filter) + self.close() + + @Slot() + def set_export_directory_path(self): + """ + Set the path of the directory to export to. + """ + export_directory_path = settings.value( + 'export_directory_path', + defaultValue=DEFAULT_SETTINGS['export_directory_path'], type=str) + if export_directory_path: + initial_directory_path = export_directory_path + elif settings.contains('directory_path'): + initial_directory_path = settings.value('directory_path', type=str) + 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 = settings.value('directory_path', type=str) + export_directory_path = Path(settings.value('export_directory_path', type=str)) + export_keep_dir_structure = settings.value('export_keep_dir_structure', type=bool) + no_overwrite = True + only_missing = True + refresh_tags = False + + try: + 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') + refresh_button = msgBox.addButton('Refresh', QMessageBox.ApplyRole) + refresh_button.setToolTip('Export only missing images, but update all captions') + overwrite_button = msgBox.addButton('Overwrite', QMessageBox.DestructiveRole) + overwrite_button.setToolTip('Overwrite all existing files') + rename_button = msgBox.addButton('Rename', QMessageBox.YesRole) + rename_button.setToolTip('Export with a new name') + only_missing_button = msgBox.addButton('Only missing', QMessageBox.AcceptRole) + only_missing_button.setToolTip('Export only missing images') + msgBox.addButton(QMessageBox.Cancel) + msgBox.setDefaultButton(refresh_button) + button = msgBox.exec_() + if button == QMessageBox.Cancel: + return + if msgBox.clickedButton() == refresh_button: + no_overwrite = False + refresh_tags = True + if msgBox.clickedButton() == overwrite_button: + no_overwrite = False + only_missing = False + if msgBox.clickedButton() == rename_button: + only_missing = False + else: + button = QMessageBox.critical( + self, + 'Path error', + 'The export directory path does not exist. Create it?', + QMessageBox.Ok, QMessageBox.Cancel) + if button == QMessageBox.Cancel: + return + os.makedirs(export_directory_path) + except Exception as e: + QMessageBox.critical(self, 'Path error', f'Error: {e}') + return + + self.progress_bar.show() + + tag_separator = settings.value('tag_separator', type=str) + if settings.value('insert_space_after_tag_separator', type=bool): + tag_separator += ' ' + caption_algorithm = settings.value('export_caption_algorithm', type=str) + separate_newline = settings.value('export_separate_newline', type=str) + filter_hashtag = settings.value('export_filter_hashtag', type=bool) + quantize_alpha = settings.value('export_quantize_alpha', type=bool) + masking_strategy = settings.value('export_masking_strategy', type=str) + masked_content = settings.value('export_masked_content', type=str) + export_format = settings.value('export_format', type=str) + quality = settings.value('export_quality', type=int) + color_space = settings.value('export_color_space', type=str) + save_profile = True + if color_space == 'sRGB (implicit, without profile)': + color_space = 'sRGB' + save_profile = False + + if masking_strategy == MaskingStrategy.MASK_FILE: + export_mask_directory_path = export_directory_path / 'mask' + export_directory_path = export_directory_path / 'image' + else: + export_mask_directory_path = Path() + + image_list = self.get_image_list() + self.progress_bar.setMaximum(len(image_list)) + for image_index, image_entry in enumerate(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_mask_path = export_mask_directory_path / relative_path + else: + export_path = export_directory_path / image_entry.path.name + export_mask_path = export_mask_directory_path / image_entry.path.name + export_path.parent.mkdir(parents=True, exist_ok=True) + export_path = export_path.with_suffix(export_format.split(' ', 1)[0]) + if masking_strategy == MaskingStrategy.MASK_FILE: + export_mask_path.parent.mkdir(parents=True, exist_ok=True) + export_mask_path = export_mask_path.with_suffix(export_format.split(' ', 1)[0]) + mask_exists = export_mask_path.exists() + else: + mask_exists = False + + image_exists = export_path.exists() + if (image_exists or mask_exists) and only_missing and not refresh_tags: + continue + + if no_overwrite: + stem = export_path.stem + counter = 0 + while (image_exists or mask_exists): + export_path = export_path.parent / f'{stem}_{counter}{export_path.suffix}' + image_exists = export_path.exists() + if masking_strategy == MaskingStrategy.MASK_FILE: + export_mask_path = export_mask_path.parent / f'{stem}_{counter}{export_mask_path.suffix}' + mask_exists = export_mask_path.exists() + counter += 1 + + # write the tag file first + if filter_hashtag: + tags = [tag for tag in image_entry.tags if tag[0] != '#' or + (separate_newline != HashNewlineHandling.IGNORE and + tag == '#newline')] + else: + tags = image_entry.tags.copy() + + if len(tags) == 0: + tags = [''] + + if separate_newline: + tag_groups = [] + temp_list = [] + for tag in tags: + if tag == '#newline': + if temp_list: + tag_groups.append(temp_list) + temp_list = [] + else: + temp_list.append(tag) + if temp_list: + tag_groups.append(temp_list) + else: + tag_groups = [tags] + + this_caption_algorithm = caption_algorithm + if caption_algorithm == CaptionStrategy.PREFIX_ENUMERATION: + prefix = tag_groups[0].pop(0) + ' ' + if len(tag_groups[0]) == 0: + tag_groups.pop(0) + this_caption_algorithm = CaptionStrategy.ENUMERATION + else: + prefix = '' + + all_tags = [] + for tags in tag_groups: + match this_caption_algorithm: + case CaptionStrategy.TAG_LIST: + tag_string = tag_separator.join(tags) + case CaptionStrategy.FIRST: + tag_string = tags[0] + case CaptionStrategy.LAST: + tag_string = tags[-1] + case CaptionStrategy.ENUMERATION: + if len(tags) == 1: + tag_string = tags[0] + elif len(tags) == 2: + tag_string = ' and '.join(tags) + else: + tag_string = ', '.join(tags[:-1]) + ', and ' + tags[-1] + tag_string = prefix + tag_string + if tag_string != '': + all_tags.append(tag_string) + + try: + if separate_newline != HashNewlineHandling.MULTIFILE: + multifile_count = None + export_path.with_suffix('.txt').write_text( + '\n'.join(all_tags), encoding='utf-8', errors='replace') + else: + multifile_count = len(all_tags) or None + for index, this_tags in enumerate(all_tags): + suffix = '.txt' if index == 0 else f'.{index}.txt' + export_path.with_suffix(suffix).write_text( + this_tags, encoding='utf-8', errors='replace') + + except OSError: + error_message_box = QMessageBox() + error_message_box.setWindowTitle('Error') + error_message_box.setIcon(QMessageBox.Icon.Critical) + error_message_box.setText(f'Failed to save tags for {image_entry.path}.') + error_message_box.exec() + + if (image_exists or mask_exists) and only_missing: + # tags were refreshed, export_path was changed when we should + # rename and not overwrite, so we can skip the image writing + continue + + # then handle the image + image_file = Image.open(image_entry.path) + export_can_alpha = export_format != ExportFormat.JPG + export_mask = masking_strategy != MaskingStrategy.IGNORE + # Preserve alpha if present: + if image_file.mode in ('RGBA', 'LA', 'PA') and export_mask: # Check for alpha channels + image_file = image_file.convert('RGBA') + else: + image_file = image_file.convert('RGB') # Otherwise, convert to RGB + + # 1. pass: add includes + for marking in image_entry.markings: + if marking.type == ImageMarking.INCLUDE and export_mask: + if image_file.mode == 'RGB': + image_file = image_file.convert('RGBA') + # completely transparent + alpha = Image.new('L', image_file.size, 0) + else: + alpha = image_file.getchannel('A') + if not quantize_alpha: + alpha.paste(255, marking.rect.adjusted(0,0,1,1).getCoords()) + image_file.putalpha(alpha) + + # 2. pass: remove excludes + for marking in image_entry.markings: + if marking.type == ImageMarking.EXCLUDE and export_mask: + if image_file.mode == 'RGB': + image_file = image_file.convert('RGBA') + # completely opaque + alpha = Image.new('L', image_file.size, 255) + else: + alpha = image_file.getchannel('A') + if not quantize_alpha: + alpha.paste(0, marking.rect.adjusted(0,0,1,1).getCoords()) + image_file.putalpha(alpha) + + if image_entry.crop is None: + grid = Grid(QRect(0, 0, *image_file.size)) + else: + grid = Grid(image_entry.crop) + visible = grid.visible + cropped_image = image_file.crop(visible.adjusted(0,0,1,1).getCoords()) + if not grid.is_visible_equal_screen_size(): + # resize with the best method available + resized_image = cropped_image.resize(grid.target.toTuple(), Image.LANCZOS) + # followed by a slight sharpening as it should be done + sharpened_image = resized_image.filter( + ImageFilter.UnsharpMask(radius = 0.5, percent = 50, threshold = 0)) + else: + sharpened_image = cropped_image + + # crop to the desired size + current_width, current_height = sharpened_image.size + crop_width = floor((current_width - image_entry.target_dimension.width()) / 2) + crop_height = floor((current_height - image_entry.target_dimension.height()) / 2) + cropped_image = sharpened_image.crop((crop_width, crop_height, + crop_width + image_entry.target_dimension.width(), + crop_height + image_entry.target_dimension.height())) + + if export_mask: + if cropped_image.mode == 'RGB': + cropped_image = cropped_image.convert('RGBA') + alpha = cropped_image.getchannel('A') + if quantize_alpha: + for marking in image_entry.markings: + if marking.type == ImageMarking.INCLUDE: + rect = QRect(grid.map(marking.rect.topLeft(), ceil), + grid.map(marking.rect.adjusted(0,0,1,1).bottomRight(), floor)) + alpha.paste(255, rect.getCoords()) + for marking in image_entry.markings: + if marking.type == ImageMarking.EXCLUDE: + rect = QRect(grid.map(marking.rect.topLeft(), floor), + grid.map(marking.rect.adjusted(0,0,1,1).bottomRight(), ceil)) + alpha.paste(0, rect.getCoords()) + + replacement = None + if masked_content in [MaskedContent.BLUR, MaskedContent.BLUR_NOISE]: + replacement = cropped_image.filter(ImageFilter.GaussianBlur(10)) + elif masked_content in [MaskedContent.GREY, MaskedContent.GREY_NOISE]: + # 126 is an 18% gray, i.e. the neutral gray, for sRGB. + # Anyway, it's masked, so there's no need to go into detail + # about different color spaces. + replacement = Image.new('RGB', cropped_image.size, (126, 126, 126)) + elif masked_content == MaskedContent.BLACK: + replacement = Image.new('RGB', cropped_image.size, (0, 0, 0)) + elif masked_content == MaskedContent.WHITE: + replacement = Image.new('RGB', cropped_image.size, (255, 255, 255)) + + if masked_content in [MaskedContent.BLUR_NOISE, MaskedContent.GREY_NOISE]: + np_image = np.array(replacement) + # Add random noise with a minimal blur + noise = np.random.normal(0, 30, np_image.shape).astype(np.int8) + noisy_image = np_image + noise + noisy_image = np.clip(noisy_image, 0, 255).astype(np.uint8) + replacement = Image.fromarray(noisy_image).filter(ImageFilter.GaussianBlur(1)) + + if replacement: + cropped_image = Image.composite(cropped_image, replacement, alpha) + + cropped_image.putalpha(alpha) + + if not export_can_alpha or masking_strategy in [MaskingStrategy.REPLACE, + MaskingStrategy.MASK_FILE]: + # remove alpha + export_image = cropped_image.convert('RGB') + else: + export_image = cropped_image + + lossless = quality > 99 + + if color_space == "feed through (don't touch)": + export_image.save(export_path, format=ExportFormatDict[export_format], + quality=quality, lossless=lossless, + icc_profile=export_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(export_image, source_profile, target_profile) + if save_profile: + final_image.save(export_path, format=ExportFormatDict[export_format], + quality=quality, lossless=lossless, + icc_profile=target_profile.tobytes()) + else: + final_image.save(export_path, format=ExportFormatDict[export_format], + quality=quality, lossless=lossless, + icc_profile=None) + if masking_strategy == MaskingStrategy.MASK_FILE: + alpha_channel = cropped_image.getchannel('A') + alpha_channel.save(export_mask_path, format=ExportFormatDict[export_format], + quality=quality, lossless=lossless, + icc_profile=None) + + if multifile_count is not None: + for index in range(1, multifile_count): + suffix = f'.{index}{export_path.suffix}' + shutil.copy(export_path, export_path.with_suffix(suffix)) + if masking_strategy == MaskingStrategy.MASK_FILE: + shutil.copy(export_mask_path, export_mask_path.with_suffix(suffix)) + self.close() + + def get_image_list(self): + image_list_view = self.image_list.list_view + if settings.value('export_filter') == ExportFilter.FILTERED: + 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 settings.value('export_filter') == ExportFilter.SELECTED: + 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/find_and_replace_dialog.py b/taggui/dialogs/find_and_replace_dialog.py index f61f9e0e..3d8dc0f5 100644 --- a/taggui/dialogs/find_and_replace_dialog.py +++ b/taggui/dialogs/find_and_replace_dialog.py @@ -5,7 +5,6 @@ QVBoxLayout) from models.image_list_model import ImageListModel, Scope -from utils.settings import get_settings from utils.settings_widgets import (SettingsBigCheckBox, SettingsComboBox, SettingsLineEdit) from utils.utils import pluralize @@ -15,7 +14,6 @@ class FindAndReplaceDialog(QDialog): def __init__(self, parent, image_list_model: ImageListModel): super().__init__(parent) self.image_list_model = image_list_model - self.settings = get_settings() self.setWindowTitle('Find and Replace') layout = QVBoxLayout(self) layout.setContentsMargins(20, 20, 20, 20) diff --git a/taggui/dialogs/settings_dialog.py b/taggui/dialogs/settings_dialog.py index 8fb2956b..88a42643 100644 --- a/taggui/dialogs/settings_dialog.py +++ b/taggui/dialogs/settings_dialog.py @@ -2,7 +2,7 @@ from PySide6.QtWidgets import (QDialog, QFileDialog, QGridLayout, QLabel, QLineEdit, QPushButton, QVBoxLayout) -from utils.settings import DEFAULT_SETTINGS, get_settings +from utils.settings import DEFAULT_SETTINGS, settings from utils.settings_widgets import (SettingsBigCheckBox, SettingsLineEdit, SettingsSpinBox) @@ -10,7 +10,6 @@ class SettingsDialog(QDialog): def __init__(self, parent): super().__init__(parent) - self.settings = get_settings() self.setWindowTitle('Settings') layout = QVBoxLayout(self) layout.setContentsMargins(20, 20, 20, 20) @@ -31,25 +30,29 @@ def __init__(self, parent): 5, 0, Qt.AlignmentFlag.AlignRight) grid_layout.addWidget(QLabel('Auto-captioning models directory'), 6, 0, Qt.AlignmentFlag.AlignRight) + grid_layout.addWidget(QLabel('Auto-marking models directory'), 8, 0, + 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) + file_types_line_edit = SettingsLineEdit( + key='image_list_file_formats') + file_types_line_edit.setMinimumWidth(400) + file_types_line_edit.textChanged.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() - tag_separator = self.settings.value( + tag_separator = settings.value( 'tag_separator', defaultValue=DEFAULT_SETTINGS['tag_separator'], type=str) if tag_separator == '\n': @@ -60,13 +63,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( @@ -75,11 +76,15 @@ def __init__(self, parent): models_directory_button.setFixedWidth( int(models_directory_button.sizeHint().width() * 1.3)) models_directory_button.clicked.connect(self.set_models_directory_path) - file_types_line_edit = SettingsLineEdit( - key='image_list_file_formats', - default=DEFAULT_SETTINGS['image_list_file_formats']) - file_types_line_edit.setMinimumWidth(400) - file_types_line_edit.textChanged.connect(self.show_restart_warning) + self.marking_models_directory_line_edit = SettingsLineEdit( + key='marking_models_directory_path') + self.marking_models_directory_line_edit.setMinimumWidth(400) + self.marking_models_directory_line_edit.setClearButtonEnabled(True) + marking_models_directory_button = QPushButton('Select Directory...') + marking_models_directory_button.setFixedWidth( + int(marking_models_directory_button.sizeHint().width() * 1.3)) + marking_models_directory_button.clicked.connect( + self.set_marking_models_directory_path) grid_layout.addWidget(font_size_spin_box, 0, 1, Qt.AlignmentFlag.AlignLeft) @@ -97,6 +102,10 @@ def __init__(self, parent): Qt.AlignmentFlag.AlignLeft) grid_layout.addWidget(models_directory_button, 7, 1, Qt.AlignmentFlag.AlignLeft) + grid_layout.addWidget(self.marking_models_directory_line_edit, 8, 1, + Qt.AlignmentFlag.AlignLeft) + grid_layout.addWidget(marking_models_directory_button, 9, 1, + Qt.AlignmentFlag.AlignLeft) layout.addLayout(grid_layout) # Prevent the grid layout from moving to the center when the warning @@ -133,18 +142,18 @@ def handle_tag_separator_change(self, tag_separator: str): self.disable_insert_space_after_tag_separator_check_box() else: self.insert_space_after_tag_separator_check_box.setEnabled(True) - self.settings.setValue('tag_separator', tag_separator) + settings.setValue('tag_separator', tag_separator) self.show_restart_warning() @Slot() def set_models_directory_path(self): - models_directory_path = self.settings.value( + models_directory_path = settings.value( 'models_directory_path', defaultValue=DEFAULT_SETTINGS['models_directory_path'], type=str) if models_directory_path: initial_directory_path = models_directory_path - elif self.settings.contains('directory_path'): - initial_directory_path = self.settings.value('directory_path') + elif settings.contains('directory_path'): + initial_directory_path = settings.value('directory_path', type=str) else: initial_directory_path = '' models_directory_path = QFileDialog.getExistingDirectory( @@ -153,3 +162,21 @@ def set_models_directory_path(self): dir=initial_directory_path) if models_directory_path: self.models_directory_line_edit.setText(models_directory_path) + + @Slot() + def set_marking_models_directory_path(self): + marking_models_directory_path = settings.value( + 'marking_models_directory_path', + defaultValue=DEFAULT_SETTINGS['marking_models_directory_path'], type=str) + if marking_models_directory_path: + initial_directory_path = marking_models_directory_path + elif settings.contains('directory_path'): + initial_directory_path = settings.value('directory_path', type=str) + else: + initial_directory_path = '' + marking_models_directory_path = QFileDialog.getExistingDirectory( + parent=self, caption='Select directory containing auto-marking ' + 'models (YOLO models)', + dir=initial_directory_path) + if marking_models_directory_path: + self.marking_models_directory_line_edit.setText(marking_models_directory_path) diff --git a/taggui/models/image_list_model.py b/taggui/models/image_list_model.py index 09e9127d..9f1791ff 100644 --- a/taggui/models/image_list_model.py +++ b/taggui/models/image_list_model.py @@ -4,18 +4,21 @@ from collections import Counter, deque from dataclasses import dataclass from enum import Enum +from math import floor, ceil from pathlib import Path +import json import exifread import imagesize -from PySide6.QtCore import (QAbstractListModel, QModelIndex, QSize, Qt, Signal, - Slot) +from PySide6.QtCore import (QAbstractListModel, QModelIndex, QPoint, QRect, + QSize, Qt, Signal, Slot) from PySide6.QtGui import QIcon, QImageReader, QPixmap from PySide6.QtWidgets import QMessageBox -from utils.image import Image -from utils.settings import DEFAULT_SETTINGS, get_settings +from utils.image import Image, ImageMarking, Marking +from utils.settings import DEFAULT_SETTINGS, settings from utils.utils import get_confirmation_dialog_reply, pluralize +import utils.target_dimension as target_dimension UNDO_STACK_SIZE = 32 @@ -37,7 +40,7 @@ def get_file_paths(directory_path: Path) -> set[Path]: @dataclass class HistoryItem: action_name: str - tags: list[list[str]] + tags: list[dict[str, list[str] | QRect | None | list[Marking]]] should_ask_for_confirmation: bool @@ -82,6 +85,15 @@ def data(self, index, role=None) -> Image | str | QIcon | QSize: image_reader = QImageReader(str(image.path)) # Rotate the image based on the orientation tag. image_reader.setAutoTransform(True) + if image.crop: + crop = image.crop + else: + crop = QRect(QPoint(0, 0), image_reader.size()) + if crop.height() > crop.width()*3: + # keep it reasonable, higher than 3x the width doesn't make sense + crop.setTop((crop.height() - crop.width()*3)//2) # center crop + crop.setHeight(crop.width()*3) + image_reader.setClipRect(crop) pixmap = QPixmap.fromImageReader(image_reader).scaledToWidth( self.image_list_image_width, Qt.TransformationMode.SmoothTransformation) @@ -91,22 +103,32 @@ def data(self, index, role=None) -> Image | str | QIcon | QSize: if role == Qt.ItemDataRole.SizeHintRole: if image.thumbnail: return image.thumbnail.availableSizes()[0] - dimensions = image.dimensions + dimensions = image.crop.size().toTuple() if image.crop else image.dimensions if not dimensions: return QSize(self.image_list_image_width, self.image_list_image_width) width, height = dimensions # Scale the dimensions to the image width. return QSize(self.image_list_image_width, - int(self.image_list_image_width * height / width)) + int(self.image_list_image_width * min(height / width, 3))) + if role == Qt.ItemDataRole.ToolTipRole: + path = image.path.relative_to(settings.value('directory_path', type=str)) + dimensions = f'{image.dimensions[0]}:{image.dimensions[1]}' + if not image.target_dimension: + if image.crop: + image.target_dimension = target_dimension.get(image.crop.size()) + else: + image.target_dimension = target_dimension.get(QSize(*image.dimensions)) + target = f'{image.target_dimension.width()}:{image.target_dimension.height()}' + return f'{path}\n{dimensions} 🠮 {target}' def load_directory(self, directory_path: Path): self.images.clear() self.undo_stack.clear() self.redo_stack.clear() self.update_undo_and_redo_actions_requested.emit() + error_messages: list[str] = [] file_paths = get_file_paths(directory_path) - settings = get_settings() image_suffixes_string = settings.value( 'image_list_file_formats', defaultValue=DEFAULT_SETTINGS['image_list_file_formats'], type=str) @@ -122,6 +144,8 @@ def load_directory(self, directory_path: Path): # strings. text_file_path_strings = {str(path) for path in file_paths if path.suffix == '.txt'} + json_file_path_strings = {str(path) for path in file_paths + if path.suffix == '.json'} for image_path in image_paths: try: dimensions = imagesize.get(image_path) @@ -139,11 +163,11 @@ def load_directory(self, directory_path: Path): for value in (5, 6, 7, 8)): dimensions = (dimensions[1], dimensions[0]) except Exception as exception: - print(f'Failed to get Exif tags for {image_path}: ' - f'{exception}', file=sys.stderr) + error_messages.append(f'Failed to get Exif tags for ' + f'{image_path}: {exception}') except (ValueError, OSError) as exception: - print(f'Failed to get dimensions for {image_path}: ' - f'{exception}', file=sys.stderr) + error_messages.append(f'Failed to get dimensions for ' + f'{image_path}: {exception}') dimensions = None tags = [] text_file_path = image_path.with_suffix('.txt') @@ -157,14 +181,58 @@ def load_directory(self, directory_path: Path): tags = [tag.strip() for tag in tags] tags = [tag for tag in tags if tag] image = Image(image_path, dimensions, tags) + json_file_path = image_path.with_suffix('.json') + if (str(json_file_path) in json_file_path_strings and + json_file_path.stat().st_size > 0): + with json_file_path.open(encoding='UTF-8') as source: + try: + meta = json.load(source) + except json.JSONDecodeError as e: + error_messages.append(f'Invalid JSON in ' + f'"{json_file_path}": {e.msg}') + break + except UnicodeDecodeError as e: + error_messages.append(f'Invalid Unicode in JSON in ' + f'"{json_file_path}": {e.reason}') + break + + if meta.get('version') == 1: + crop = meta.get('crop') + if crop and type(crop) is list and len(crop) == 4: + image.crop = QRect(*crop) + rating = meta.get('rating') + if rating: + image.rating = rating + markings = meta.get('markings') + if markings and type(markings) is list: + for marking in markings: + marking = Marking(label=marking.get('label'), + type=ImageMarking[marking.get('type')], + rect=QRect(*marking.get('rect')), + confidence=marking.get('confidence', 1.0)) + image.markings.append(marking) + else: + error_messages.append(f'Invalid version ' + f'"{meta.get('version')}" in ' + f'"{json_file_path}"') self.images.append(image) self.images.sort(key=lambda image_: image_.path) self.modelReset.emit() + if len(error_messages) > 0: + print('\n'.join(error_messages), file=sys.stderr) + error_message_box = QMessageBox() + error_message_box.setWindowTitle('Directory reading error') + error_message_box.setIcon(QMessageBox.Icon.Warning) + error_message_box.setText('\n'.join(error_messages)) + error_message_box.exec() def add_to_undo_stack(self, action_name: str, should_ask_for_confirmation: bool): """Add the current state of the image tags to the undo stack.""" - tags = [image.tags.copy() for image in self.images] + tags = [{'tags': image.tags.copy(), + 'rating': image.rating, + 'crop': QRect(image.crop) if image.crop is not None else None, + 'markings': image.markings.copy()} for image in self.images] self.undo_stack.append(HistoryItem(action_name, tags, should_ask_for_confirmation)) self.redo_stack.clear() @@ -182,6 +250,26 @@ def write_image_tags_to_disk(self, image: Image): error_message_box.setText(f'Failed to save tags for {image.path}.') error_message_box.exec() + def write_meta_to_disk(self, image: Image): + does_exist = image.path.with_suffix('.json').exists() + meta: dict[str, any] = {'version': 1, 'rating': image.rating} + if image.crop is not None: + meta['crop'] = image.crop.getRect() + meta['markings'] = [{'label': marking.label, + 'type': marking.type.name, + 'confidence': marking.confidence, + 'rect': marking.rect.getRect()} for marking in image.markings] + if does_exist or len(meta.keys()) > 1: + try: + with image.path.with_suffix('.json').open('w', encoding='UTF-8') as meta_file: + json.dump(meta, meta_file) + except OSError: + error_message_box = QMessageBox() + error_message_box.setWindowTitle('Error') + error_message_box.setIcon(QMessageBox.Icon.Critical) + error_message_box.setText(f'Failed to save JSON for {image.path}.') + error_message_box.exec() + def restore_history_tags(self, is_undo: bool): if is_undo: source_stack = self.undo_stack @@ -202,18 +290,28 @@ def restore_history_tags(self, is_undo: bool): if reply != QMessageBox.StandardButton.Yes: return source_stack.pop() - tags = [image.tags for image in self.images] + tags = [{'tags': image.tags.copy(), + 'rating': image.rating, + 'crop': QRect(image.crop) if image.crop is not None else None, + 'markings': image.markings.copy()} for image in self.images] destination_stack.append(HistoryItem( history_item.action_name, tags, history_item.should_ask_for_confirmation)) changed_image_indices = [] for image_index, (image, history_image_tags) in enumerate( zip(self.images, history_item.tags)): - if image.tags == history_image_tags: + if (image.tags == history_image_tags['tags'] and + image.rating == history_image_tags['rating'] and + image.crop == history_image_tags['crop'] and + image.markings == history_image_tags['markings']): continue changed_image_indices.append(image_index) - image.tags = history_image_tags + image.tags = history_image_tags['tags'] + image.rating = history_image_tags['rating'] + image.crop = history_image_tags['crop'] + image.markings = history_image_tags['markings'] self.write_image_tags_to_disk(image) + self.write_meta_to_disk(image) if changed_image_indices: self.dataChanged.emit(self.index(changed_image_indices[0]), self.index(changed_image_indices[-1])) @@ -384,6 +482,36 @@ def shuffle_tags(self, do_not_reorder_first_tag: bool): self.dataChanged.emit(self.index(changed_image_indices[0]), self.index(changed_image_indices[-1])) + def sort_sentences_down(self, separate_newline: bool): + """Sort the tags so that the sentences are on the bottom.""" + self.add_to_undo_stack(action_name='Sort Sentence Tags', + should_ask_for_confirmation=True) + changed_image_indices = [] + for image_index, image in enumerate(self.images): + changed_image_indices.append(image_index) + sentence_tags = [] + non_sentence_tags = [] + for tag in image.tags: + if separate_newline and tag == '#newline': + continue + if tag.endswith('.'): + sentence_tags.append(tag) + else: + non_sentence_tags.append(tag) + if separate_newline: + if len(sentence_tags) > 0: + non_sentence_tags.append(sentence_tags.pop()) + for tag in sentence_tags: + non_sentence_tags.append('#newline') + non_sentence_tags.append(tag) + else: + non_sentence_tags.extend(sentence_tags) + image.tags = non_sentence_tags + self.write_image_tags_to_disk(image) + if changed_image_indices: + self.dataChanged.emit(self.index(changed_image_indices[0]), + self.index(changed_image_indices[-1])) + def move_tags_to_front(self, tags_to_move: list[str]): """ Move one or more tags to the front of the tags list for each image. @@ -539,3 +667,21 @@ def delete_tags(self, tags: list[str], if changed_image_indices: self.dataChanged.emit(self.index(changed_image_indices[0]), self.index(changed_image_indices[-1])) + + def add_image_markings(self, image_index: QModelIndex, markings: list[dict]): + image: Image = self.data(image_index, Qt.ItemDataRole.UserRole) + for marking in markings: + marking_type = { + 'hint': ImageMarking.HINT, + 'include': ImageMarking.INCLUDE, + 'exclude': ImageMarking.EXCLUDE}[marking['type']] + box = marking['box'] + top_left = QPoint(floor(box[0]), floor(box[1])) + bottom_right = QPoint(ceil(box[2]), ceil(box[3])) + image.markings.append(Marking(label=marking['label'], + type=marking_type, + rect=QRect(top_left, bottom_right), + confidence=marking['confidence'])) + if len(markings) > 0: + self.dataChanged.emit(image_index, image_index) + self.write_meta_to_disk(image) diff --git a/taggui/models/proxy_image_list_model.py b/taggui/models/proxy_image_list_model.py index b730c059..da593b8d 100644 --- a/taggui/models/proxy_image_list_model.py +++ b/taggui/models/proxy_image_list_model.py @@ -1,14 +1,29 @@ import operator +import re from fnmatch import fnmatchcase -from PySide6.QtCore import QModelIndex, QSortFilterProxyModel, Qt +from PySide6.QtCore import (QModelIndex, QSortFilterProxyModel, Qt, QRect, + QSize, Signal) from transformers import PreTrainedTokenizerBase from models.image_list_model import ImageListModel from utils.image import Image +import utils.target_dimension as target_dimension + +comparison_operators = { + '=': operator.eq, + '==': operator.eq, + '!=': operator.ne, + '<': operator.lt, + '>': operator.gt, + '<=': operator.le, + '>=': operator.ge +} class ProxyImageListModel(QSortFilterProxyModel): + filter_changed = Signal() + def __init__(self, image_list_model: ImageListModel, tokenizer: PreTrainedTokenizerBase, tag_separator: str): super().__init__() @@ -17,8 +32,15 @@ def __init__(self, image_list_model: ImageListModel, self.tag_separator = tag_separator self.filter: list | None = None + def set_filter(self, new_filter: list | None): + self.filter = new_filter + self.invalidateFilter() + self.filter_changed.emit() + def does_image_match_filter(self, image: Image, - filter_: list | str) -> bool: + filter_: list | str | None) -> bool: + if filter_ is None: + return True if isinstance(filter_, str): return (fnmatchcase(self.tag_separator.join(image.tags), f'*{filter_}*') @@ -33,25 +55,58 @@ def does_image_match_filter(self, image: Image, if filter_[0] == 'caption': caption = self.tag_separator.join(image.tags) return fnmatchcase(caption, f'*{filter_[1]}*') + if filter_[0] == 'marking': + last_colon_index = filter_[1].rfind(':') + if last_colon_index < 0: + return any(fnmatchcase(marking.label, filter_[1]) + for marking in image.markings) + else: + label = filter_[1][:last_colon_index] + confidence = filter_[1][last_colon_index + 1:] + pattern =r'^(<=|>=|==|<|>|=)\s*(0?[.,][0-9]+)' + match = re.match(pattern, confidence) + if not match or len(match.group(2)) == 0: + return False + comparison_operator = comparison_operators[match.group(1)] + confidence_target = float(match.group(2).replace(',', '.')) + return any((fnmatchcase(marking.label, label) and + comparison_operator(marking.confidence, + confidence_target)) + for marking in image.markings) + if filter_[0] == 'crops': + crop = image.crop if image.crop is not None else QRect(0, 0, *image.dimensions) + return any(fnmatchcase(marking.label, filter_[1]) and + marking.rect.intersects(crop) and not crop.contains(marking.rect) + for marking in image.markings) + if filter_[0] == 'visible': + crop = image.crop if image.crop is not None else QRect(0, 0, *image.dimensions) + return any(fnmatchcase(marking.label, filter_[1]) and + marking.rect.intersects(crop) + for marking in image.markings) if filter_[0] == 'name': 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_dimension is None: + image.target_dimension = target_dimension.get(QSize(*image.dimensions)) + return (len(dimension) == 2 + and dimension[0] == str(image.target_dimension.width()) + and dimension[1] == str(image.target_dimension.height())) if filter_[1] == 'AND': return (self.does_image_match_filter(image, filter_[0]) and self.does_image_match_filter(image, filter_[2:])) if filter_[1] == 'OR': return (self.does_image_match_filter(image, filter_[0]) or self.does_image_match_filter(image, filter_[2:])) - comparison_operators = { - '=': operator.eq, - '==': operator.eq, - '!=': operator.ne, - '<': operator.lt, - '>': operator.gt, - '<=': operator.le, - '>=': operator.ge - } comparison_operator = comparison_operators[filter_[1]] number_to_compare = None if filter_[0] == 'tags': @@ -63,6 +118,14 @@ 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] == 'stars': + number_to_compare = image.rating * 5.0 + elif filter_[0] == 'width': + number_to_compare = image.dimensions[0] + elif filter_[0] == 'height': + number_to_compare = image.dimensions[1] + elif filter_[0] == 'area': + number_to_compare = image.dimensions[0] * image.dimensions[1] return comparison_operator(number_to_compare, int(filter_[2])) def filterAcceptsRow(self, source_row: int, @@ -78,3 +141,7 @@ def filterAcceptsRow(self, source_row: int, def is_image_in_filtered_images(self, image: Image) -> bool: return (self.filter is None or self.does_image_match_filter(image, self.filter)) + + def get_list(self) -> list[Image]: + return [self.data(self.index(row, 0, QModelIndex()), Qt.UserRole) + for row in range(self.rowCount())] diff --git a/taggui/models/proxy_tag_counter_model.py b/taggui/models/proxy_tag_counter_model.py index 21becf30..d473f3d4 100644 --- a/taggui/models/proxy_tag_counter_model.py +++ b/taggui/models/proxy_tag_counter_model.py @@ -21,8 +21,19 @@ def lessThan(self, left: QModelIndex, right: QModelIndex) -> bool: left.row()] right_tag, right_count = self.tag_counter_model.most_common_tags[ right.row()] + if self.tag_counter_model.most_common_tags_filtered is None: + left_cnt_f = 0 + right_cnt_f = 0 + else: + left_cnt_f = self.tag_counter_model.most_common_tags_filtered[left_tag] + right_cnt_f = self.tag_counter_model.most_common_tags_filtered[right_tag] + if self.sort_by == AllTagsSortBy.FREQUENCY: - return left_count < right_count + return left_count < right_count or (left_count == right_count and + left_cnt_f < right_cnt_f) + elif self.sort_by == AllTagsSortBy.FREQUENCY_FILTERED: + return left_cnt_f < right_cnt_f or (left_cnt_f == right_cnt_f and + left_count < right_count) elif self.sort_by == AllTagsSortBy.NAME: return left_tag < right_tag elif self.sort_by == AllTagsSortBy.LENGTH: diff --git a/taggui/models/tag_counter_model.py b/taggui/models/tag_counter_model.py index 0d0b5c1e..95cda5c8 100644 --- a/taggui/models/tag_counter_model.py +++ b/taggui/models/tag_counter_model.py @@ -14,6 +14,7 @@ def __init__(self): super().__init__() self.tag_counter = Counter() self.most_common_tags = [] + self.most_common_tags_filtered = None self.all_tags_list = None def rowCount(self, parent=None) -> int: @@ -24,7 +25,10 @@ def data(self, index, role=None) -> tuple[str, int] | str: if role == Qt.ItemDataRole.UserRole: return tag, count if role == Qt.ItemDataRole.DisplayRole: - return f'{tag} ({count})' + if self.most_common_tags_filtered is None: + return f'{tag} ({count})' + else: + return f'{tag} ({self.most_common_tags_filtered[tag]}/{count})' if role == Qt.ItemDataRole.EditRole: return tag @@ -69,7 +73,18 @@ def setData(self, index, value: str, @Slot() def count_tags(self, images: list[Image]): self.tag_counter.clear() + self.most_common_tags_filtered = None for image in images: self.tag_counter.update(image.tags) self.most_common_tags = self.tag_counter.most_common() self.modelReset.emit() + + @Slot() + def count_tags_filtered(self, images: list[Image] | None): + if images is None: + self.most_common_tags_filtered = None + else: + self.most_common_tags_filtered = Counter() + for image in images: + self.most_common_tags_filtered.update(image.tags) + self.modelReset.emit() diff --git a/taggui/run_gui.py b/taggui/run_gui.py index 46b74a8c..5f19d722 100644 --- a/taggui/run_gui.py +++ b/taggui/run_gui.py @@ -8,7 +8,7 @@ from PySide6.QtGui import QImageReader from PySide6.QtWidgets import QApplication, QMessageBox -from utils.settings import get_settings +from utils.settings import settings from widgets.main_window import MainWindow @@ -49,7 +49,6 @@ def run_gui(): try: run_gui() except Exception as exception: - settings = get_settings() settings.clear() error_message_box = QMessageBox() error_message_box.setWindowTitle('Error') diff --git a/taggui/utils/ModelThread.py b/taggui/utils/ModelThread.py new file mode 100644 index 00000000..4b1e3a3f --- /dev/null +++ b/taggui/utils/ModelThread.py @@ -0,0 +1,133 @@ +from abc import abstractmethod +from datetime import datetime +from time import perf_counter + +import numpy as np +from PIL import UnidentifiedImageError +from transformers import BatchFeature +from PySide6.QtCore import QModelIndex, QThread, Qt, Signal + +from utils.image import Image +from models.image_list_model import ImageListModel + +def format_duration(seconds: float) -> str: + seconds_per_minute = 60 + seconds_per_hour = 60 * seconds_per_minute + seconds_per_day = 24 * seconds_per_hour + if seconds < seconds_per_minute: + return f'{seconds:.1f} seconds' + if seconds < seconds_per_hour: + minutes = seconds / seconds_per_minute + return f'{minutes:.1f} minutes' + if seconds < seconds_per_day: + hours = seconds / seconds_per_hour + return f'{hours:.1f} hours' + days = seconds / seconds_per_day + return f'{days:.1f} days' + + +class ModelThread(QThread): + """Base class for all model running threads""" + text_outputted = Signal(str) + clear_console_text_edit_requested = Signal() + progress_bar_update_requested = Signal(int) + + def __init__(self, parent, image_list_model: ImageListModel, + selected_image_indices: list[QModelIndex]): + super().__init__(parent) + self.image_list_model = image_list_model + self.selected_image_indices = selected_image_indices + self.is_error = False + self.error_message = '' + self.is_canceled = False + self.device = 'default' + self.text = { + 'Generating': 'Generating', + 'generating': 'generating' + } + + def run_generating(self): + self.load_model() + if self.is_error: + self.clear_console_text_edit_requested.emit() + print(self.error_message) + return + if self.is_canceled: + print(f'Canceled {self.text['generating']}.') + return + self.clear_console_text_edit_requested.emit() + selected_image_count = len(self.selected_image_indices) + are_multiple_images_selected = selected_image_count > 1 + generating_start_datetime = datetime.now() + generating_message = self.get_generating_message( + are_multiple_images_selected, generating_start_datetime) + print(generating_message) + for i, image_index in enumerate(self.selected_image_indices): + start_time = perf_counter() + if self.is_canceled: + print(f'Canceled {self.text['generating']}.') + return + image: Image = self.image_list_model.data(image_index, + Qt.ItemDataRole.UserRole) + try: + image_prompt, model_inputs = self.get_model_inputs(image) + except UnidentifiedImageError: + print(f'Skipping {image.path.name} because its file format is ' + 'not supported or it is a corrupted image.') + continue + console_output_caption = self.generate_output(image_index, image, + image_prompt, model_inputs) + if are_multiple_images_selected: + self.progress_bar_update_requested.emit(i + 1) + if i == 0 and not are_multiple_images_selected: + self.clear_console_text_edit_requested.emit() + print(f'{image.path.name} ({perf_counter() - start_time:.1f} s):\n' + f'{console_output_caption}') + if are_multiple_images_selected: + generating_end_datetime = datetime.now() + total_generating_duration = ((generating_end_datetime + - generating_start_datetime) + .total_seconds()) + average_generating_duration = (total_generating_duration / + selected_image_count) + print(f'Finished {self.text['generating']} {selected_image_count} images in ' + f'{format_duration(total_generating_duration)} ' + f'({average_generating_duration:.1f} s/image) at ' + f'{generating_end_datetime.strftime("%Y-%m-%d %H:%M:%S")}.') + + @abstractmethod + def load_model(self): + """Load the model for the generating task.""" + pass + + def get_generating_message(self, are_multiple_images_selected: bool, + generating_start_datetime: datetime) -> str: + if are_multiple_images_selected: + generating_start_datetime_string = ( + generating_start_datetime.strftime('%Y-%m-%d %H:%M:%S')) + return (f'{self.text['Generating']}... (device: {self.device}, ' + f'start time: {generating_start_datetime_string})') + return f'{self.text['Generating']}... (device: {self.device})' + + @abstractmethod + def get_model_inputs(self, image: Image) -> tuple[ + str | None, BatchFeature | dict | np.ndarray]: + pass + + @abstractmethod + def generate_output(self, image_index, + image: Image, + image_prompt: str | None, + model_inputs: BatchFeature | dict | np.ndarray) -> str: + pass + + def run(self): + try: + self.run_generating() + except Exception as exception: + self.is_error = True + # Show the error message in the console text edit. + raise exception + + def write(self, text: str): + self.text_outputted.emit(text) diff --git a/taggui/utils/big_widgets.py b/taggui/utils/big_widgets.py index 098026bf..b0c7a5ab 100644 --- a/taggui/utils/big_widgets.py +++ b/taggui/utils/big_widgets.py @@ -1,6 +1,6 @@ from PySide6.QtWidgets import QCheckBox, QPushButton -from utils.settings import DEFAULT_SETTINGS, get_settings +from utils.settings import DEFAULT_SETTINGS, settings class BigPushButton(QPushButton): @@ -20,7 +20,6 @@ def __init__(self, text: str): class BigCheckBox(QCheckBox): def __init__(self, text: str | None = None): super().__init__(text) - settings = get_settings() font_size = settings.value( 'font_size', defaultValue=DEFAULT_SETTINGS['font_size'], type=int) new_size = font_size * 1.5 diff --git a/taggui/utils/enums.py b/taggui/utils/enums.py index 21cb60c8..163eb2f3 100644 --- a/taggui/utils/enums.py +++ b/taggui/utils/enums.py @@ -4,6 +4,7 @@ # `StrEnum` is a Python 3.11 feature that can be used here. class AllTagsSortBy(str, Enum): FREQUENCY = 'Frequency' + FREQUENCY_FILTERED = 'Frequency (filtered)' NAME = 'Name' LENGTH = 'Length' @@ -24,3 +25,79 @@ class CaptionPosition(str, Enum): class CaptionDevice(str, Enum): GPU = 'GPU if available' CPU = 'CPU' + + +class ExportFilter(str, Enum): + NONE = 'All images' + FILTERED = 'Filtered images' + SELECTED = 'Selected images' + + +Presets = { + 'manual': (0, 0, 1, '1:1, 2:1, 3:2, 4:3, 16:9, 21:9'), + 'Direct feed through': (0, 1, 1, '1:1, 2:1, 3:2, 4:3, 16:9, 21:9'), + 'SD1': (512, 64, 8, '512:512, 640:320, 576:384, 512:384, 640:384, 768:320'), + 'SDXL, SD3, Flux': (1024, 64, 8, '1024:1024, 1408:704, 1216:832, 1152:896, 1344:768, 1536:640') +} + +class MaskingStrategy(str, Enum): + IGNORE = 'ignore' + REPLACE = 'replace' + REMOVE = 'remove' + MASK_FILE = 'create mask files' + + +class MaskedContent(str, Enum): + ORIGINAL = 'original' + BLUR = 'blur' + BLUR_NOISE = 'blur + noise' + GREY = 'grey' + GREY_NOISE = 'grey + noise' + BLACK = 'black' + WHITE = 'white' + + +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 CaptionStrategy(str, Enum): + TAG_LIST = 'tag list (using tag separator)' + FIRST = 'only first tag' + LAST = 'only last tag' + ENUMERATION = 'enumeration ("t1, t2, t3, and t4")' + PREFIX_ENUMERATION = 'prefixed enumeration ("t1 t2, t3, and t4")' + + +class HashNewlineHandling(str, Enum): + IGNORE = 'No special handling' + MULTILINE = 'Create additional line' + MULTIFILE = 'Create additional file' diff --git a/taggui/utils/grid.py b/taggui/utils/grid.py new file mode 100644 index 00000000..99e7646b --- /dev/null +++ b/taggui/utils/grid.py @@ -0,0 +1,94 @@ +from math import ceil, floor +from PySide6.QtCore import (QPoint, QPointF, QRect, QSize) +from utils.enums import BucketStrategy +from utils.settings import settings +import utils.target_dimension as target_dimension + +class Grid: + """Span a grid inside the screen. + + The screen is adjusted according to bucket strategy and scaled to the + target dimension and then the grid respects the latent size. + """ + + def __init__(self, screen: QRect): + # the full image or the user cropped part of it + self.screen: QRect + # the visible part of the screen, i.e. the bucket cropped part of it + self.visible: QRect + # the size of the exported image + self.target: QSize + self.scale_x: float + self.scale_y: float + self.aspect_ratio: tuple[int, int, float] | None = None + + self.update(screen) + + def update(self, screen: QRect | None = None): + assert screen == None or isinstance(screen, QRect) + bucket_strategy = settings.value('export_bucket_strategy', type=str) + if screen != None: + self.screen = screen + + if self.screen.width() == 0 or self.screen.height() == 0: + self.visible = self.screen + self.target = QSize(1, 1) + self.scale_x = 1 + self.scale_y = 1 + return + + vis_size = self.screen.size() + self.target = target_dimension.get(vis_size) + aspect_ratio = self.target.width() / self.target.height() + + if (bucket_strategy == BucketStrategy.CROP or + bucket_strategy == BucketStrategy.CROP_SCALE): + if (self.screen.height() * self.target.width() + < self.target.height() * self.screen.width()): # too wide + vis_size.setWidth(floor(self.screen.height() * aspect_ratio)) + else: + vis_size.setHeight(floor(self.screen.width() / aspect_ratio)) + if bucket_strategy == BucketStrategy.CROP_SCALE: + vis_size.setWidth(floor((self.screen.width() + vis_size.width())/2)) + vis_size.setHeight(floor((self.screen.height() + vis_size.height())/2)) + + delta = self.screen.size() - vis_size + self.visible = self.screen.adjusted(floor(delta.width()/2), + floor(delta.height()/2), + -ceil(delta.width()/2), + -ceil(delta.height()/2)) + + self.scale_x = self.target.width() / self.visible.width() + self.scale_y = self.target.height() / self.visible.height() + + self.aspect_ratio = target_dimension.get_noteable_aspect_ratio( + self.target.width(), self.target.height()) + + def is_visible_equal_screen_size(self) -> bool: + return self.screen.size() == self.visible + + def map_raw(self, point: QPoint) -> QPointF: + """Translate point into screen coordinates.""" + assert isinstance(point, QPoint) + return QPointF((point.x()-self.visible.x())*self.scale_x, + (point.y()-self.visible.y())*self.scale_y) + + + def map(self, point: QPoint, method = round) -> QPoint: + """Align the point to the closest position on the grid aligned at + `base_point` and with a step width of `grid`. + """ + assert isinstance(point, QPoint) + latent_size = max(settings.value('export_latent_size', type=int), 1) + raw = self.map_raw(point) + return QPoint(method(raw.x()/latent_size)*latent_size, + method(raw.y()/latent_size)*latent_size) + + def snap(self, point: QPoint, method = round) -> QPointF: + """Align the point to the closest position on the grid but in + screen coordinates. + """ + assert isinstance(point, QPoint) + mapped = self.map(point, method) + return QPointF(mapped.x()/self.scale_x + self.visible.x(), + mapped.y()/self.scale_y + self.visible.y()) diff --git a/taggui/utils/icons.py b/taggui/utils/icons.py new file mode 100644 index 00000000..9e22f495 --- /dev/null +++ b/taggui/utils/icons.py @@ -0,0 +1,52 @@ +from pathlib import Path + +from PySide6.QtGui import QColor, QPixmap, QIcon, QPainter, QPen, Qt, QPainterPath, QImage +from PySide6.QtCore import QRect + +from taggui.utils.utils import get_resource_path + +ICON_PATH = Path('images/icon.ico') +TOGGLE_MARKING_ICON_PATH = Path('images/toggle_marking.png') +SHOW_MARKINGS_ICON_PATH = Path('images/show_marking.png') +SHOW_LABELS_ICON_PATH = Path('images/show_label.png') +SHOW_MARKING_LATENT_ICON_PATH = Path('images/show_marking_latent.png') + +def taggui_icon(): + return QIcon(str(get_resource_path(ICON_PATH))) + +def toggle_marking_icon(): + return QIcon(str(get_resource_path(TOGGLE_MARKING_ICON_PATH))) + +def show_markings_icon(): + return QIcon(str(get_resource_path(SHOW_MARKINGS_ICON_PATH))) + +def show_labels_icon(): + return QIcon(str(get_resource_path(SHOW_LABELS_ICON_PATH))) + +def show_marking_latent_icon(): + return QIcon(str(get_resource_path(SHOW_MARKING_LATENT_ICON_PATH))) + +def create_add_box_icon(color: QColor) -> QPixmap: + """Create a QPixmap for an icon""" + pixmap = QPixmap(32, 32) + pixmap.fill(QColor('transparent')) + + # Create a painter to draw on the pixmap + painter = QPainter(pixmap) + + # Draw a bordered rectangle in the specified color + rect = QRect(2, 2, 28, 28) + painter.setPen(QPen(color, 2)) + painter.drawRect(rect) + + # Draw a plus sign in the middle + painter.setPen(QPen(Qt.black, 1)) + path = QPainterPath() + path.moveTo(16, 10) + path.lineTo(16, 22) + path.moveTo(10, 16) + path.lineTo(22, 16) + painter.drawPath(path) + painter.end() + + return pixmap diff --git a/taggui/utils/image.py b/taggui/utils/image.py index f2637da4..1c68e355 100644 --- a/taggui/utils/image.py +++ b/taggui/utils/image.py @@ -1,12 +1,34 @@ +from enum import Enum from dataclasses import dataclass, field from pathlib import Path +from PySide6.QtCore import QRect, QSize from PySide6.QtGui import QIcon +class ImageMarking(str, Enum): + CROP = 'crop' + HINT = 'hint' + INCLUDE = 'include in mask' + EXCLUDE = 'exclude from mask' + NONE = 'no marking' + + +@dataclass +class Marking: + label: str + type: ImageMarking + rect: QRect + confidence: float = 1.0 + + @dataclass class Image: path: Path dimensions: tuple[int, int] | None tags: list[str] = field(default_factory=list) + target_dimension: QSize | None = None + crop: QRect | None = None + markings: list[Marking] = field(default_factory=list) + rating: float = 0.0 thumbnail: QIcon | None = None diff --git a/taggui/utils/rect.py b/taggui/utils/rect.py new file mode 100644 index 00000000..a7768878 --- /dev/null +++ b/taggui/utils/rect.py @@ -0,0 +1,166 @@ +from enum import Enum +from math import floor, ceil + +from PySide6.QtCore import QPoint, QRect, QSize, QPointF, QRectF, Qt + + +class RectPosition(str, Enum): + TL = 'top left' + TOP = 'top' + TR = 'top right' + RIGHT = 'right' + BR = 'bottom right' + BOTTOM = 'bottom' + BL = 'bottom left' + LEFT = 'left' + NONE = 'none' + +def map_rect_position_to_cursor(handle: RectPosition) -> Qt.CursorShape | None: + if handle == RectPosition.TL or handle == RectPosition.BR: + return Qt.CursorShape.SizeFDiagCursor + elif handle == RectPosition.TR or handle == RectPosition.BL: + return Qt.CursorShape.SizeBDiagCursor + elif handle == RectPosition.TOP or handle == RectPosition.BOTTOM: + return Qt.CursorShape.SizeVerCursor + elif handle == RectPosition.LEFT or handle == RectPosition.RIGHT: + return Qt.CursorShape.SizeHorCursor + return None + +def get_rect_position(left: bool, right: bool, top: bool, bottom: bool) -> RectPosition: + if top: + if left: + return RectPosition.TL + elif right: + return RectPosition.TR + return RectPosition.TOP + elif bottom: + if left: + return RectPosition.BL + elif right: + return RectPosition.BR + return RectPosition.BOTTOM + if left: + return RectPosition.LEFT + elif right: + return RectPosition.RIGHT + + return RectPosition.NONE + +def flip_rect_position(pos: RectPosition, h_flip: bool, v_flip: bool) -> RectPosition: + if pos == RectPosition.NONE: + return RectPosition.NONE + + if pos == RectPosition.TL or pos == RectPosition.TOP or pos == RectPosition.TR: + v = 2 if v_flip else 0 + elif pos == RectPosition.LEFT or pos == RectPosition.RIGHT: + v = 1 + else: + v = 0 if v_flip else 2 + + if pos == RectPosition.TL or pos == RectPosition.LEFT or pos == RectPosition.BL: + h = 2 if h_flip else 0 + elif pos == RectPosition.TOP or pos == RectPosition.BOTTOM: + h = 1 + else: + h = 0 if h_flip else 2 + + return { + 0: RectPosition.TL, 1: RectPosition.TOP, 2: RectPosition.TR, + 10: RectPosition.LEFT, 12: RectPosition.RIGHT, + 20: RectPosition.BL, 21: RectPosition.BOTTOM, 22: RectPosition.BR, + }[h+10*v] + +def change_rect(rect: QRect, rect_pos: RectPosition, pos: QPoint) -> QRect: + """Change `rect` to move `rect_pos` at `pos`.""" + if rect_pos == RectPosition.TL: + rect.setTopLeft(pos) + elif rect_pos == RectPosition.TOP: + rect.setTop(pos.y()) + elif rect_pos == RectPosition.TR: + rect.setTopRight(pos) + elif rect_pos == RectPosition.RIGHT: + rect.setRight(pos.x() - 1) + elif rect_pos == RectPosition.BR: + rect.setBottomRight(pos - QPoint(1, 1)) + elif rect_pos == RectPosition.BOTTOM: + rect.setBottom(pos.y() - 1) + elif rect_pos == RectPosition.BL: + rect.setBottomLeft(pos) + elif rect_pos == RectPosition.LEFT: + rect.setLeft(pos.x()) + return rect + +def change_rectF(rect: QRectF, rect_pos: RectPosition, pos: QPointF) -> QRectF: + """Change `rect` to move `rect_pos` at `pos`.""" + if rect_pos == RectPosition.TL: + rect.setTopLeft(pos) + elif rect_pos == RectPosition.TOP: + rect.setTop(pos.y()) + elif rect_pos == RectPosition.TR: + rect.setTopRight(pos) + elif rect_pos == RectPosition.RIGHT: + rect.setRight(pos.x()) + elif rect_pos == RectPosition.BR: + rect.setBottomRight(pos) + elif rect_pos == RectPosition.BOTTOM: + rect.setBottom(pos.y()) + elif rect_pos == RectPosition.BL: + rect.setBottomLeft(pos) + elif rect_pos == RectPosition.LEFT: + rect.setLeft(pos.x()) + return rect + +def change_rect_round(rect: QRect, rect_pos: RectPosition, pos: QPointF, grow: bool) -> QRect: + """Change `rect` to move `rect_pos` at `pos` and round by growing or shrinking the rect as `grow` demands.""" + round_tl = floor if grow else ceil + round_br = ceil if grow else floor + + if rect_pos == RectPosition.TL: + rect.setTopLeft(QPoint(round_tl(pos.x()), round_tl(pos.y()))) + elif rect_pos == RectPosition.TOP: + rect.setTop(round_tl(pos.y())) + elif rect_pos == RectPosition.TR: + rect.setTopRight(QPoint(round_br(pos.x()), round_tl(pos.y()))) + elif rect_pos == RectPosition.RIGHT: + rect.setRight(round_br(pos.x())) + elif rect_pos == RectPosition.BR: + rect.setBottomRight(QPoint(round_br(pos.x()), round_br(pos.y()))) + elif rect_pos == RectPosition.BOTTOM: + rect.setBottom(round_br(pos.y())) + elif rect_pos == RectPosition.BL: + rect.setBottomLeft(QPoint(round_tl(pos.x()), round_br(pos.y()))) + elif rect_pos == RectPosition.LEFT: + rect.setLeft(round_tl(pos.x())) + return rect + +def change_rect_to_match_size(rect: QRectF, rect_pos: RectPosition, size: QSize) -> QRect: + """Change the `rect` at place `rect_pos` so that the size matches `size`. + + Moving one side will ignore the value in the size of the perpendicular side. + """ + rect_new = QRectF(rect) + if rect_pos == RectPosition.TL: + rect_new.setSize(size) + rect_new.moveBottomRight(rect.bottomRight()) + elif rect_pos == RectPosition.TOP: + rect_new.setHeight(size.height()) + rect_new.moveBottom(rect.bottom()) + elif rect_pos == RectPosition.TR: + rect_new.setSize(size) + rect_new.moveBottomLeft(rect.bottomLeft()) + elif rect_pos == RectPosition.RIGHT: + rect_new.setWidth(size.width()) + rect_new.moveLeft(rect.left()) + elif rect_pos == RectPosition.BR: + rect_new.setSize(size) + rect_new.moveTopLeft(rect.topLeft()) + elif rect_pos == RectPosition.BOTTOM: + rect_new.setHeight(size.height()) + rect_new.moveTop(rect.top()) + elif rect_pos == RectPosition.BL: + rect_new.setSize(size) + rect_new.moveTopRight(rect.topRight()) + elif rect_pos == RectPosition.LEFT: + rect_new.setWidth(size.width()) + rect_new.moveRight(rect.right()) + return rect_new.toRect() diff --git a/taggui/utils/settings.py b/taggui/utils/settings.py index 88e1f0f5..5358a035 100644 --- a/taggui/utils/settings.py +++ b/taggui/utils/settings.py @@ -1,4 +1,4 @@ -from PySide6.QtCore import QSettings +from PySide6.QtCore import QSettings, Signal # Defaults for settings that are accessed from multiple places. DEFAULT_SETTINGS = { @@ -9,17 +9,46 @@ 'tag_separator': ',', 'insert_space_after_tag_separator': True, 'autocomplete_tags': True, - 'models_directory_path': '' + 'models_directory_path': '', + 'marking_models_directory_path': '', + 'export_filter': 'All images', + 'export_preset': 'SDXL, SD3, Flux', + 'export_resolution': 1024, + 'export_bucket_res_size': 64, + 'export_latent_size': 8, + 'export_quantize_alpha': True, + 'export_masking_strategy': 'remove', + 'export_masked_content': 'blur + noise', + 'export_preferred_sizes' : '1024:1024, 1408:704, 1216:832, 1152:896, 1344:768, 1536:640', + 'export_upscaling': False, + 'export_bucket_strategy': 'crop', + 'export_format': '.png - PNG', + 'export_quality': 100, + 'export_color_space': 'sRGB', + 'export_caption_algorithm': 'tag list (using tag separator)', + 'export_separate_newline': 'Create additional line', + 'export_directory_path': '', + 'export_keep_dir_structure': False, + 'export_filter_hashtag': True } -def get_settings() -> QSettings: - settings = QSettings('taggui', 'taggui') - return settings +class Settings(QSettings): + # Signal that shows that the setting with the given string was changes + change = Signal(str, object, name='settingsChanged') + + def __init__(self): + super().__init__('taggui', 'taggui') + + def setValue(self, key, value): + super().setValue(key, value) + self.change.emit(key, value) + +# Common shared instance to ensure the Signal is also shared +settings = Settings() def get_tag_separator() -> str: - settings = get_settings() tag_separator = settings.value( 'tag_separator', defaultValue=DEFAULT_SETTINGS['tag_separator'], type=str) diff --git a/taggui/utils/settings_widgets.py b/taggui/utils/settings_widgets.py index 43b0b5ec..3c5ab8bf 100644 --- a/taggui/utils/settings_widgets.py +++ b/taggui/utils/settings_widgets.py @@ -4,14 +4,15 @@ 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, 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)) @@ -21,14 +22,14 @@ class SettingsComboBox(QComboBox): def __init__(self, key: str, default: str | None = None): super().__init__() self.key = key - self.default = default - self.settings = get_settings() + if not settings.contains(key): + 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 = settings.value(self.key, type=str) super().addItems(texts) self.currentTextChanged.connect( - lambda text: self.settings.setValue(self.key, text)) + lambda text: settings.setValue(self.key, text)) if setting: self.setCurrentText(setting) @@ -44,17 +45,17 @@ def __init__(self, key: str, default: float, minimum: float, # The range must be set here so that the setting value is not clamped # by the default range. self.setRange(minimum, maximum) - settings = get_settings() self.setValue(settings.value(key, default, type=float)) self.valueChanged.connect(lambda value: settings.setValue(key, value)) 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 +64,19 @@ 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..8be9367c --- /dev/null +++ b/taggui/utils/target_dimension.py @@ -0,0 +1,203 @@ +import sys +from math import floor, sqrt +import re + +from PySide6.QtCore import QSize + +from utils.settings import DEFAULT_SETTINGS, settings + +# singleton data store +_preferred_sizes : list[QSize] = [] +notable_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 = notable_aspect_ratios.copy() + +settings.change.connect(lambda: _preferred_sizes.clear()) + +def get_preferred_sizes(): + global _preferred_sizes + if not _preferred_sizes: + prepare() + return _preferred_sizes + + +def prepare() -> list[tuple[int, int, float]] | None: + """ + Prepare by parsing the user supplied preferred sizes. + + Return + ------ + The same list of aspect ratios (when supplied) but extrended by the real + aspect ratios of the preferred sizes. + """ + global _preferred_sizes + global aspect_ratios + _preferred_sizes = [] + aspect_ratios = notable_aspect_ratios.copy() + for res_str in re.split(r'\s*,\s*', + settings.value('export_preferred_sizes', type=str) 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 is not 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: + ar_delta = abs(ar[2] - aspect_ratio) + if ar_delta < 1e-4: + # already included + break + if ar_delta < 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 calculate_cropped_area(width: int, height: int, test_width: int, test_height: int) -> int: + original_aspect_ratio = width / height + target_aspect_ratio = test_width / test_height + + if original_aspect_ratio > target_aspect_ratio: + # Crop horizontally + cropped_area = height * (width - (height * test_width) / test_height) + else: + # Crop vertically + cropped_area = width * (height - (width * test_height) / test_width) + return int(cropped_area) + +def get(dimensions: QSize) -> QSize: + """ + 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: QSize + The width and height of the image + """ + global _preferred_sizes + width, height = dimensions.toTuple() + # The target resolution of the AI model. The target image pixels + # will not exceed the square of this number + resolution = settings.value('export_resolution', defaultValue=DEFAULT_SETTINGS['export_resolution'], type=int) + # Is upscaling of images allowed? + upscaling = settings.value('export_upscaling', defaultValue=DEFAULT_SETTINGS['export_upscaling'], type=bool) + # The resolution of the buckets + bucket_res = 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 QSize((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 = floor(resolution * sqrt(width/height)) + opt_height = floor(resolution * sqrt(height/width)) + + loss = 1e10 + for dx, dy in [(0,0), (0,1), (1,0), (1,1)]: + opt_width += dx + opt_height += dy + + 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 + test_width = max(opt_width // bucket_res, 1) * bucket_res + test_height = max((height * test_width / width) // bucket_res, 1) * bucket_res + test_loss = calculate_cropped_area(width, height, test_width, test_height) + if (test_width, test_height) in _preferred_sizes: + test_loss *= preferred_sizes_bonus + if test_loss < loss or (test_loss == loss and + (candidate_width < test_width or candidate_height < test_height)): + candidate_width = test_width + candidate_height = test_height + loss = test_loss + # 1.2: exact height + test_height = max(opt_height // bucket_res, 1) * bucket_res + test_width = max((width * test_height / height) // bucket_res, 1) * bucket_res + test_loss = calculate_cropped_area(width, height, test_width, test_height) + if (test_height, test_width) in _preferred_sizes: + test_loss *= preferred_sizes_bonus + if test_loss < loss or (test_loss == loss and + (candidate_width < test_width or candidate_height < test_height)): + 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 = max(opt_width // bucket_res + delta, 1) * bucket_res + test_height = max((height * test_width / width) // bucket_res, 1) * bucket_res + if test_width * test_height > max_pixels: + break + if (test_width > width or test_height > height) and not upscaling: + break + test_loss = calculate_cropped_area(width, height, test_width, test_height) + if (test_height, test_width) in _preferred_sizes: + test_loss *= preferred_sizes_bonus + if test_loss < loss or (test_loss == loss and + (candidate_width < test_width or candidate_height < test_height)): + candidate_width = test_width + candidate_height = test_height + loss = test_loss + # 2.2: exact height + for delta in range(1, 10): + test_height = max(opt_height // bucket_res + delta, 1) * bucket_res + test_width = max((width * test_height / height) // bucket_res, 1) * bucket_res + if test_width * test_height > max_pixels: + break + if (test_width > width or test_height > height) and not upscaling: + break + test_loss = calculate_cropped_area(width, height, test_width, test_height) + if (test_height, test_width) in _preferred_sizes: + test_loss *= preferred_sizes_bonus + if test_loss < loss or (test_loss == loss and + (candidate_width < test_width or candidate_height < test_height)): + candidate_width = test_width + candidate_height = test_height + loss = test_loss + + return QSize(candidate_width, candidate_height) + +def get_noteable_aspect_ratio(width: float|int, height: float|int) -> tuple[int, int, bool] | None: + """Test whether the aspect_ratio is noteable and return it.""" + aspect_ratio = width / height if height > 0 else 100 + for ar in aspect_ratios: + if abs(ar[2] - aspect_ratio) < 1e-3: + return ar[0], ar[1], (width, height) in _preferred_sizes + elif abs(1/ar[2] - aspect_ratio) < 1e-3: + return ar[1], ar[0], (width, height) in _preferred_sizes + return None diff --git a/taggui/widgets/auto_captioner.py b/taggui/widgets/auto_captioner.py index 58935904..8cd5320b 100644 --- a/taggui/widgets/auto_captioner.py +++ b/taggui/widgets/auto_captioner.py @@ -15,7 +15,7 @@ from models.image_list_model import ImageListModel from utils.big_widgets import TallPushButton from utils.enums import CaptionDevice, CaptionPosition -from utils.settings import DEFAULT_SETTINGS, get_settings, get_tag_separator +from utils.settings import DEFAULT_SETTINGS, settings, get_tag_separator from utils.settings_widgets import (FocusedScrollSettingsComboBox, FocusedScrollSettingsDoubleSpinBox, FocusedScrollSettingsSpinBox, @@ -50,7 +50,6 @@ def __init__(self): class CaptionSettingsForm(QVBoxLayout): def __init__(self): super().__init__() - self.settings = get_settings() try: import bitsandbytes self.is_bitsandbytes_available = True @@ -74,6 +73,15 @@ def __init__(self): self.caption_position_combo_box = FocusedScrollSettingsComboBox( key='caption_position') self.caption_position_combo_box.addItems(list(CaptionPosition)) + self.skip_hash_container = QWidget() + skip_hash_layout = QHBoxLayout() + skip_hash_layout.setAlignment(Qt.AlignmentFlag.AlignLeft) + skip_hash_layout.setContentsMargins(0, 0, 0, 0) + self.skip_hash_check_box = SettingsBigCheckBox( + key='skip_hash', default=True) + skip_hash_layout.addWidget(QLabel('Skip hash tags when inserting in prompt')) + skip_hash_layout.addWidget(self.skip_hash_check_box) + self.skip_hash_container.setLayout(skip_hash_layout) self.device_combo_box = FocusedScrollSettingsComboBox(key='device') self.device_combo_box.addItems(list(CaptionDevice)) self.load_in_4_bit_container = QWidget() @@ -85,6 +93,15 @@ def __init__(self): load_in_4_bit_layout.addWidget(QLabel('Load in 4-bit')) load_in_4_bit_layout.addWidget(self.load_in_4_bit_check_box) self.load_in_4_bit_container.setLayout(load_in_4_bit_layout) + self.limit_to_crop_container = QWidget() + limit_to_crop_layout = QHBoxLayout() + limit_to_crop_layout.setAlignment(Qt.AlignmentFlag.AlignLeft) + limit_to_crop_layout.setContentsMargins(0, 0, 0, 0) + self.limit_to_crop_check_box = SettingsBigCheckBox( + key='limit_to_crop', default=True) + limit_to_crop_layout.addWidget(QLabel('Limit to crop')) + limit_to_crop_layout.addWidget(self.limit_to_crop_check_box) + self.limit_to_crop_container.setLayout(limit_to_crop_layout) self.remove_tag_separators_container = QWidget() remove_tag_separators_layout = QHBoxLayout( self.remove_tag_separators_container) @@ -105,10 +122,12 @@ def __init__(self): self.caption_start_line_edit) basic_settings_form.addRow('Caption position', self.caption_position_combo_box) + basic_settings_form.addRow(self.skip_hash_container) self.device_label = QLabel('Device') basic_settings_form.addRow(self.device_label, self.device_combo_box) basic_settings_form.addRow(self.load_in_4_bit_container) basic_settings_form.addRow(self.remove_tag_separators_container) + basic_settings_form.addRow(self.limit_to_crop_container) self.wd_tagger_settings_form_container = QWidget() wd_tagger_settings_form = QFormLayout( @@ -241,7 +260,7 @@ def __init__(self): self.load_in_4_bit_check_box.setChecked(False) def get_local_model_paths(self) -> list[str]: - models_directory_path = self.settings.value( + models_directory_path = settings.value( 'models_directory_path', defaultValue=DEFAULT_SETTINGS['models_directory_path'], type=str) if not models_directory_path: @@ -267,6 +286,7 @@ def show_settings_for_model(self, model_id: str): non_wd_tagger_widgets = [ self.prompt_label, self.prompt_text_edit, + self.skip_hash_container, self.caption_start_label, self.caption_start_line_edit, self.device_label, @@ -310,11 +330,13 @@ def get_caption_settings(self) -> dict: return { 'model_id': self.model_combo_box.currentText(), 'prompt': self.prompt_text_edit.toPlainText(), + 'skip_hash': self.skip_hash_check_box.isChecked(), 'caption_start': self.caption_start_line_edit.text(), 'caption_position': self.caption_position_combo_box.currentText(), 'device': self.device_combo_box.currentText(), 'gpu_index': self.gpu_index_spin_box.value(), 'load_in_4_bit': self.load_in_4_bit_check_box.isChecked(), + 'limit_to_crop': self.limit_to_crop_check_box.isChecked(), 'remove_tag_separators': self.remove_tag_separators_check_box.isChecked(), 'bad_words': self.bad_words_line_edit.text(), @@ -357,7 +379,6 @@ def __init__(self, image_list_model: ImageListModel, super().__init__() self.image_list_model = image_list_model self.image_list = image_list - self.settings = get_settings() self.is_captioning = False self.captioning_thread = None self.processor = None @@ -482,7 +503,7 @@ def generate_captions(self): self.progress_bar.setValue(0) self.progress_bar.show() tag_separator = get_tag_separator() - models_directory_path = self.settings.value( + models_directory_path = settings.value( 'models_directory_path', defaultValue=DEFAULT_SETTINGS['models_directory_path'], type=str) models_directory_path = (Path(models_directory_path) diff --git a/taggui/widgets/auto_markings.py b/taggui/widgets/auto_markings.py new file mode 100644 index 00000000..f544ac66 --- /dev/null +++ b/taggui/widgets/auto_markings.py @@ -0,0 +1,314 @@ +import sys +from pathlib import Path + +from PySide6.QtCore import Signal, QModelIndex, Qt, Slot +from PySide6.QtGui import QTextCursor +from PySide6.QtWidgets import (QDockWidget, QProgressBar, QPlainTextEdit, + QWidget, QVBoxLayout, QScrollArea, + QAbstractScrollArea, QFrame, QFormLayout, + QMessageBox, QTableWidget, QHeaderView, + QTableWidgetItem, QComboBox) + +from utils.icons import create_add_box_icon +from models.image_list_model import ImageListModel +from utils.utils import pluralize +from utils.big_widgets import TallPushButton +from utils.settings import settings, DEFAULT_SETTINGS +from utils.settings_widgets import (FocusedScrollSettingsComboBox, + FocusedScrollSettingsDoubleSpinBox, + FocusedScrollSettingsSpinBox) +from widgets.auto_captioner import (set_text_edit_height, + restore_stdout_and_stderr, HorizontalLine) +from widgets.image_list import ImageList +from auto_marking.marking_thread import MarkingThread +from dialogs.caption_multiple_images_dialog import CaptionMultipleImagesDialog + + +class MarkingSettingsForm(QVBoxLayout): + model_selected = Signal(bool) + + def __init__(self): + super().__init__() + basic_settings_form = QFormLayout() + basic_settings_form.setRowWrapPolicy( + QFormLayout.RowWrapPolicy.WrapAllRows) + basic_settings_form.setFieldGrowthPolicy( + QFormLayout.FieldGrowthPolicy.ExpandingFieldsGrow) + self.model_combo_box = FocusedScrollSettingsComboBox(key='marking_model_id') + self.model_combo_box.setPlaceholderText('Set marking model directory in "Settings..."') + self.model_combo_box.activated.connect(lambda _: self.model_selected.emit(True)) + self.get_local_model_paths() + settings.change.connect(lambda key, value: self.get_local_model_paths() + if key == 'marking_models_directory_path' else 0) + basic_settings_form.addRow('Model', self.model_combo_box) + + self.class_table = QTableWidget(0, 2) + self.class_table.setHorizontalHeaderLabels(['Class', 'Marking']) + self.class_table.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch) + self.class_table.verticalHeader().setSectionResizeMode(QHeaderView.ResizeToContents) + basic_settings_form.addRow('Classes', self.class_table) + + self.toggle_advanced_settings_form_button = TallPushButton( + 'Show Advanced Settings') + + self.advanced_settings_form_container = QWidget() + advanced_settings_form = QFormLayout( + self.advanced_settings_form_container) + advanced_settings_form.setLabelAlignment(Qt.AlignmentFlag.AlignRight) + advanced_settings_form.setFieldGrowthPolicy( + QFormLayout.FieldGrowthPolicy.ExpandingFieldsGrow) + # Sets the minimum confidence threshold for detections. + # Objects detected with confidence below this threshold will be + # disregarded. Adjusting this value can help reduce false positives. + self.confidence_spin_box = FocusedScrollSettingsDoubleSpinBox( + key='confidence', default=0.25, minimum=0.01, maximum=1.0) + self.confidence_spin_box.setSingleStep(0.01) + advanced_settings_form.addRow('Confidence', + self.confidence_spin_box) + # Intersection Over Union (IoU) threshold for Non-Maximum Suppression + # (NMS). Lower values result in fewer detections by eliminating + # overlapping boxes, useful for reducing duplicates. + self.iou_spin_box = FocusedScrollSettingsDoubleSpinBox( + key='iou', default=0.7, minimum=0.01, maximum=1.0) + self.iou_spin_box.setSingleStep(0.01) + advanced_settings_form.addRow('Intersection Over Union (IoU)', + self.iou_spin_box) + # Maximum number of detections allowed per image. + # Limits the total number of objects the model can detect in a single + # inference, preventing excessive outputs in dense scenes. + self.max_det_spin_box = FocusedScrollSettingsSpinBox( + key='max_det', default=300, minimum=1, maximum=500) + advanced_settings_form.addRow('Maximum number of detections', self.max_det_spin_box) + self.advanced_settings_form_container.hide() + + self.addLayout(basic_settings_form) + self.horizontal_line = HorizontalLine() + self.addWidget(self.horizontal_line) + self.addWidget(self.toggle_advanced_settings_form_button) + self.addWidget(self.advanced_settings_form_container) + + self.toggle_advanced_settings_form_button.clicked.connect( + self.toggle_advanced_settings_form) + + def get_local_model_paths(self): + models_directory_path = settings.value( + 'marking_models_directory_path', + defaultValue=DEFAULT_SETTINGS['marking_models_directory_path'], + type=str) + if not models_directory_path: + return + models_directory_path = Path(models_directory_path) + print(f'Loading local auto-marking model paths under ' + f'{models_directory_path}...') + config_paths = set(models_directory_path.glob('**/*.pt')) + self.model_selected.emit(False) + self.model_combo_box.clear() + if len(config_paths) == 0: + self.model_combo_box.setPlaceholderText( + 'Set marking model directory in "Settings..."') + else: + self.model_combo_box.setPlaceholderText('Select marking model') + for path in config_paths: + self.model_combo_box.addItem( + str(path.relative_to(models_directory_path)), userData=path) + + @Slot() + def toggle_advanced_settings_form(self): + if self.advanced_settings_form_container.isHidden(): + self.advanced_settings_form_container.show() + self.toggle_advanced_settings_form_button.setText( + 'Hide Advanced Settings') + else: + self.advanced_settings_form_container.hide() + self.toggle_advanced_settings_form_button.setText( + 'Show Advanced Settings') + + def get_marking_settings(self) -> dict: + return { + 'model_path': self.model_combo_box.currentData(), + 'conf': self.confidence_spin_box.value(), + 'iou': self.iou_spin_box.value(), + 'max_det': self.max_det_spin_box.value(), + 'classes': [] + } + +class AutoMarkings(QDockWidget): + marking_generated = Signal(QModelIndex, list) + + def __init__(self, image_list_model: ImageListModel, + image_list: ImageList, parent): + super().__init__(parent) + self.image_list_model = image_list_model + self.image_list = image_list + self.is_marking = False + self.marking_thread = None + self.show_alert_when_finished = False + # Whether the last block of text in the console text edit should be + # replaced with the next block of text that is outputted. + self.replace_last_console_text_edit_block = False + # Each `QDockWidget` needs a unique object name for saving its state. + self.setObjectName('auto_markings') + self.setWindowTitle('Auto-Markings') + self.setAllowedAreas(Qt.DockWidgetArea.LeftDockWidgetArea | + Qt.DockWidgetArea.RightDockWidgetArea) + + self.start_cancel_button = TallPushButton('Start Auto-Marking') + self.start_cancel_button.setEnabled(False) + self.progress_bar = QProgressBar() + self.progress_bar.setFormat('%v / %m images marked (%p%)') + self.progress_bar.hide() + self.console_text_edit = QPlainTextEdit() + set_text_edit_height(self.console_text_edit, 4) + self.console_text_edit.setReadOnly(True) + self.console_text_edit.hide() + container = QWidget() + layout = QVBoxLayout(container) + layout.addWidget(self.start_cancel_button) + layout.addWidget(self.progress_bar) + layout.addWidget(self.console_text_edit) + self.marking_settings_form = MarkingSettingsForm() + layout.addLayout(self.marking_settings_form) + scroll_area = QScrollArea() + scroll_area.setWidgetResizable(True) + scroll_area.setSizeAdjustPolicy( + QAbstractScrollArea.SizeAdjustPolicy.AdjustToContents) + scroll_area.setFrameShape(QFrame.Shape.NoFrame) + scroll_area.setWidget(container) + self.setWidget(scroll_area) + + self.start_cancel_button.clicked.connect( + self.start_or_cancel_marking) + self.marking_settings_form.model_selected.connect(lambda _: self.prepare_generation()) + self.marking_settings_form.model_selected.connect(self.start_cancel_button.setEnabled) + + @Slot() + def start_or_cancel_marking(self): + if self.is_marking: + # Cancel marking. + self.marking_thread.is_canceled = True + self.start_cancel_button.setEnabled(False) + self.start_cancel_button.setText('Canceling Auto-Marking...') + else: + # Start marking. + self.generate_markings() + + def set_is_marking(self, is_marking: bool): + self.is_marking = is_marking + button_text = ('Cancel Auto-Marking' if is_marking + else 'Start Auto-Marking') + self.start_cancel_button.setText(button_text) + + @Slot(str) + def update_console_text_edit(self, text: str): + # '\x1b[A' is the ANSI escape sequence for moving the cursor up. + if text == '\x1b[A': + self.replace_last_console_text_edit_block = True + return + text = text.strip() + if not text: + return + if self.console_text_edit.isHidden(): + self.console_text_edit.show() + if self.replace_last_console_text_edit_block: + self.replace_last_console_text_edit_block = False + # Select and remove the last block of text. + self.console_text_edit.moveCursor(QTextCursor.MoveOperation.End) + self.console_text_edit.moveCursor( + QTextCursor.MoveOperation.StartOfBlock, + QTextCursor.MoveMode.KeepAnchor) + self.console_text_edit.textCursor().removeSelectedText() + # Delete the newline. + self.console_text_edit.textCursor().deletePreviousChar() + self.console_text_edit.appendPlainText(text) + + @Slot() + def show_alert(self): + if self.marking_thread.is_canceled: + return + if self.marking_thread.is_error: + icon = QMessageBox.Icon.Critical + text = ('An error occurred during marking. See the ' + 'Auto-Marking console for more information.') + else: + icon = QMessageBox.Icon.Information + text = 'Marking has finished.' + alert = QMessageBox() + alert.setIcon(icon) + alert.setText(text) + alert.exec() + + def prepare_generation(self): + selected_image_indices = self.image_list.get_selected_image_indices() + marking_settings = self.marking_settings_form.get_marking_settings() + self.marking_thread = MarkingThread( + self, self.image_list_model, selected_image_indices, + marking_settings) + self.marking_thread.text_outputted.connect( + self.update_console_text_edit) + self.marking_thread.clear_console_text_edit_requested.connect( + self.console_text_edit.clear) + self.marking_thread.marking_generated.connect( + self.marking_generated) + self.marking_thread.progress_bar_update_requested.connect( + self.progress_bar.setValue) + self.marking_thread.finished.connect( + lambda: self.set_is_marking(False)) + self.marking_thread.finished.connect(restore_stdout_and_stderr) + self.marking_thread.finished.connect(self.progress_bar.hide) + self.marking_thread.finished.connect( + lambda: self.start_cancel_button.setEnabled(True)) + if self.show_alert_when_finished: + self.marking_thread.finished.connect(self.show_alert) + self.marking_thread.preload_model() + self.marking_settings_form.class_table.setRowCount( + len(self.marking_thread.model.names)) + for row, (class_id, class_name) in enumerate( + self.marking_thread.model.names.items()): + self.marking_settings_form.class_table.setItem( + row, 0, QTableWidgetItem(class_name)) + combo = QComboBox() + combo.addItem('ignore') + combo.addItem(create_add_box_icon(Qt.gray), 'hint') + combo.addItem(create_add_box_icon(Qt.red), 'exclude') + combo.addItem(create_add_box_icon(Qt.green), 'include') + self.marking_settings_form.class_table.setCellWidget(row, 1, combo) + # NOTE: As this thread has no place to display the output, we keep + # `stdout` and `stderr`. + # Redirect `stdout` and `stderr` so that the outputs are displayed in + # the console text edit. + ###sys.stdout = self.marking_thread + ###sys.stderr = self.marking_thread + + @Slot() + def generate_markings(self): + selected_image_indices = self.image_list.get_selected_image_indices() + if self.marking_thread is None: + self.prepare_generation() + self.marking_thread.selected_image_indices = selected_image_indices + self.marking_thread.marking_settings = self.marking_settings_form.get_marking_settings() + classes = {} + for row, (class_id, class_name) in enumerate( + self.marking_thread.model.names.items()): + combo = self.marking_settings_form.class_table.cellWidget(row, 1).currentText() + if combo != 'ignore': + classes[class_id] = (class_name, combo) + self.marking_thread.marking_settings['classes'] = classes + selected_image_count = len(selected_image_indices) + self.image_list_model.add_to_undo_stack( + action_name=f'Generate ' + f'{pluralize('Marking', selected_image_count)}', + should_ask_for_confirmation=selected_image_count > 1) + if selected_image_count > 1: + confirmation_dialog = CaptionMultipleImagesDialog( + selected_image_count, 'Mark', 'Markings') + reply = confirmation_dialog.exec() + if reply != QMessageBox.StandardButton.Yes: + return + self.show_alert_when_finished = (confirmation_dialog + .show_alert_check_box.isChecked()) + self.set_is_marking(True) + if selected_image_count > 1: + self.progress_bar.setRange(0, selected_image_count) + self.progress_bar.setValue(0) + self.progress_bar.show() + self.marking_thread.start() diff --git a/taggui/widgets/image_list.py b/taggui/widgets/image_list.py index 18af0eb0..30fe3849 100644 --- a/taggui/widgets/image_list.py +++ b/taggui/widgets/image_list.py @@ -6,21 +6,22 @@ from PySide6.QtCore import (QFile, QItemSelection, QItemSelectionModel, QItemSelectionRange, QModelIndex, QSize, QUrl, Qt, - Signal, Slot) -from PySide6.QtGui import QDesktopServices + Signal, Slot, QPersistentModelIndex) +from PySide6.QtGui import QDesktopServices, QColor from PySide6.QtWidgets import (QAbstractItemView, QApplication, QDockWidget, QFileDialog, QHBoxLayout, QLabel, QLineEdit, QListView, QMenu, QMessageBox, QVBoxLayout, - QWidget) + QWidget, QStyledItemDelegate) from pyparsing import (CaselessKeyword, CaselessLiteral, Group, OpAssoc, ParseException, QuotedString, Suppress, Word, infix_notation, nums, one_of, printables) from models.proxy_image_list_model import ProxyImageListModel from utils.image import Image -from utils.settings import get_settings +from utils.settings import settings from utils.settings_widgets import SettingsComboBox from utils.utils import get_confirmation_dialog_reply, pluralize +from utils.grid import Grid def replace_filter_wildcards(filter_: str | list) -> str | list: @@ -48,12 +49,14 @@ def __init__(self): | QuotedString(quote_char="'", esc_char='\\') | Word(printables, exclude_chars='()')) - string_filter_keys = ['tag', 'caption', 'name', 'path'] + string_filter_keys = ['tag', 'caption', 'marking', 'crops', 'visible', + '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', 'stars', 'width', + 'height', 'area'] number_filter_expressions = [Group(CaselessLiteral(key) + Suppress(':') + comparison_operator + Word(nums)) for key in number_filter_keys] @@ -96,6 +99,35 @@ class SelectionMode(str, Enum): TOGGLE = 'Toggle' +class ImageDelegate(QStyledItemDelegate): + def __init__(self, parent=None): + super().__init__(parent) + self.labels = {} + + def sizeHint(self, option, index): + return index.data(Qt.ItemDataRole.SizeHintRole) + + def paint(self, painter, option, index): + super().paint(painter, option, index) + p_index = QPersistentModelIndex(index) + if p_index in self.labels: + label_text = self.labels[p_index] + painter.setBrush(QColor(255, 255, 255, 163)) + painter.drawRect(option.rect) + painter.drawText(option.rect, label_text, Qt.AlignCenter) + + def update_label(self, index: QModelIndex, label: str): + p_index = QPersistentModelIndex(index) + self.labels[p_index] = label + self.parent().update(p_index) + + def remove_label(self, index: QPersistentModelIndex): + p_index = QPersistentModelIndex(index) + if p_index in self.labels: + del self.labels[p_index] + self.parent().update(index) + + class ImageListView(QListView): tags_paste_requested = Signal(list, list) directory_reload_requested = Signal() @@ -106,6 +138,8 @@ def __init__(self, parent, proxy_image_list_model: ProxyImageListModel, self.proxy_image_list_model = proxy_image_list_model self.tag_separator = tag_separator self.setModel(proxy_image_list_model) + self.delegate = ImageDelegate(self) + self.setItemDelegate(self.delegate) self.setWordWrap(True) # If the actual height of the image is greater than 3 times the width, # the image will be scaled down to fit. @@ -167,6 +201,25 @@ def __init__(self, parent, proxy_image_list_model: ProxyImageListModel, def contextMenuEvent(self, event): self.context_menu.exec_(event.globalPos()) + @Slot(Grid) + def show_crop_size(self, grid): + index = self.currentIndex() + if index.isValid(): + image = index.data(Qt.ItemDataRole.UserRole) + if grid is None: + self.delegate.remove_label(index) + else: + crop_delta = grid.screen.size() - grid.visible.size() + crop_fit = max(crop_delta.width(), crop_delta.height()) + crop_fit_text = f' (-{crop_fit})' if crop_fit > 0 else '' + label = f'image: {image.dimensions[0]}x{image.dimensions[1]}\n'\ + f'crop: {grid.screen.width()}x{grid.screen.height()}{crop_fit_text}\n'\ + f'target: {grid.target.width()}x{grid.target.height()}' + if grid.aspect_ratio is not None: + label += '✅' if grid.aspect_ratio[2] else '' + label += f' {grid.aspect_ratio[0]}:{grid.aspect_ratio[1]}' + self.delegate.update_label(index, label) + @Slot() def invert_selection(self): selected_proxy_rows = {index.row() for index in self.selectedIndexes()} @@ -235,7 +288,6 @@ def move_selected_images(self): caption = (f'Select directory to move {selected_image_count} selected ' f'{pluralize("Image", selected_image_count)} and ' f'{pluralize("caption", selected_image_count)} to') - settings = get_settings() move_directory_path = QFileDialog.getExistingDirectory( parent=self, caption=caption, dir=settings.value('directory_path', type=str)) @@ -262,7 +314,6 @@ def copy_selected_images(self): caption = (f'Select directory to copy {selected_image_count} selected ' f'{pluralize("Image", selected_image_count)} and ' f'{pluralize("caption", selected_image_count)} to') - settings = get_settings() copy_directory_path = QFileDialog.getExistingDirectory( parent=self, caption=caption, dir=settings.value('directory_path', type=str)) diff --git a/taggui/widgets/image_tags_editor.py b/taggui/widgets/image_tags_editor.py index 65a85e05..de517778 100644 --- a/taggui/widgets/image_tags_editor.py +++ b/taggui/widgets/image_tags_editor.py @@ -9,7 +9,7 @@ from models.proxy_image_list_model import ProxyImageListModel from models.tag_counter_model import TagCounterModel from utils.image import Image -from utils.settings import DEFAULT_SETTINGS, get_settings +from utils.settings import DEFAULT_SETTINGS, settings from utils.text_edit_item_delegate import TextEditItemDelegate from utils.utils import get_confirmation_dialog_reply from widgets.image_list import ImageList @@ -30,7 +30,6 @@ def __init__(self, image_tag_list_model: QStringListModel, self.setPlaceholderText('Add Tag') self.setStyleSheet('padding: 8px;') - settings = get_settings() autocomplete_tags = settings.value( 'autocomplete_tags', defaultValue=DEFAULT_SETTINGS['autocomplete_tags'], type=bool) @@ -227,7 +226,8 @@ def reload_image_tags_if_changed(self, first_changed_index: QModelIndex, Reload the tags for the current image if its index is in the range of changed indices. """ - if (first_changed_index.row() <= self.image_index.row() + if (self.image_index and + first_changed_index.row() <= self.image_index.row() <= last_changed_index.row()): proxy_image_index = self.proxy_image_list_model.mapFromSource( self.image_index) diff --git a/taggui/widgets/image_viewer.py b/taggui/widgets/image_viewer.py index 0b67d3cc..b3b156c5 100644 --- a/taggui/widgets/image_viewer.py +++ b/taggui/widgets/image_viewer.py @@ -1,51 +1,963 @@ -from pathlib import Path +import re +from math import ceil, floor, sqrt +from PySide6.QtCore import (QModelIndex, QPersistentModelIndex, QPoint, QPointF, + QRect, QRectF, QSize, Qt, Signal, Slot) +from PySide6.QtGui import (QAction, QActionGroup, QColor, QIcon, + QPainter, QPainterPath, QPen, QPixmap, QTransform, + QMouseEvent) +from PySide6.QtWidgets import (QGraphicsItem, QGraphicsLineItem, + QGraphicsPixmapItem, QGraphicsRectItem, + QGraphicsTextItem, QGraphicsScene, QGraphicsView, + QMenu, QVBoxLayout, QWidget) +from utils.settings import settings +from models.proxy_image_list_model import ProxyImageListModel +from utils.image import Image, ImageMarking, Marking +import utils.target_dimension as target_dimension +from utils.grid import Grid +from utils.rect import (change_rect, change_rect_to_match_size, + flip_rect_position, get_rect_position, + map_rect_position_to_cursor, RectPosition) -from PySide6.QtCore import QModelIndex, QSize, Qt, Slot -from PySide6.QtGui import QImageReader, QPixmap, QResizeEvent -from PySide6.QtWidgets import QLabel, QSizePolicy, QVBoxLayout, QWidget +# The (inverse) golden ratio for showing hints during cropping +golden_ratio = 2 / (1 + sqrt(5)) -from models.proxy_image_list_model import ProxyImageListModel -from utils.image import Image +# Grid for alignment to latent space +grid = Grid(QRect(0, 0, 1, 1)) +marking_colors = { + ImageMarking.CROP: Qt.blue, + ImageMarking.HINT: Qt.gray, + ImageMarking.INCLUDE: Qt.green, + ImageMarking.EXCLUDE: Qt.red, +} -class ImageLabel(QLabel): - def __init__(self): - super().__init__() - self.image_path = None - self.setAlignment(Qt.AlignmentFlag.AlignCenter) - self.setSizePolicy(QSizePolicy.Policy.Expanding, - QSizePolicy.Policy.Expanding) - # This allows the label to shrink. - self.setMinimumSize(QSize(1, 1)) - - def resizeEvent(self, event: QResizeEvent): - """Reload the image whenever the label is resized.""" - if self.image_path: - self.load_image(self.image_path) - - def load_image(self, image_path: Path): - self.image_path = image_path - image_reader = QImageReader(str(image_path)) - # Rotate the image according to the orientation tag. - image_reader.setAutoTransform(True) - pixmap = QPixmap.fromImageReader(image_reader) - pixmap.setDevicePixelRatio(self.devicePixelRatio()) - pixmap = pixmap.scaled( - self.size() * pixmap.devicePixelRatio(), - Qt.AspectRatioMode.KeepAspectRatio, - Qt.TransformationMode.SmoothTransformation) - self.setPixmap(pixmap) +def calculate_grid(content: QRect): + global grid + grid = Grid(content) + +# Organization of Z value: +# 0: the image +# 1: the semi transparent include/exclude areas +# 2: the borders of the markings +# 3: the HUD +# 4: the currently manipulated marking +# 5: the cross hair lines + +class MarkingItem(QGraphicsRectItem): + # the halved size of the pen in local coordinates to make sure it stays the + # same during zooming + pen_half_width: float = 1.0 + # the minimal size of the active area in scene coordinates + handle_half_size: int = 5 + zoom_factor: float = 1.0 + # The size of the image this rect belongs to + image_size: QRect = QRect(0, 0, 1, 1) + # Static link to the single ImageGraphicsView in this application + image_view: bool = None + show_marking_latent: bool = True + handle_selected: RectPosition = RectPosition.NONE + show_crop_hint: bool = True + + def __init__(self, rect: QRect, rect_type: ImageMarking, interactive: bool, + parent = None): + super().__init__(rect.toRectF(), parent) + self.setFlag(QGraphicsItem.ItemIsSelectable, True) + self.rect_type = rect_type + self.label: MarkingLabel | None = None + self.color = marking_colors[rect_type] + self.setZValue(2) + if rect_type in [ImageMarking.INCLUDE, ImageMarking.EXCLUDE]: + self.area = QGraphicsRectItem(self) + self.area.setVisible(self.show_marking_latent) + self.area.setFlag(QGraphicsItem.ItemStacksBehindParent) + self.area.setZValue(1) + area_color = QColor(self.color) + area_color.setAlpha(127) + self.area.setBrush(area_color) + self.area.setPen(Qt.NoPen) + self.move() + if interactive: + MarkingItem.handle_selected = RectPosition.BR + + def move(self): + if self.rect_type == ImageMarking.CROP: + self.image_view.image_viewer.hud_item.setValues(self.rect(), MarkingItem.handle_selected) + elif self.rect_type == ImageMarking.INCLUDE: + self.area.setRect(QRectF(grid.snap(self.rect().toRect().topLeft(), ceil), + grid.snap(self.rect().toRect().adjusted(0,0,1,1).bottomRight(), floor))) + elif self.rect_type == ImageMarking.EXCLUDE: + self.area.setRect(QRectF(grid.snap(self.rect().toRect().topLeft(), floor), + grid.snap(self.rect().toRect().adjusted(0,0,1,1).bottomRight(), ceil))) + + def handleAt(self, point: QPointF) -> RectPosition: + handle_space = -min(self.pen_half_width - self.handle_half_size, + 0)/self.zoom_factor + return get_rect_position(point.x() < self.rect().left() + handle_space, + point.x() > self.rect().right() - handle_space, + point.y() < self.rect().top() + handle_space, + point.y() > self.rect().bottom() - handle_space) + + def mousePressEvent(self, event): + self.show_crop_hint = ((event.modifiers() & Qt.KeyboardModifier.AltModifier) != + Qt.KeyboardModifier.AltModifier) + MarkingItem.handle_selected = self.handleAt(event.pos()) + if (event.button() == Qt.MouseButton.LeftButton and + MarkingItem.handle_selected != RectPosition.NONE): + self.image_view.image_viewer.proxy_image_index.model().sourceModel().add_to_undo_stack( + action_name=f'Change marking geometry', should_ask_for_confirmation=False) + self.setZValue(4) + self.move() + elif (event.button() == Qt.MouseButton.RightButton and + MarkingItem.handle_selected != RectPosition.NONE): + pass + else: + event.ignore() + + def mouseMoveEvent(self, event): + if MarkingItem.handle_selected != RectPosition.NONE: + self.show_crop_hint = ((event.modifiers() & Qt.KeyboardModifier.AltModifier) != + Qt.KeyboardModifier.AltModifier) + if ((event.modifiers() & Qt.KeyboardModifier.ShiftModifier) == + Qt.KeyboardModifier.ShiftModifier): + if self.rect_type == ImageMarking.CROP: + bucket_res = settings.value('export_bucket_res_size', type=int) + rect_pre = change_rect(self.rect(), + MarkingItem.handle_selected, + event.pos()) + target_size = target_dimension.get(rect_pre.toRect().size()) + # target is the final size, so anticipate the scaling + scale = min(rect_pre.width() / target_size.width(), + rect_pre.height() / target_size.height()) + target = target_size.toSizeF() * scale + target = QSize(max(bucket_res, ceil(target.width())), + max(bucket_res, ceil(target.height()))) + rect = change_rect_to_match_size(self.rect(), + MarkingItem.handle_selected, + target) + else: + rect = change_rect(self.rect(), + MarkingItem.handle_selected, + event.pos()) + + round_tl = round + round_br = round + if self.rect_type == ImageMarking.EXCLUDE: + round_tl = floor + round_br = ceil + elif self.rect_type == ImageMarking.INCLUDE: + round_tl = ceil + round_br = floor + rect = QRectF(grid.snap(rect.toRect().topLeft(), round_tl), + grid.snap(rect.toRect().bottomRight(), round_br)) + rect = QRect(QPoint(round_br(rect.topLeft().x()), + round_br(rect.topLeft().y())), + QPoint(round_tl(rect.bottomRight().x()), + round_tl(rect.bottomRight().y()))) + else: + pos_quantized = event.pos().toPoint() + rect = change_rect(self.rect().toRect(), + MarkingItem.handle_selected, + pos_quantized) + + MarkingItem.handle_selected = flip_rect_position(self.handle_selected, + rect.width() < 0, + rect.height() < 0) + + if rect.width() == 0 or rect.height() == 0: + self.setRect(rect) + else: + rect = rect.intersected(self.image_size) + self.setRect(rect) + self.size_changed() + + self.move() + super().mouseMoveEvent(event) + + def mouseReleaseEvent(self, event): + MarkingItem.handle_selected = RectPosition.NONE + self.move() + self.setZValue(2) + if ((event.modifiers() & Qt.KeyboardModifier.ControlModifier) == + Qt.KeyboardModifier.ControlModifier): + self.image_view.set_insertion_mode(self.rect_type) + self.ungrabMouse() + super().mouseReleaseEvent(event) + self.image_view.image_viewer.marking_changed(self) + + def paint(self, painter, option, widget=None): + if self.rect_type == ImageMarking.CROP: + if (self.show_crop_hint and + MarkingItem.handle_selected != RectPosition.NONE and + self==self.scene().mouseGrabberItem()): + hint_line_crossings = [ + self.rect().center(), + self.rect().topLeft() + QPointF(self.rect().width()*golden_ratio, + self.rect().height()*golden_ratio), + self.rect().bottomRight() - QPointF(self.rect().width()*golden_ratio, + self.rect().height()*golden_ratio), + self.rect().topLeft() + QPointF(self.rect().width()/3, + self.rect().height()/3), + self.rect().bottomRight() - QPointF(self.rect().width()/3, + self.rect().height()/3)] + lint_line_style = [Qt.SolidLine, Qt.DotLine, Qt.DotLine, Qt.DashLine, Qt.DashLine] + for crossing, style in zip(hint_line_crossings, lint_line_style): + path = QPainterPath() + path.moveTo(self.rect().x(), crossing.y()) + path.lineTo(self.rect().right(), crossing.y()) + path.moveTo(crossing.x(), self.rect().y()) + path.lineTo(crossing.x(), self.rect().bottom()) + painter.setPen(QPen(QColor(255, 255, 255, 127), 3 / self.zoom_factor)) + painter.drawPath(path) + painter.setPen(QPen(QColor(0, 0, 0), 1 / self.zoom_factor, style)) + painter.drawPath(path) + painter.setPen(Qt.NoPen) + painter.setBrush(QColor(255, 0, 0, 127)) + path = QPainterPath() + path.addRect(self.rect()) + path.addRect(grid.visible) + painter.drawPath(path) + + pen_half_width = self.pen_half_width / self.zoom_factor + pen = QPen(self.color, 2*pen_half_width, Qt.SolidLine, Qt.RoundCap, + Qt.RoundJoin) + painter.setPen(pen) + painter.setBrush(Qt.NoBrush) + painter.drawRect(self.rect().adjusted(-pen_half_width, -pen_half_width, + pen_half_width, pen_half_width)) + + if self.isSelected(): + s_rect = self.rect().adjusted(-2*pen_half_width, -2*pen_half_width, + 2*pen_half_width, 2*pen_half_width) + painter.setPen(QPen(Qt.white, 1.5 / self.zoom_factor, Qt.SolidLine)) + painter.drawRect(s_rect) + painter.setPen(QPen(Qt.black, 1.5 / self.zoom_factor, Qt.DotLine)) + painter.drawRect(s_rect) + + def shape(self): + path = super().shape() + adjust = (self.pen_half_width + max(self.pen_half_width, + self.handle_half_size))/self.zoom_factor + path.addRect(self.rect().adjusted(-adjust, -adjust, adjust, adjust)) + return path + + def boundingRect(self): + adjust = (self.pen_half_width + max(self.pen_half_width, + self.handle_half_size))/self.zoom_factor + bbox = self.rect().adjusted(-adjust, -adjust, adjust, adjust) + return bbox + + def size_changed(self): + if self.rect_type == ImageMarking.CROP: + old_grid = grid + calculate_grid(self.rect().toRect()) + if old_grid != grid: + self.image_view.image_viewer.recalculate_markings(self) + self.adjust_layout() + + def adjust_layout(self): + if self.label is not None: + self.label.changeZoom(self.zoom_factor) + pen_half_width = self.pen_half_width / self.zoom_factor + if self.rect().y() > self.label.boundingRect().height(): + self.label.setPos(self.rect().adjusted( + -2 * pen_half_width, + -1.8*pen_half_width + - self.label.boundingRect().height() / self.zoom_factor, + 0, 0).topLeft()) + self.label.parentItem().setRect(self.label.sceneBoundingRect()) + else: + self.label.setPos(self.rect().adjusted( + -pen_half_width, -pen_half_width, 0, 0).topLeft()) + self.label.parentItem().setRect(self.label.sceneBoundingRect()) + + +class MarkingLabel(QGraphicsTextItem): + editingFinished = Signal() + + def __init__(self, text, confidence, parent): + if 0 <= confidence < 1: + super().__init__(f'{text}: {confidence:.3f}', parent) + else: + super().__init__(text, parent) + self.setDefaultTextColor(Qt.black) + self.setTextInteractionFlags(Qt.TextEditorInteraction) + + def focusOutEvent(self, event): + super().focusOutEvent(event) + self.editingFinished.emit() + + def keyPressEvent(self, event): + if event.key() in (Qt.Key_Enter, Qt.Key_Return): + self.clearFocus() + self.editingFinished.emit() + else: + super().keyPressEvent(event) + self.parentItem().setRect(self.sceneBoundingRect()) + + def insertFromMimeData(self, source): + if source.hasText(): + # Insert only the plain text + cursor = self.textCursor() + cursor.insertText(source.text()) + else: + super().insertFromMimeData(source) + self.parentItem().setRect(self.sceneBoundingRect()) + + def changeZoom(self, zoom_factor): + self.setScale(1/zoom_factor) + self.parentItem().setRect(self.sceneBoundingRect()) + + +class ResizeHintHUD(QGraphicsItem): + zoom_factor = 1.0 + + def __init__(self, boundingRect: QRect, parent=None): + super().__init__(parent) + self._boundingRect = boundingRect + self.rect = QRectF(0, 0, 1, 1) + self.path_ar = QPainterPath() + self.path_size = QPainterPath() + self.setCacheMode(QGraphicsItem.DeviceCoordinateCache) + self.setZValue(3) + self.last_point: QPointF | float = QPointF(-1, -1) + self.last_pos = RectPosition.NONE + + @Slot(QRectF, RectPosition) + def setValues(self, rect: QRectF, pos: RectPosition): + if self.rect == rect and self.isVisible() == (pos != RectPosition.NONE): + return + + self.rect = rect + self.setVisible(pos != RectPosition.NONE) + pos_change = self.last_pos != pos + self.last_pos = pos + + self.path_ar = QPainterPath() + self.path_size = QPainterPath() + do_update = False + + if pos == RectPosition.TL: + do_update = self.add_hyperbola_limit(self.rect.bottomRight(), -1, -1, pos_change) + elif pos == RectPosition.TOP: + do_update = self.add_line_limit_lr(self.rect.bottom(), -1, pos_change) + elif pos == RectPosition.TR: + do_update = self.add_hyperbola_limit(self.rect.bottomLeft(), 1, -1, pos_change) + elif pos == RectPosition.RIGHT: + do_update = self.add_line_limit_td(self.rect.x(), 1, pos_change) + elif pos == RectPosition.BR: + do_update = self.add_hyperbola_limit(self.rect.topLeft(), 1, 1, pos_change) + elif pos == RectPosition.BOTTOM: + do_update = self.add_line_limit_lr(self.rect.y(), 1, pos_change) + elif pos == RectPosition.BL: + do_update = self.add_hyperbola_limit(self.rect.topRight(), -1, 1, pos_change) + elif pos == RectPosition.LEFT: + do_update = self.add_line_limit_td(self.rect.right(), -1, pos_change) + + if do_update: + self.update() + + def add_line_limit_td(self, x: float, lr: int, pos_change: bool) -> bool: + if self.last_point == x and not pos_change: + return False + width = settings.value('export_resolution', type=int)**2 / self.rect.height() + res_size = max(settings.value('export_bucket_res_size', type=int), 1) + self.path_size.moveTo(x + lr * width, self.rect.y() ) + self.path_size.lineTo(x + lr * width, self.rect.y() + self.rect.height()) + + for ar in target_dimension.get_preferred_sizes(): + s = max(res_size / ar[0], res_size / ar[1]) + f = max(self._boundingRect.width() / ar[0], + self._boundingRect.height() / ar[1], 2) + self.path_ar.moveTo(x + lr * ar[0] * s, self.rect.y() + ar[1] * s) + self.path_ar.lineTo(x + lr * ar[0] * f, self.rect.y() + ar[1] * f) + self.path_ar.moveTo(x + lr * ar[0] * s, self.rect.bottom() - ar[1] * s) + self.path_ar.lineTo(x + lr * ar[0] * f, self.rect.bottom() - ar[1] * f) + self.last_pos = x + return True + + def add_line_limit_lr(self, y: float, td: int, pos_change: bool) -> bool: + if self.last_point == y and not pos_change: + return False + height = settings.value('export_resolution', type=int)**2 / self.rect.width() + res_size = max(settings.value('export_bucket_res_size', type=int), 1) + self.path_size.moveTo(self.rect.x(), y + td * height) + self.path_size.lineTo(self.rect.x() + self.rect.width(), y + td * height) + + for ar in target_dimension.get_preferred_sizes(): + s = max(res_size / ar[0], res_size / ar[1]) + f = max(self._boundingRect.width() / ar[0], + self._boundingRect.height() / ar[1], 2) + self.path_ar.moveTo(self.rect.x() + ar[0] * s, y + td * ar[1] * s) + self.path_ar.lineTo(self.rect.x() + ar[0] * f, y + td * ar[1] * f) + self.path_ar.moveTo(self.rect.right() - ar[0] * s, y + td * ar[1] * s) + self.path_ar.lineTo(self.rect.right() - ar[0] * f, y + td * ar[1] * f) + self.last_pos = y + return True + + def add_hyperbola_limit(self, pos: QPointF, lr: int, td: int, pos_change: bool) -> bool: + if self.last_point == pos and not pos_change: + return False + target_area = max(settings.value('export_resolution', type=int)**2, 1) + res_size = max(settings.value('export_bucket_res_size', type=int), 1) + if td < 0: + distance_x = target_area / (pos.y() - self._boundingRect.y()) + else: + distance_x = target_area / (self._boundingRect.bottom() - pos.y()) + x = self._boundingRect.x() if lr < 0 else pos.x() + distance_x + end_x = pos.x() - distance_x if lr < 0 else self._boundingRect.right() + first = True + while x < end_x + 50: + p = QPointF(x, pos.y() + td * target_area / (lr * (x - pos.x()))) + self.path_size.moveTo(p) if first else self.path_size.lineTo(p) + first = False + x += 50 + + for ar in target_dimension.get_preferred_sizes(): + s = max(res_size / ar[0], res_size / ar[1]) + f = max(self._boundingRect.width() / ar[0], + self._boundingRect.height() / ar[1], 2) + self.path_ar.moveTo(pos.x() + lr * ar[0] * s, pos.y() + td * ar[1] * s) + self.path_ar.lineTo(pos.x() + lr * ar[0] * f, pos.y() + td * ar[1] * f) + self.last_pos = pos + return True + + def boundingRect(self): + return self._boundingRect + + def paint(self, painter, option, widget=None): + clip_path = QPainterPath() + clip_path.addRect(self._boundingRect) + painter.setClipPath(clip_path) + pen = QPen(QColor(255, 255, 255, 127), 3 / self.zoom_factor) + painter.setPen(pen) + painter.drawPath(self.path_size) + painter.drawPath(self.path_ar) + pen = QPen(QColor(0, 255, 0), 1 / self.zoom_factor) + painter.setPen(pen) + painter.drawPath(self.path_size) + pen = QPen(QColor(0, 0, 0), 1 / self.zoom_factor) + painter.setPen(pen) + painter.drawPath(self.path_ar) + +class ImageGraphicsView(QGraphicsView): + def __init__(self, scene, image_viewer): + super().__init__(scene) + self.setContextMenuPolicy(Qt.CustomContextMenu) + self.customContextMenuRequested.connect(self.showContextMenu) + self.setRenderHint(QPainter.Antialiasing) + self.setDragMode(QGraphicsView.DragMode.ScrollHandDrag) + self.setTransformationAnchor(QGraphicsView.ViewportAnchor.AnchorUnderMouse) + self.setResizeAnchor(QGraphicsView.ViewportAnchor.AnchorUnderMouse) + self.image_viewer = image_viewer + MarkingItem.image_view = self + self.setMouseTracking(True) + self.last_pos = None + self.clear_scene() + + def showContextMenu(self, pos): + scene_pos = self.mapToScene(pos) + item = self.scene().itemAt(scene_pos, self.transform()) + if isinstance(item, MarkingLabel): + item = item.parentItem().parentItem() + if isinstance(item, MarkingItem) and MarkingItem.handle_selected != RectPosition.NONE: + menu = QMenu() + if item.rect_type != ImageMarking.NONE: + if item.rect_type != ImageMarking.CROP: + marking_group = QActionGroup(menu) + change_to_hint_action = QAction('Hint', marking_group) + change_to_hint_action.setCheckable(True) + change_to_hint_action.setChecked(item.rect_type == ImageMarking.HINT) + change_to_hint_action.triggered.connect( + lambda: self.image_viewer.change_marking([item], ImageMarking.HINT)) + menu.addAction(change_to_hint_action) + change_to_exclude_action = QAction('Exclude', marking_group) + change_to_exclude_action.setCheckable(True) + change_to_exclude_action.setChecked(item.rect_type == ImageMarking.EXCLUDE) + change_to_exclude_action.triggered.connect( + lambda: self.image_viewer.change_marking([item], ImageMarking.EXCLUDE)) + menu.addAction(change_to_exclude_action) + change_to_include_action = QAction('Include', marking_group) + change_to_include_action.setCheckable(True) + change_to_include_action.setChecked(item.rect_type == ImageMarking.INCLUDE) + change_to_include_action.triggered.connect( + lambda: self.image_viewer.change_marking([item], ImageMarking.INCLUDE)) + menu.addAction(change_to_include_action) + menu.addSeparator() + delete_marking_action = QAction( + QIcon.fromTheme('edit-delete'), 'Delete', self) + delete_marking_action.triggered.connect( + lambda: self.image_viewer.delete_markings([item])) + menu.addAction(delete_marking_action) + menu.exec(self.mapToGlobal(pos)) + + def clear_scene(self): + """Use this and not scene().clear() due to resource management.""" + self.insertion_mode = False + self.horizontal_line = None + self.vertical_line = None + self.scene().clear() + + def set_insertion_mode(self, marking: ImageMarking): + old_insertion_mode = self.insertion_mode + self.insertion_mode = marking != ImageMarking.NONE + if self.insertion_mode: + if not old_insertion_mode: + self.setDragMode(QGraphicsView.DragMode.NoDrag) + self.horizontal_line = QGraphicsLineItem() + self.horizontal_line.setZValue(5) + self.vertical_line = QGraphicsLineItem() + self.vertical_line.setZValue(5) + self.scene().addItem(self.horizontal_line) + self.scene().addItem(self.vertical_line) + self.update_lines_pos() + self.image_viewer.marking.emit(marking) + else: + self.setDragMode(QGraphicsView.DragMode.ScrollHandDrag) + if self.horizontal_line: + self.scene().removeItem(self.horizontal_line) + self.horizontal_line = None + self.scene().removeItem(self.vertical_line) + self.vertical_line = None + self.image_viewer.marking.emit(ImageMarking.NONE) + + def update_lines_pos(self): + """Show the hint lines at the position self.last_pos. + + Note: do not use a position parameter as then the key event couldn't + immediately show them as the mouse position would be missing then. + """ + if self.last_pos: + view_rect = self.mapToScene(self.viewport().rect()).boundingRect() + self.horizontal_line.setLine(view_rect.left(), self.last_pos.y(), + view_rect.right(), self.last_pos.y()) + self.vertical_line.setLine(self.last_pos.x(), view_rect.top(), + self.last_pos.x(), view_rect.bottom()) + + def mousePressEvent(self, event: QMouseEvent): + if self.insertion_mode and event.button() == Qt.MouseButton.LeftButton: + rect_type = self.image_viewer.marking_to_add + if rect_type == ImageMarking.NONE: + if ((event.modifiers() & Qt.KeyboardModifier.AltModifier) == + Qt.KeyboardModifier.AltModifier): + rect_type = ImageMarking.EXCLUDE + else: + rect_type = ImageMarking.HINT + + self.image_viewer.proxy_image_index.model().sourceModel().add_to_undo_stack( + action_name=f'Add {rect_type.value}', should_ask_for_confirmation=False) + + self.image_viewer.add_rectangle(QRect(self.last_pos, QSize(0, 0)), + rect_type, interactive=True) + self.set_insertion_mode(ImageMarking.NONE) + self.setDragMode(QGraphicsView.DragMode.NoDrag) + return + super().mousePressEvent(event) + + def mouseMoveEvent(self, event: QMouseEvent): + scene_pos = self.mapToScene(event.position().toPoint()) + items = self.scene().items(scene_pos) + cursor = None + + if self.insertion_mode: + cursor = Qt.CursorShape.CrossCursor + elif MarkingItem.handle_selected != RectPosition.NONE: + cursor = map_rect_position_to_cursor(MarkingItem.handle_selected) + else: + for item in items: + if isinstance(item, MarkingItem): + handle = item.handleAt(scene_pos) + if handle == RectPosition.NONE: + continue + cursor = map_rect_position_to_cursor(handle) + break + if cursor is None: + self.setDragMode(QGraphicsView.DragMode.ScrollHandDrag) + else: + self.setDragMode(QGraphicsView.DragMode.NoDrag) + self.setCursor(cursor) + + if ((event.modifiers() & Qt.KeyboardModifier.ShiftModifier) == + Qt.KeyboardModifier.ShiftModifier): + self.last_pos = grid.snap(scene_pos.toPoint()).toPoint() + else: + self.last_pos = scene_pos.toPoint() + + if self.insertion_mode: + self.update_lines_pos() + else: + super().mouseMoveEvent(event) + + def keyPressEvent(self, event): + if event.key() == Qt.Key.Key_Delete: + edited_item = self.scene().focusItem() + if not (isinstance(edited_item, MarkingLabel) and + edited_item.textInteractionFlags() == Qt.TextEditorInteraction): + # Delete marking only when not editing the label + self.image_viewer.delete_markings() + else: + if MarkingItem.handle_selected == RectPosition.NONE: + if ((event.modifiers() & Qt.KeyboardModifier.ControlModifier) == + Qt.KeyboardModifier.ControlModifier): + if ((event.modifiers() & Qt.KeyboardModifier.AltModifier) == + Qt.KeyboardModifier.AltModifier): + self.set_insertion_mode(ImageMarking.EXCLUDE) + else: + self.set_insertion_mode(ImageMarking.HINT) + super().keyPressEvent(event) + + def keyReleaseEvent(self, event): + if MarkingItem.handle_selected == RectPosition.NONE: + if ((event.modifiers() & Qt.KeyboardModifier.ControlModifier) == + Qt.KeyboardModifier.ControlModifier): + if ((event.modifiers() & Qt.KeyboardModifier.AltModifier) == + Qt.KeyboardModifier.AltModifier): + self.set_insertion_mode(ImageMarking.EXCLUDE) + else: + self.set_insertion_mode(ImageMarking.HINT) + else: + self.set_insertion_mode(ImageMarking.NONE) + super().keyReleaseEvent(event) + + def resizeEvent(self, event): + super().resizeEvent(event) + if self.image_viewer.is_zoom_to_fit: + self.image_viewer.zoom_fit() class ImageViewer(QWidget): + zoom = Signal(float, name='zoomChanged') + marking = Signal(ImageMarking, name='markingToAdd') + accept_crop_addition = Signal(bool, name='allowAdditionOfCrop') + crop_changed = Signal(Grid, name='cropChanged') + rating_changed = Signal(float, name='ratingChanged') + def __init__(self, proxy_image_list_model: ProxyImageListModel): super().__init__() + self.inhibit_reload_image = False self.proxy_image_list_model = proxy_image_list_model - self.image_label = ImageLabel() - QVBoxLayout(self).addWidget(self.image_label) + MarkingItem.pen_half_width = round(self.devicePixelRatio()) + MarkingItem.zoom_factor = 1.0 + self.is_zoom_to_fit = True + self.show_marking_state = True + self.show_label_state = True + self.show_marking_latent_state = True + self.marking_to_add = ImageMarking.NONE + self.scene = QGraphicsScene() + self.view = ImageGraphicsView(self.scene, self) + self.view.setOptimizationFlags(QGraphicsView.DontSavePainterState) + self.crop_marking: ImageMarking | None = None + settings.change.connect(self.setting_change) + + layout = QVBoxLayout() + layout.addWidget(self.view) + self.setLayout(layout) + + self.proxy_image_index: QPersistentModelIndex = None + self.marking_items: list[MarkingItem] = [] + + self.view.wheelEvent = self.wheelEvent + + @Slot() + def load_image(self, proxy_image_index: QModelIndex, is_complete = True): + persistent_image_index = QPersistentModelIndex(proxy_image_index) + if ((not persistent_image_index.isValid()) or + (self.inhibit_reload_image and + persistent_image_index == self.proxy_image_index)): + return + self.proxy_image_index = persistent_image_index + + image: Image = self.proxy_image_index.data(Qt.ItemDataRole.UserRole) + self.rating_changed.emit(image.rating) + + if is_complete: + self.marking_items.clear() + self.view.clear_scene() + pixmap = QPixmap(str(image.path)) + image_item = QGraphicsPixmapItem(pixmap) + image_item.setZValue(0) + self.scene.setSceneRect(image_item.boundingRect() + .adjusted(-1, -1, 1, 1)) # space for rect border + self.scene.addItem(image_item) + MarkingItem.image_size = image_item.boundingRect().toRect() + self.zoom_fit() + + self.hud_item = ResizeHintHUD(MarkingItem.image_size, image_item) + else: + for item in self.marking_items: + self.scene.removeItem(item) + self.marking_items.clear() + + self.marking_to_add = ImageMarking.NONE + self.marking.emit(ImageMarking.NONE) + self.accept_crop_addition.emit(image.crop is None) + if image.crop is not None: + self.add_rectangle(image.crop, ImageMarking.CROP, interactive=False) + else: + calculate_grid(MarkingItem.image_size) + for marking in image.markings: + self.add_rectangle(marking.rect, marking.type, interactive=False, + name=marking.label, confidence=marking.confidence) + + def rating_change(self, rating: float): + if self.proxy_image_index.isValid(): + image: Image = self.proxy_image_index.data(Qt.ItemDataRole.UserRole) + if image.rating != rating: + image.rating = rating + self.proxy_image_list_model.sourceModel().write_meta_to_disk(image) + + @Slot() + def setting_change(self, key, value): + if key in ['export_resolution', 'export_bucket_res_size', + 'export_latent_size', 'export_upscaling', + 'export_bucket_strategy']: + self.recalculate_markings() + + def recalculate_markings(self, ignore: MarkingItem | None = None): + if self.crop_marking: + calculate_grid(self.crop_marking.rect().toRect()) + if MarkingItem.handle_selected != RectPosition.NONE: + # currently editing the crop marking -> update display + self.crop_changed.emit(grid) + else: + calculate_grid(MarkingItem.image_size) + for marking in self.marking_items: + if marking != ignore: + marking.size_changed() + self.scene.invalidate() + + @Slot() + def zoom_in(self, center_pos: QPoint = None): + MarkingItem.zoom_factor = min(MarkingItem.zoom_factor * 1.25, 16) + self.is_zoom_to_fit = False + self.zoom_emit() + + @Slot() + def zoom_out(self, center_pos: QPoint = None): + view = self.view.viewport().size() + scene = self.scene.sceneRect() + if scene.width() < 1 or scene.height() < 1: + return + limit = min(view.width()/scene.width(), view.height()/scene.height()) + MarkingItem.zoom_factor = max(MarkingItem.zoom_factor / 1.25, limit) + self.is_zoom_to_fit = MarkingItem.zoom_factor == limit + self.zoom_emit() + + @Slot() + def zoom_original(self): + MarkingItem.zoom_factor = 1.0 + self.is_zoom_to_fit = False + self.zoom_emit() + + @Slot() + def zoom_fit(self): + self.view.fitInView(self.scene.sceneRect(), Qt.KeepAspectRatio) + MarkingItem.zoom_factor = self.view.transform().m11() + self.is_zoom_to_fit = True + self.zoom_emit() + + def zoom_emit(self): + ResizeHintHUD.zoom_factor = MarkingItem.zoom_factor + transform = self.view.transform() + self.view.setTransform(QTransform( + MarkingItem.zoom_factor, transform.m12(), transform.m13(), + transform.m21(), MarkingItem.zoom_factor, transform.m23(), + transform.m31(), transform.m32(), transform.m33())) + for marking in self.marking_items: + marking.adjust_layout() + if self.is_zoom_to_fit: + self.zoom.emit(-1) + else: + self.zoom.emit(MarkingItem.zoom_factor) + + @Slot(ImageMarking) + def add_marking(self, marking: ImageMarking): + self.marking_to_add = marking + self.view.set_insertion_mode(marking) + + @Slot() + def change_marking(self, items: list[MarkingItem] | None = None, + new_marking: ImageMarking = ImageMarking.NONE): + self.proxy_image_index.model().sourceModel().add_to_undo_stack( + action_name=f'Change marking', should_ask_for_confirmation=False) + if items is None: + items = self.scene.selectedItems() + for item in items: + if new_marking == ImageMarking.NONE: + # default: toggle between all types + item.rect_type = {ImageMarking.HINT: ImageMarking.EXCLUDE, + ImageMarking.INCLUDE: ImageMarking.HINT, + ImageMarking.EXCLUDE: ImageMarking.INCLUDE + }[item.rect_type] + else: + item.rect_type = new_marking + item.color = marking_colors[item.rect_type] + item.label.parentItem().setBrush(item.color) + self.marking_changed(item) + item.update() + + @Slot(bool) + def show_marking(self, checked: bool): + self.show_marking_state = checked + for marking in self.marking_items: + marking.setVisible(checked) + + @Slot(bool) + def show_label(self, checked: bool): + self.show_label_state = checked + for marking in self.marking_items: + if marking.label: + marking.label.setVisible(checked) + marking.label.parentItem().setVisible(checked) + + @Slot(bool) + def show_marking_latent(self, checked: bool): + MarkingItem.show_marking_latent = checked + for marking in self.marking_items: + if marking.rect_type in [ImageMarking.INCLUDE, ImageMarking.EXCLUDE]: + marking.area.setVisible(checked) + + def wheelEvent(self, event): + old_pos = self.view.mapToScene(event.position().toPoint()) + + if event.angleDelta().y() > 0: + self.zoom_in() + elif event.angleDelta().y() < 0: + self.zoom_out() + else: + return + + new_pos = self.view.mapToScene(event.position().toPoint()) + delta = new_pos - old_pos + self.view.translate(delta.x(), delta.y()) + + def add_rectangle(self, rect: QRect, rect_type: ImageMarking, + interactive: bool, size: QSize = None, name: str = '', + confidence: float = 1.0): + self.marking_to_add = ImageMarking.NONE + marking_item = MarkingItem(rect, rect_type, interactive, size) + marking_item.setVisible(self.show_marking_state) + if rect_type == ImageMarking.CROP: + self.crop_marking = marking_item + marking_item.size_changed() # call after self.crop_marking was set! + elif name == '' and rect_type != ImageMarking.NONE: + image: Image = self.proxy_image_index.data(Qt.ItemDataRole.UserRole) + name = {ImageMarking.HINT: 'hint', + ImageMarking.INCLUDE: 'include', + ImageMarking.EXCLUDE: 'exclude'}[rect_type] + image.markings.append(Marking(name, rect_type, rect, confidence)) + marking_item.setData(0, name) + marking_item.setData(1, confidence) + if rect_type != ImageMarking.CROP and rect_type != ImageMarking.NONE: + label_background = QGraphicsRectItem(marking_item) + label_background.setZValue(2) + label_background.setBrush(marking_item.color) + label_background.setPen(Qt.NoPen) + label_background.setVisible(self.show_label_state) + marking_item.label = MarkingLabel(name, confidence, label_background) + marking_item.label.setZValue(2) + marking_item.label.setVisible(self.show_label_state) + marking_item.label.editingFinished.connect(self.label_changed) + marking_item.adjust_layout() + self.scene.addItem(marking_item) + self.marking_items.append(marking_item) + if interactive: + self.scene.clearSelection() + marking_item.grabMouse() + if rect_type == ImageMarking.CROP: + self.accept_crop_addition.emit(False) + + @Slot() + def label_changed(self): + """Slot to call when a marking label was changed to sync the information + in the image.""" + self.proxy_image_index.model().sourceModel().add_to_undo_stack( + action_name=f'Change label', should_ask_for_confirmation=False) + image: Image = self.proxy_image_index.data(Qt.ItemDataRole.UserRole) + image.markings.clear() + for marking in self.marking_items: + if marking.rect_type != ImageMarking.CROP: + label = marking.label.toPlainText() + match = re.match(r'^(.*):\s*(\d*\.\d+)$', label) + if match: + label = match.group(1) + confidence = float(match.group(2)) + else: + confidence = 1.0 + marking.label.parentItem().parentItem().setData(0, label) + marking.label.parentItem().parentItem().setData(1, confidence) + image.markings.append(Marking(label=label, + type=marking.rect_type, + rect=marking.rect().toRect(), + confidence=confidence)) + self.proxy_image_list_model.sourceModel().write_meta_to_disk(image) + + @Slot(QGraphicsRectItem) + def marking_changed(self, marking: QGraphicsRectItem): + """Slot to call when a marking was changed to sync the information + in the image.""" + assert self.proxy_image_index != None + assert self.proxy_image_index.isValid() + image: Image = self.proxy_image_index.data(Qt.ItemDataRole.UserRole) + + if marking.rect_type == ImageMarking.CROP: + self.inhibit_reload_image = True + self.proxy_image_list_model.sourceModel().layoutAboutToBeChanged.emit() + image.thumbnail = None + image.crop = marking.rect().toRect() # ensure int! + image.target_dimension = grid.target + if not self.proxy_image_list_model.does_image_match_filter( + image, self.proxy_image_list_model.filter): + # don't call .invalidate() as the displayed list shouldn't + # update + self.proxy_image_list_model.filter = [['path', str(image.path)], + 'OR', + self.proxy_image_list_model.filter] + self.crop_changed.emit(None) + self.proxy_image_list_model.sourceModel().changePersistentIndex( + self.proxy_image_index, self.proxy_image_index) + + self.proxy_image_list_model.sourceModel().dataChanged.emit( + self.proxy_image_index, self.proxy_image_index, + [Qt.ItemDataRole.DecorationRole, Qt.ItemDataRole.SizeHintRole, + Qt.ToolTipRole, Qt.ItemDataRole.UserRole]) + self.proxy_image_list_model.sourceModel().layoutChanged.emit() + self.inhibit_reload_image = False + else: + image.markings = [Marking(m.data(0), + m.rect_type, + m.rect().toRect(), + m.data(1)) + for m in self.marking_items if m.rect_type != ImageMarking.CROP] + self.proxy_image_list_model.sourceModel().write_meta_to_disk(image) + + def get_selected_type(self) -> ImageMarking: + if len(self.scene.selectedItems()) > 0: + return self.scene.selectedItems()[0].rect_type + return ImageMarking.NONE @Slot() - def load_image(self, proxy_image_index: QModelIndex): - image: Image = self.proxy_image_list_model.data( - proxy_image_index, Qt.ItemDataRole.UserRole) - self.image_label.load_image(image.path) + def delete_markings(self, items: list[MarkingItem] | None = None): + """Slot to delete the list of items or when items = None all currently + selected marking items.""" + self.proxy_image_index.model().sourceModel().add_to_undo_stack( + action_name=f'Delete marking', should_ask_for_confirmation=False) + image: Image = self.proxy_image_index.data(Qt.ItemDataRole.UserRole) + if items is None: + items = self.scene.selectedItems() + for item in items: + if item.rect_type == ImageMarking.CROP: + self.crop_marking = None + image.thumbnail = None + image.crop = None + image.target_dimension = None + self.accept_crop_addition.emit(True) + calculate_grid(MarkingItem.image_size) + self.proxy_image_list_model.sourceModel().dataChanged.emit( + self.proxy_image_index, self.proxy_image_index, + [Qt.ItemDataRole.DecorationRole, Qt.ItemDataRole.SizeHintRole, + Qt.ToolTipRole, Qt.ItemDataRole.UserRole]) + else: + self.marking_items.remove(item) + self.label_changed() + self.proxy_image_list_model.sourceModel().write_meta_to_disk(image) + self.scene.removeItem(item) diff --git a/taggui/widgets/main_window.py b/taggui/widgets/main_window.py index 78a8cf98..1bb1dc42 100644 --- a/taggui/widgets/main_window.py +++ b/taggui/widgets/main_window.py @@ -1,33 +1,39 @@ from pathlib import Path from PySide6.QtCore import QKeyCombination, QModelIndex, QUrl, Qt, Slot -from PySide6.QtGui import (QAction, QCloseEvent, QDesktopServices, QIcon, - QKeySequence, QPixmap, QShortcut) +from PySide6.QtGui import (QAction, QActionGroup, QCloseEvent, QDesktopServices, + QIcon, QKeySequence, QShortcut, QMouseEvent) from PySide6.QtWidgets import (QApplication, QFileDialog, QMainWindow, - QMessageBox, QStackedWidget, QVBoxLayout, - QWidget) + QMessageBox, QStackedWidget, QToolBar, + QVBoxLayout, QWidget, QSizePolicy, QHBoxLayout, + QLabel) + from transformers import AutoTokenizer 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 from models.proxy_image_list_model import ProxyImageListModel from models.tag_counter_model import TagCounterModel +from utils.icons import (taggui_icon, create_add_box_icon, toggle_marking_icon, + show_markings_icon, show_labels_icon, + show_marking_latent_icon) from utils.big_widgets import BigPushButton from utils.image import Image from utils.key_press_forwarder import KeyPressForwarder -from utils.settings import DEFAULT_SETTINGS, get_settings, get_tag_separator +from utils.settings import DEFAULT_SETTINGS, settings, get_tag_separator from utils.shortcut_remover import ShortcutRemover from utils.utils import get_resource_path, pluralize from widgets.all_tags_editor import AllTagsEditor from widgets.auto_captioner import AutoCaptioner +from widgets.auto_markings import AutoMarkings from widgets.image_list import ImageList from widgets.image_tags_editor import ImageTagsEditor -from widgets.image_viewer import ImageViewer +from widgets.image_viewer import ImageViewer, ImageMarking -ICON_PATH = Path('images/icon.ico') GITHUB_REPOSITORY_URL = 'https://github.com/jhc13/taggui' TOKENIZER_DIRECTORY_PATH = Path('clip-vit-base-patch32') @@ -36,11 +42,12 @@ class MainWindow(QMainWindow): def __init__(self, app: QApplication): super().__init__() self.app = app - self.settings = get_settings() # The path of the currently loaded directory. This is set later when a # directory is loaded. self.directory_path = None - image_list_image_width = self.settings.value( + self.is_running = True + app.aboutToQuit.connect(lambda: setattr(self, 'is_running', False)) + image_list_image_width = settings.value( 'image_list_image_width', defaultValue=DEFAULT_SETTINGS['image_list_image_width'], type=int) tag_separator = get_tag_separator() @@ -55,7 +62,7 @@ def __init__(self, app: QApplication): self.tag_counter_model = TagCounterModel() self.image_tag_list_model = ImageTagListModel() - self.setWindowIcon(QIcon(QPixmap(get_resource_path(ICON_PATH)))) + self.setWindowIcon(taggui_icon()) # Not setting this results in some ugly colors. self.setPalette(self.app.style().standardPalette()) # The font size must be set before creating the widgets to ensure that @@ -63,6 +70,94 @@ def __init__(self, app: QApplication): self.set_font_size() self.image_viewer = ImageViewer(self.proxy_image_list_model) self.create_central_widget() + + self.toolbar = QToolBar('Main toolbar', self) + self.toolbar.setObjectName('Main toolbar') + self.toolbar.setFloatable(True) + self.addToolBar(self.toolbar) + self.zoom_fit_best_action = QAction(QIcon.fromTheme('zoom-fit-best'), + 'Zoom to fit', self) + self.zoom_fit_best_action.setCheckable(True) + self.toolbar.addAction(self.zoom_fit_best_action) + self.zoom_in_action = QAction(QIcon.fromTheme('zoom-in'), + 'Zoom in', self) + self.toolbar.addAction(self.zoom_in_action) + self.zoom_original_action = QAction(QIcon.fromTheme('zoom-original'), + 'Original size', self) + self.zoom_original_action.setCheckable(True) + self.toolbar.addAction(self.zoom_original_action) + self.zoom_out_action = QAction(QIcon.fromTheme('zoom-out'), + 'Zoom out', self) + self.toolbar.addAction(self.zoom_out_action) + self.toolbar.addSeparator() + self.add_action_group = QActionGroup(self) + self.add_action_group.setExclusionPolicy(QActionGroup.ExclusiveOptional) + self.add_crop_action = QAction(create_add_box_icon(Qt.blue), + 'Add crop', self.add_action_group) + self.add_crop_action.setCheckable(True) + self.toolbar.addAction(self.add_crop_action) + self.add_hint_action = QAction(create_add_box_icon(Qt.gray), + 'Add hint', self.add_action_group) + self.add_hint_action.setCheckable(True) + self.toolbar.addAction(self.add_hint_action) + self.add_exclude_action = QAction(create_add_box_icon(Qt.red), + 'Add exclude mask', self.add_action_group) + self.add_exclude_action.setCheckable(True) + self.toolbar.addAction(self.add_exclude_action) + self.add_include_action = QAction(create_add_box_icon(Qt.green), + 'Add include mask', self.add_action_group) + self.add_include_action.setCheckable(True) + self.toolbar.addAction(self.add_include_action) + self.delete_marking_action = QAction(QIcon.fromTheme('edit-delete'), + 'Delete marking', self) + self.delete_marking_action.setEnabled(False) + self.toolbar.addAction(self.delete_marking_action) + self.add_toggle_marking_action = QAction(toggle_marking_icon(), + 'Change marking type', self) + self.add_toggle_marking_action.setEnabled(False) + self.toolbar.addAction(self.add_toggle_marking_action) + self.add_show_marking_action = QAction(show_markings_icon(), + 'Show markings', self) + self.add_show_marking_action.setCheckable(True) + self.add_show_marking_action.setChecked(True) + self.toolbar.addAction(self.add_show_marking_action) + self.add_show_labels_action = QAction(show_labels_icon(), + 'Show labels', self) + self.add_show_labels_action.setCheckable(True) + self.add_show_labels_action.setChecked(True) + self.toolbar.addAction(self.add_show_labels_action) + self.add_show_marking_latent_action = QAction(show_marking_latent_icon(), + 'Show marking in latent space', self) + self.add_show_marking_latent_action.setCheckable(True) + self.add_show_marking_latent_action.setChecked(True) + self.toolbar.addAction(self.add_show_marking_latent_action) + spacer = QWidget() + spacer.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Preferred) + self.toolbar.addWidget(spacer) + star_widget = QWidget() + star_layout = QHBoxLayout(star_widget) + star_layout.setContentsMargins(0, 0, 0, 0) + star_layout.setSpacing(0) + self.rating = 0 + self.star_labels = [] + for i in range(6): + shortcut = QShortcut(QKeySequence(f'Ctrl+{i}'), self) + shortcut.activated.connect(lambda checked=False, rating=i: + self.set_rating(2*rating, False)) + if i == 0: + continue + star_label = QLabel('☆', self) + star_label.setEnabled(False) + star_label.setAlignment(Qt.AlignCenter) + star_label.setStyleSheet('QLabel { font-size: 22px; }') + star_label.setToolTip(f'Ctrl+{i}') + star_label.mousePressEvent = lambda event, rating=i: ( + self.set_rating(rating/5.0, True, event)) + self.star_labels.append(star_label) + star_layout.addWidget(star_label) + self.image_viewer.rating_changed.connect(self.set_rating) + self.toolbar.addWidget(star_widget) + self.image_list = ImageList(self.proxy_image_list_model, tag_separator, image_list_image_width) self.addDockWidget(Qt.DockWidgetArea.LeftDockWidgetArea, @@ -82,7 +177,12 @@ def __init__(self, app: QApplication): self.image_list) self.addDockWidget(Qt.DockWidgetArea.RightDockWidgetArea, self.auto_captioner) + self.auto_markings = AutoMarkings(self.image_list_model, + self.image_list, self) + self.addDockWidget(Qt.DockWidgetArea.RightDockWidgetArea, + self.auto_markings) self.tabifyDockWidget(self.all_tags_editor, self.auto_captioner) + self.tabifyDockWidget(self.auto_captioner, self.auto_markings) self.all_tags_editor.raise_() # Set default widths for the dock widgets. # Temporarily set a size for the window so that the dock widgets can be @@ -101,22 +201,27 @@ def __init__(self, app: QApplication): self.reload_directory_action.setDisabled(True) self.undo_action = QAction('Undo', parent=self) self.redo_action = QAction('Redo', parent=self) + self.toggle_toolbar_action = QAction('Toolbar', parent=self) self.toggle_image_list_action = QAction('Images', parent=self) self.toggle_image_tags_editor_action = QAction('Image Tags', parent=self) self.toggle_all_tags_editor_action = QAction('All Tags', parent=self) self.toggle_auto_captioner_action = QAction('Auto-Captioner', parent=self) + self.toggle_auto_markings_action = QAction('Auto-Markings', + parent=self) self.create_menus() self.image_list_selection_model = (self.image_list.list_view .selectionModel()) self.image_list_model.image_list_selection_model = ( self.image_list_selection_model) + self.connect_toolbar_signals() self.connect_image_list_signals() self.connect_image_tags_editor_signals() self.connect_all_tags_editor_signals() self.connect_auto_captioner_signals() + self.connect_auto_markings_signals() # Forward any unhandled image changing key presses to the image list. key_press_forwarder = KeyPressForwarder( parent=self, target=self.image_list.list_view, @@ -182,13 +287,13 @@ def __init__(self, app: QApplication): def closeEvent(self, event: QCloseEvent): """Save the window geometry and state before closing.""" - self.settings.setValue('geometry', self.saveGeometry()) - self.settings.setValue('window_state', self.saveState()) + settings.setValue('geometry', self.saveGeometry()) + settings.setValue('window_state', self.saveState()) super().closeEvent(event) def set_font_size(self): font = self.app.font() - font_size = self.settings.value( + font_size = settings.value( 'font_size', defaultValue=DEFAULT_SETTINGS['font_size'], type=int) font.setPointSize(font_size) self.app.setFont(font) @@ -206,11 +311,23 @@ def create_central_widget(self): central_widget.addWidget(self.image_viewer) self.setCentralWidget(central_widget) + @Slot() + def zoom(self, factor): + if factor < 0: + self.zoom_fit_best_action.setChecked(True) + self.zoom_original_action.setChecked(False) + elif factor == 1.0: + self.zoom_fit_best_action.setChecked(False) + self.zoom_original_action.setChecked(True) + else: + self.zoom_fit_best_action.setChecked(False) + self.zoom_original_action.setChecked(False) + def load_directory(self, path: Path, select_index: int = 0, save_path_to_settings: bool = False): self.directory_path = path.resolve() if save_path_to_settings: - self.settings.setValue('directory_path', str(self.directory_path)) + settings.setValue('directory_path', str(self.directory_path)) self.setWindowTitle(path.name) self.image_list_model.load_directory(path) self.image_list.filter_line_edit.clear() @@ -245,7 +362,7 @@ def reload_directory(self): select_index_key = ('image_index' if self.proxy_image_list_model.filter is None else 'filtered_image_index') - select_index = self.settings.value(select_index_key, type=int) or 0 + select_index = settings.value(select_index_key, type=int) or 0 self.load_directory(self.directory_path) self.image_list.filter_line_edit.setText(filter_text) # If the selected image index is out of bounds due to images being @@ -255,6 +372,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 +436,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) @@ -355,10 +481,14 @@ def create_menus(self): edit_menu.addAction(remove_empty_tags_action) view_menu = menu_bar.addMenu('View') + self.toggle_toolbar_action.setCheckable(True) self.toggle_image_list_action.setCheckable(True) self.toggle_image_tags_editor_action.setCheckable(True) self.toggle_all_tags_editor_action.setCheckable(True) self.toggle_auto_captioner_action.setCheckable(True) + self.toggle_auto_markings_action.setCheckable(True) + self.toggle_toolbar_action.triggered.connect( + lambda is_checked: self.toolbar.setVisible(is_checked)) self.toggle_image_list_action.triggered.connect( lambda is_checked: self.image_list.setVisible(is_checked)) self.toggle_image_tags_editor_action.triggered.connect( @@ -367,10 +497,14 @@ def create_menus(self): lambda is_checked: self.all_tags_editor.setVisible(is_checked)) self.toggle_auto_captioner_action.triggered.connect( lambda is_checked: self.auto_captioner.setVisible(is_checked)) + self.toggle_auto_markings_action.triggered.connect( + lambda is_checked: self.auto_markings.setVisible(is_checked)) + view_menu.addAction(self.toggle_toolbar_action) view_menu.addAction(self.toggle_image_list_action) view_menu.addAction(self.toggle_image_tags_editor_action) view_menu.addAction(self.toggle_all_tags_editor_action) view_menu.addAction(self.toggle_auto_captioner_action) + view_menu.addAction(self.toggle_auto_markings_action) help_menu = menu_bar.addMenu('Help') open_github_repository_action = QAction('GitHub', parent=self) @@ -398,9 +532,8 @@ def update_undo_and_redo_actions(self): @Slot() def set_image_list_filter(self): filter_ = self.image_list.filter_line_edit.parse_filter_text() - self.proxy_image_list_model.filter = filter_ - # Apply the new filter. - self.proxy_image_list_model.invalidateFilter() + self.proxy_image_list_model.set_filter(filter_) + self.proxy_image_list_model.filter_changed.emit() if filter_ is None: all_tags_list_selection_model = (self.all_tags_editor .all_tags_list.selectionModel()) @@ -409,7 +542,7 @@ def set_image_list_filter(self): self.all_tags_editor.all_tags_list.setCurrentIndex(QModelIndex()) # Select the previously selected image in the unfiltered image # list. - select_index = self.settings.value('image_index', type=int) or 0 + select_index = settings.value('image_index', type=int) or 0 self.image_list.list_view.setCurrentIndex( self.proxy_image_list_model.index(select_index, 0)) else: @@ -423,7 +556,80 @@ def save_image_index(self, proxy_image_index: QModelIndex): settings_key = ('image_index' if self.proxy_image_list_model.filter is None else 'filtered_image_index') - self.settings.setValue(settings_key, proxy_image_index.row()) + settings.setValue(settings_key, proxy_image_index.row()) + + def connect_toolbar_signals(self): + self.toolbar.visibilityChanged.connect( + lambda: self.toggle_toolbar_action.setChecked( + self.toolbar.isVisible())) + self.image_viewer.zoom.connect(self.zoom) + self.zoom_fit_best_action.triggered.connect( + self.image_viewer.zoom_fit) + self.zoom_in_action.triggered.connect( + self.image_viewer.zoom_in) + self.zoom_original_action.triggered.connect( + self.image_viewer.zoom_original) + self.zoom_out_action.triggered.connect( + self.image_viewer.zoom_out) + self.add_action_group.triggered.connect( + lambda action: self.image_viewer.add_marking( + ImageMarking.NONE if not action.isChecked() else + ImageMarking.CROP if action == self.add_crop_action else + ImageMarking.HINT if action == self.add_hint_action else + ImageMarking.EXCLUDE if action == self.add_exclude_action else + ImageMarking.INCLUDE)) + self.image_viewer.marking.connect(lambda marking: + self.add_crop_action.setChecked(True) if marking == ImageMarking.CROP else + self.add_hint_action.setChecked(True) if marking == ImageMarking.HINT else + self.add_exclude_action.setChecked(True) if marking == ImageMarking.EXCLUDE else + self.add_include_action.setChecked(True) if marking == ImageMarking.INCLUDE else + self.add_action_group.checkedAction() and + self.add_action_group.checkedAction().setChecked(False)) + self.image_viewer.scene.selectionChanged.connect(lambda: + self.is_running and self.add_toggle_marking_action.setEnabled( + self.image_viewer.get_selected_type() not in [ImageMarking.NONE, + ImageMarking.CROP])) + self.image_viewer.accept_crop_addition.connect(self.add_crop_action.setEnabled) + self.image_viewer.scene.selectionChanged.connect(lambda: + self.is_running and self.delete_marking_action.setEnabled( + self.image_viewer.get_selected_type() != ImageMarking.NONE)) + self.delete_marking_action.triggered.connect(lambda: self.image_viewer.delete_markings()) + self.add_show_marking_action.toggled.connect(self.image_viewer.show_marking) + self.add_show_marking_action.toggled.connect(self.add_action_group.setEnabled) + self.add_show_marking_action.toggled.connect(lambda toggled: + self.add_toggle_marking_action.setEnabled(toggled and + self.image_viewer.get_selected_type() != ImageMarking.NONE)) + self.add_show_marking_action.toggled.connect(self.add_show_labels_action.setEnabled) + self.add_show_marking_action.toggled.connect(self.add_show_marking_latent_action.setEnabled) + self.add_toggle_marking_action.triggered.connect(lambda: self.image_viewer.change_marking()) + self.add_show_labels_action.toggled.connect(self.image_viewer.show_label) + self.add_show_marking_latent_action.toggled.connect(self.image_viewer.show_marking_latent) + + @Slot(float) + def set_rating(self, rating: float, interactive: bool = False, + event: QMouseEvent|None = None): + """Set the rating from 0.0 to 1.0. + + In the future, half-stars '⯪' might be included, but right now it's + causing display issues.""" + if event is not None and (event.modifiers() & Qt.ControlModifier) == Qt.ControlModifier: + # don't set the image but instead the filter + is_shift = (event.modifiers() & Qt.ShiftModifier) == Qt.ShiftModifier + stars = f'stars:{'>=' if is_shift else '='}{round(rating*5)}' + self.image_list.filter_line_edit.setText(stars) + return + + if interactive and rating == 2.0/10.0 and self.rating == rating: + rating = 0.0 + self.rating = rating + for i, label in enumerate(self.star_labels): + label.setEnabled(True) + label.setText('★' if 2*i+1 < 10.0*rating else '☆') + if interactive: + self.image_list_model.add_to_undo_stack( + action_name='Change rating', should_ask_for_confirmation=False) + self.image_viewer.rating_change(rating) + self.proxy_image_list_model.set_filter(self.proxy_image_list_model.filter) def connect_image_list_signals(self): self.image_list.filter_line_edit.textChanged.connect( @@ -433,7 +639,7 @@ def connect_image_list_signals(self): self.image_list_selection_model.currentChanged.connect( self.image_list.update_image_index_label) self.image_list_selection_model.currentChanged.connect( - self.image_viewer.load_image) + lambda current, previous: self.image_viewer.load_image(current)) self.image_list_selection_model.currentChanged.connect( self.image_tags_editor.load_image_tags) self.image_list_model.modelReset.connect( @@ -444,16 +650,20 @@ def connect_image_list_signals(self): self.image_list_model.images)) self.image_list_model.dataChanged.connect( self.image_tags_editor.reload_image_tags_if_changed) + self.image_list_model.dataChanged.connect( + lambda start, end, roles: + self.image_viewer.load_image(self.image_viewer.proxy_image_index, + False) + if (start.row() <= self.image_viewer.proxy_image_index.row() <= end.row()) else 0) self.image_list_model.update_undo_and_redo_actions_requested.connect( self.update_undo_and_redo_actions) - # Rows are inserted or removed from the proxy image list model when the - # filter is changed. - self.proxy_image_list_model.rowsInserted.connect( - lambda: self.image_list.update_image_index_label( - self.image_list.list_view.currentIndex())) - self.proxy_image_list_model.rowsRemoved.connect( + self.proxy_image_list_model.filter_changed.connect( lambda: self.image_list.update_image_index_label( self.image_list.list_view.currentIndex())) + self.proxy_image_list_model.filter_changed.connect( + lambda: self.tag_counter_model.count_tags_filtered( + self.proxy_image_list_model.get_list() if + len(self.proxy_image_list_model.filter or [])>0 else None)) self.image_list.list_view.directory_reload_requested.connect( self.reload_directory) self.image_list.list_view.tags_paste_requested.connect( @@ -463,6 +673,7 @@ def connect_image_list_signals(self): self.image_list.visibilityChanged.connect( lambda: self.toggle_image_list_action.setChecked( self.image_list.isVisible())) + self.image_viewer.crop_changed.connect(self.image_list.list_view.show_crop_size) @Slot() def update_image_tags(self): @@ -555,21 +766,29 @@ def connect_auto_captioner_signals(self): lambda: self.toggle_auto_captioner_action.setChecked( self.auto_captioner.isVisible())) + def connect_auto_markings_signals(self): + self.auto_markings.marking_generated.connect( + lambda image_index, markings: + self.image_list_model.add_image_markings(image_index, markings)) + self.auto_markings.visibilityChanged.connect( + lambda: self.toggle_auto_markings_action.setChecked( + self.auto_markings.isVisible())) + def restore(self): # Restore the window geometry and state. - if self.settings.contains('geometry'): - self.restoreGeometry(self.settings.value('geometry', type=bytes)) + if settings.contains('geometry'): + self.restoreGeometry(settings.value('geometry', type=bytes)) else: self.showMaximized() - self.restoreState(self.settings.value('window_state', type=bytes)) + self.restoreState(settings.value('window_state', type=bytes)) # Get the last index of the last selected image. - if self.settings.contains('image_index'): - image_index = self.settings.value('image_index', type=int) + if settings.contains('image_index'): + image_index = settings.value('image_index', type=int) else: image_index = 0 # Load the last loaded directory. - if self.settings.contains('directory_path'): - directory_path = Path(self.settings.value('directory_path', + if settings.contains('directory_path'): + directory_path = Path(settings.value('directory_path', type=str)) if directory_path.is_dir(): self.load_directory(directory_path, select_index=image_index)