From 8feb29de4035bccb23f24bd9516cdd97b4ddc0cb Mon Sep 17 00:00:00 2001
From: toshiaki1729 <116595002+toshiaki1729@users.noreply.github.com>
Date: Mon, 20 May 2024 19:07:40 +0900
Subject: [PATCH] fix #97 (#98)

and a bug: caption file extension not work on save
---
 scripts/dataset_tag_editor/dte_logic.py | 16 +++++++++++-----
 scripts/tag_editor_ui/block_toprow.py   |  6 +++---
 2 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/scripts/dataset_tag_editor/dte_logic.py b/scripts/dataset_tag_editor/dte_logic.py
index 5e0d4eb..10b2a81 100644
--- a/scripts/dataset_tag_editor/dte_logic.py
+++ b/scripts/dataset_tag_editor/dte_logic.py
@@ -768,7 +768,13 @@ class DatasetTagEditor(Singleton):
             taglists = load_captions(imgpaths)
         
         interrogate_tags = {img_path : [] for img_path in imgpaths}
-        if interrogate_method != self.InterrogateMethod.NONE:
+        
+        img_to_interrogate = [
+        img_path for i, img_path in enumerate(imgpaths) 
+            if (not taglists[i] or interrogate_method != self.InterrogateMethod.PREFILL)
+        ]
+
+        if interrogate_method != self.InterrogateMethod.NONE and img_to_interrogate:
             logger.write("Preprocess images...")
             max_workers = shared.opts.dataset_editor_num_cpu_workers
             if max_workers < 0:
@@ -781,7 +787,7 @@ class DatasetTagEditor(Singleton):
             
             from concurrent.futures import ThreadPoolExecutor
             with ThreadPoolExecutor(max_workers=max_workers) as executor:
-                result = list(executor.map(convert_rgb, gen_data(imgpaths, self.images)))
+                result = list(executor.map(convert_rgb, gen_data(img_to_interrogate, self.images)))
             logger.write("Preprocess completed")
             
             for tg, th in tqdm(tagger_thresholds):
@@ -798,10 +804,10 @@ class DatasetTagEditor(Singleton):
                     continue
                 try:
                     if use_pipe:
-                        for img_path, tags in tqdm(zip(imgpaths, tg.predict_pipe(result, th)), desc=tg.name(), total=len(imgpaths)):
+                        for img_path, tags in tqdm(zip(img_to_interrogate, tg.predict_pipe(result, th)), desc=tg.name(), total=len(img_to_interrogate)):
                             interrogate_tags[img_path] += tags
                     else:
-                        for img_path, data in tqdm(zip(imgpaths, result), desc=tg.name(), total=len(imgpaths)):
+                        for img_path, data in tqdm(zip(img_to_interrogate, result), desc=tg.name(), total=len(img_to_interrogate)):
                             interrogate_tags[img_path] += tg.predict(data, th)
                 except Exception as e:
                     tb = sys.exc_info()[2]
@@ -814,7 +820,7 @@ class DatasetTagEditor(Singleton):
                 tags = interrogate_tags[img_path]
             elif interrogate_method == self.InterrogateMethod.PREPEND:
                 tags = interrogate_tags[img_path] + tags
-            else:
+            elif interrogate_method != self.InterrogateMethod.PREFILL:
                 tags = tags + interrogate_tags[img_path]
 
             self.set_tags_by_image_path(img_path, tags)
diff --git a/scripts/tag_editor_ui/block_toprow.py b/scripts/tag_editor_ui/block_toprow.py
index 9f47454..23ab893 100644
--- a/scripts/tag_editor_ui/block_toprow.py
+++ b/scripts/tag_editor_ui/block_toprow.py
@@ -33,14 +33,14 @@ class ToprowUI(UIBase):
     
     def set_callbacks(self, load_dataset:LoadDatasetUI):
         
-        def save_all_changes(backup: bool, save_kohya_metadata:bool, metadata_output:str, metadata_input:str, metadata_overwrite:bool, metadata_as_caption:bool, metadata_use_fullpath:bool):
+        def save_all_changes(backup: bool, save_kohya_metadata:bool, metadata_output:str, metadata_input:str, metadata_overwrite:bool, metadata_as_caption:bool, metadata_use_fullpath:bool, caption_file_ext:str):
             if not metadata_input:
                 metadata_input = None
-            dte_instance.save_dataset(backup, load_dataset.caption_file_ext, save_kohya_metadata, metadata_output, metadata_input, metadata_overwrite, metadata_as_caption, metadata_use_fullpath)
+            dte_instance.save_dataset(backup, caption_file_ext, save_kohya_metadata, metadata_output, metadata_input, metadata_overwrite, metadata_as_caption, metadata_use_fullpath)
 
         self.btn_save_all_changes.click(
             fn=save_all_changes,
-            inputs=[self.cb_backup, self.cb_save_kohya_metadata, self.tb_metadata_output, self.tb_metadata_input, self.cb_metadata_overwrite, self.cb_metadata_as_caption, self.cb_metadata_use_fullpath]
+            inputs=[self.cb_backup, self.cb_save_kohya_metadata, self.tb_metadata_output, self.tb_metadata_input, self.cb_metadata_overwrite, self.cb_metadata_as_caption, self.cb_metadata_use_fullpath, load_dataset.tb_caption_file_ext]
         )
 
         self.cb_save_kohya_metadata.change(