From c1c78f597679abb008f3e8f1add84620e1aec35e Mon Sep 17 00:00:00 2001 From: Chris Mekarski Date: Sun, 2 Apr 2023 07:15:32 -0400 Subject: [PATCH] Add max image resolution option (#70) * Add max image resolution option * disable webUI's custom feature when max_res is specified --------- Co-authored-by: toshiaki1729 <116595002+toshiaki1729@users.noreply.github.com> --- scripts/dataset_tag_editor/dte_logic.py | 7 +++++-- scripts/main.py | 1 + scripts/ui/block_load_dataset.py | 2 +- 3 files changed, 7 insertions(+), 3 deletions(-) diff --git a/scripts/dataset_tag_editor/dte_logic.py b/scripts/dataset_tag_editor/dte_logic.py index 3fa2630..b2d862c 100644 --- a/scripts/dataset_tag_editor/dte_logic.py +++ b/scripts/dataset_tag_editor/dte_logic.py @@ -482,7 +482,7 @@ class DatasetTagEditor(Singleton): print(e) - def load_dataset(self, img_dir:str, caption_ext:str, recursive:bool, load_caption_from_filename:bool, replace_new_line:bool, interrogate_method:InterrogateMethod, interrogator_names:List[str], threshold_booru:float, threshold_waifu:float, use_temp_dir:bool, kohya_json_path:Optional[str]): + def load_dataset(self, img_dir:str, caption_ext:str, recursive:bool, load_caption_from_filename:bool, replace_new_line:bool, interrogate_method:InterrogateMethod, interrogator_names:List[str], threshold_booru:float, threshold_waifu:float, use_temp_dir:bool, kohya_json_path:Optional[str], max_res:float): self.clear() img_dir_obj = Path(img_dir) @@ -511,11 +511,14 @@ class DatasetTagEditor(Singleton): continue try: img = Image.open(img_path) + if (max_res > 0): + img_res = int(max_res), int(max_res) + img.thumbnail(img_res) except: continue else: abs_path = str(img_path.absolute()) - if not use_temp_dir: + if not use_temp_dir and max_res <= 0: img.already_saved_as = abs_path images[abs_path] = img diff --git a/scripts/main.py b/scripts/main.py index 75c4f74..35e0f8c 100644 --- a/scripts/main.py +++ b/scripts/main.py @@ -339,6 +339,7 @@ def on_ui_tabs(): def on_ui_settings(): section = ('dataset-tag-editor', "Dataset Tag Editor") shared.opts.add_option("dataset_editor_image_columns", shared.OptionInfo(6, "Number of columns on image gallery", section=section)) + shared.opts.add_option("dataset_editor_max_res", shared.OptionInfo(0, "Max resolution of temporary files", section=section)) shared.opts.add_option("dataset_editor_use_temp_files", shared.OptionInfo(False, "Force image gallery to use temporary files", section=section)) shared.opts.add_option("dataset_editor_use_raw_clip_token", shared.OptionInfo(True, "Use raw CLIP token to calculate token count (without emphasis or embeddings)", section=section)) diff --git a/scripts/ui/block_load_dataset.py b/scripts/ui/block_load_dataset.py index 0078871..9437a4d 100644 --- a/scripts/ui/block_load_dataset.py +++ b/scripts/ui/block_load_dataset.py @@ -77,7 +77,7 @@ class LoadDatasetUI(UIBase): threshold_booru = custom_threshold_booru if use_custom_threshold_booru else shared.opts.interrogate_deepbooru_score_threshold threshold_waifu = custom_threshold_waifu if use_custom_threshold_waifu else -1 - dte_instance.load_dataset(dir, caption_file_ext, recursive, load_caption_from_filename, replace_new_line, interrogate_method, use_interrogator_names, threshold_booru, threshold_waifu, opts.dataset_editor_use_temp_files, kohya_json_path if use_kohya_metadata else None) + dte_instance.load_dataset(dir, caption_file_ext, recursive, load_caption_from_filename, replace_new_line, interrogate_method, use_interrogator_names, threshold_booru, threshold_waifu, opts.dataset_editor_use_temp_files, kohya_json_path if use_kohya_metadata else None, opts.dataset_editor_max_res) imgs = dte_instance.get_filtered_imgs(filters=[]) img_indices = dte_instance.get_filtered_imgindices(filters=[]) return [