Fix

Changed `categorized_output` to `final_categorized_output`.

I realized that `categorized_output_string` was being accumulated for all images in the `categorized_output_strings` list, and then a final concatenated string was created from this list resulting in the tags from all images being merged together.

I fixed this issue by removing the Accumulation of categorized_output_strings; Instead of accumulating all `categorized_output_strings` into a single list, store the `categorized_output_string` for each image separately in the `tag_results` dictionary. Also updated the `get_selection_from_gallery` to ensure that this function correctly retrieves the `categorized_output_string` for the selected image from the `tag_results` dictionary.

app.py

@@ -75,7 +75,7 @@ Features:
 - Supports batch processing of multiple images.
 - Tags images with multiple categories: general tags, character tags, and ratings.
 - Displays categorized tags in a structured format.
-- Includes a separate tab for image captioning using Florence 2. This model supports CUDA, MPS or CPU if one of them is available.
+- Includes a separate tab for image captioning using Florence 2. This supports CUDA, MPS or CPU if one of them is available.
 - Supports various captioning tasks (e.g., Caption, Detailed Caption, Object Detection), as well it can display output text and images for tasks that generate visual outputs.
 
 Example image by [me.](https://huggingface.co/Werli)
@@ -448,176 +448,196 @@ class Predictor:
         additional_tags_append,
         tag_results,
         progress=gr.Progress()
-    ):
-        gallery_len = len(gallery)
-        print(f"Predict load model: {model_repo}, gallery length: {gallery_len}")
-
-        timer = Timer()  # Create a timer
-        progressRatio = 0.5 if llama3_reorganize_model_repo else 1
-        progressTotal = gallery_len + 1
-        current_progress = 0
-
-        self.load_model(model_repo)
-        current_progress += progressRatio/progressTotal;
-        progress(current_progress, desc="Initialize wd model finished")
-        timer.checkpoint(f"Initialize wd model")
-
-        # Result
-        txt_infos = []
-        output_dir = tempfile.mkdtemp()
-        if not os.path.exists(output_dir):
-            os.makedirs(output_dir)
-
-        sorted_general_strings = ""
-        rating = None
-        character_res = None
-        general_res = None
-
-        if llama3_reorganize_model_repo:
-            print(f"Llama3 reorganize load model {llama3_reorganize_model_repo}")
-            llama3_reorganize = Llama3Reorganize(llama3_reorganize_model_repo, loadModel=True)
-            current_progress += progressRatio/progressTotal;
-            progress(current_progress, desc="Initialize llama3 model finished")
-            timer.checkpoint(f"Initialize llama3 model")
-            
-        timer.report()
-
-        prepend_list = [tag.strip() for tag in additional_tags_prepend.split(",") if tag.strip()]
-        append_list = [tag.strip() for tag in additional_tags_append.split(",") if tag.strip()]
-        if prepend_list and append_list:
-            append_list = [item for item in append_list if item not in prepend_list]
-            
-        # Dictionary to track counters for each filename
-        name_counters = defaultdict(int)
-        # New code to create categorized output string
-        categorized_output_strings = []
-        for idx, value in enumerate(gallery):
-            try:
-                image_path = value[0]
-                image_name = os.path.splitext(os.path.basename(image_path))[0]
-
-                # Increment the counter for the current name
-                name_counters[image_name] += 1
-                
-                if name_counters[image_name] > 1:
-                    image_name = f"{image_name}_{name_counters[image_name]:02d}"
+        ):
+            # Clear tag_results before starting a new prediction
+            tag_results.clear()
 
-                image = self.prepare_image(image_path)
+            gallery_len = len(gallery)
+            print(f"Predict load model: {model_repo}, gallery length: {gallery_len}")
 
-                input_name = self.model.get_inputs()[0].name
-                label_name = self.model.get_outputs()[0].name
-                print(f"Gallery {idx:02d}: Starting run wd model...")
-                preds = self.model.run([label_name], {input_name: image})[0]
+            timer = Timer()  # Create a timer
+            progressRatio = 0.5 if llama3_reorganize_model_repo else 1
+            progressTotal = gallery_len + 1
+            current_progress = 0
 
-                labels = list(zip(self.tag_names, preds[0].astype(float)))
+            self.load_model(model_repo)
+            current_progress += progressRatio/progressTotal;
+            progress(current_progress, desc="Initialize wd model finished")
+            timer.checkpoint(f"Initialize wd model")
+
+            # Result
+            txt_infos = []
+            output_dir = tempfile.mkdtemp()
+            if not os.path.exists(output_dir):
+                os.makedirs(output_dir)
+
+            sorted_general_strings = ""
+            # New code to create categorized output string
+            categorized_output_strings = []
+            rating = None
+            character_res = None
+            general_res = None
+
+            if llama3_reorganize_model_repo:
+                print(f"Llama3 reorganize load model {llama3_reorganize_model_repo}")
+                llama3_reorganize = Llama3Reorganize(llama3_reorganize_model_repo, loadModel=True)
+                current_progress += progressRatio/progressTotal;
+                progress(current_progress, desc="Initialize llama3 model finished")
+                timer.checkpoint(f"Initialize llama3 model")
+                
+            timer.report()
 
-                # First 4 labels are actually ratings: pick one with argmax
-                ratings_names = [labels[i] for i in self.rating_indexes]
-                rating = dict(ratings_names)
+            prepend_list = [tag.strip() for tag in additional_tags_prepend.split(",") if tag.strip()]
+            append_list = [tag.strip() for tag in additional_tags_append.split(",") if tag.strip()]
+            if prepend_list and append_list:
+                append_list = [item for item in append_list if item not in prepend_list]
+                
+            # Dictionary to track counters for each filename
+            name_counters = defaultdict(int)
+            
+            for idx, value in enumerate(gallery):
+                try:
+                    image_path = value[0]
+                    image_name = os.path.splitext(os.path.basename(image_path))[0]
 
-                # Then we have general tags: pick any where prediction confidence > threshold
-                general_names = [labels[i] for i in self.general_indexes]
+                    # Increment the counter for the current name
+                    name_counters[image_name] += 1
+                    
+                    if name_counters[image_name] > 1:
+                        image_name = f"{image_name}_{name_counters[image_name]:02d}"
 
-                if general_mcut_enabled:
-                    general_probs = np.array([x[1] for x in general_names])
-                    general_thresh = mcut_threshold(general_probs)
+                    image = self.prepare_image(image_path)
 
-                general_res = [x for x in general_names if x[1] > general_thresh]
-                general_res = dict(general_res)
+                    input_name = self.model.get_inputs()[0].name
+                    label_name = self.model.get_outputs()[0].name
+                    print(f"Gallery {idx:02d}: Starting run wd model...")
+                    preds = self.model.run([label_name], {input_name: image})[0]
 
-                # Everything else is characters: pick any where prediction confidence > threshold
-                character_names = [labels[i] for i in self.character_indexes]
+                    labels = list(zip(self.tag_names, preds[0].astype(float)))
 
-                if character_mcut_enabled:
-                    character_probs = np.array([x[1] for x in character_names])
-                    character_thresh = mcut_threshold(character_probs)
-                    character_thresh = max(0.15, character_thresh)
+                    # First 4 labels are actually ratings: pick one with argmax
+                    ratings_names = [labels[i] for i in self.rating_indexes]
+                    rating = dict(ratings_names)
 
-                character_res = [x for x in character_names if x[1] > character_thresh]
-                character_res = dict(character_res)
-                character_list = list(character_res.keys())
+                    # Then we have general tags: pick any where prediction confidence > threshold
+                    general_names = [labels[i] for i in self.general_indexes]
 
-                sorted_general_list = sorted(
-                    general_res.items(),
-                    key=lambda x: x[1],
-                    reverse=True,
-                )
-                sorted_general_list = [x[0] for x in sorted_general_list]
-                #Remove values from character_list that already exist in sorted_general_list
-                character_list = [item for item in character_list if item not in sorted_general_list]
-                #Remove values from sorted_general_list that already exist in prepend_list or append_list
-                if prepend_list:
-                    sorted_general_list = [item for item in sorted_general_list if item not in prepend_list]
-                if append_list:
-                    sorted_general_list = [item for item in sorted_general_list if item not in append_list]
+                    if general_mcut_enabled:
+                        general_probs = np.array([x[1] for x in general_names])
+                        general_thresh = mcut_threshold(general_probs)
 
-                sorted_general_list = prepend_list + sorted_general_list + append_list
+                    general_res = [x for x in general_names if x[1] > general_thresh]
+                    general_res = dict(general_res)
 
-                sorted_general_strings = ", ".join((character_list if characters_merge_enabled else []) + sorted_general_list).replace("(", "\(").replace(")", "\)")
+                    # Everything else is characters: pick any where prediction confidence > threshold
+                    character_names = [labels[i] for i in self.character_indexes]
 
-                classified_tags, unclassified_tags = classify_tags(sorted_general_list)
-                
-                # Create a single string of all categorized tags
-                categorized_output_string = ', '.join([', '.join(tags) for tags in classified_tags.values()])
-                categorized_output_strings.append(categorized_output_string)
+                    if character_mcut_enabled:
+                        character_probs = np.array([x[1] for x in character_names])
+                        character_thresh = mcut_threshold(character_probs)
+                        character_thresh = max(0.15, character_thresh)
 
-                current_progress += progressRatio/progressTotal;
-                progress(current_progress, desc=f"image{idx:02d}, predict finished")
-                timer.checkpoint(f"image{idx:02d}, predict finished")
-                
-                if llama3_reorganize_model_repo:
-                    print(f"Starting reorganize with llama3...")
-                    reorganize_strings = llama3_reorganize.reorganize(sorted_general_strings)
-                    reorganize_strings = re.sub(r" *Title: *", "", reorganize_strings)
-                    reorganize_strings = re.sub(r"\n+", ",", reorganize_strings)
-                    reorganize_strings = re.sub(r",,+", ",", reorganize_strings)
-                    sorted_general_strings += "," + reorganize_strings
+                    character_res = [x for x in character_names if x[1] > character_thresh]
+                    character_res = dict(character_res)
+                    character_list = list(character_res.keys())
 
+                    sorted_general_list = sorted(
+                        general_res.items(),
+                        key=lambda x: x[1],
+                        reverse=True,
+                    )
+                    sorted_general_list = [x[0] for x in sorted_general_list]
+                    #Remove values from character_list that already exist in sorted_general_list
+                    character_list = [item for item in character_list if item not in sorted_general_list]
+                    #Remove values from sorted_general_list that already exist in prepend_list or append_list
+                    if prepend_list:
+                        sorted_general_list = [item for item in sorted_general_list if item not in prepend_list]
+                    if append_list:
+                        sorted_general_list = [item for item in sorted_general_list if item not in append_list]
+
+                    sorted_general_list = prepend_list + sorted_general_list + append_list
+
+                    sorted_general_strings = ", ".join((character_list if characters_merge_enabled else []) + sorted_general_list).replace("(", "\(").replace(")", "\)")
+
+                    classified_tags, unclassified_tags = classify_tags(sorted_general_list)
+                    
+                    # Create a single string of ALL categorized tags for the current image
+                    categorized_output_string = ', '.join([', '.join(tags) for tags in classified_tags.values()])
+                    categorized_output_strings.append(categorized_output_string)
+                    # Collect all categorized output strings into a single string
+                    final_categorized_output = ', '.join(categorized_output_strings)
+                    
                     current_progress += progressRatio/progressTotal;
-                    progress(current_progress, desc=f"image{idx:02d}, llama3 reorganize finished")
-                    timer.checkpoint(f"image{idx:02d}, llama3 reorganize finished")
-
-                txt_file = self.create_file(sorted_general_strings, output_dir, image_name + ".txt")
-                txt_infos.append({"path":txt_file, "name": image_name + ".txt"})
-
-                tag_results[image_path] = { "strings": sorted_general_strings, "classified_tags": classified_tags, "rating": rating, "character_res": character_res, "general_res": general_res, "unclassified_tags": unclassified_tags }
-                timer.report()
-            except Exception as e:
-                print(traceback.format_exc())
-                print("Error predict: " + str(e))
-        # Result
-        download = []
-        if txt_infos is not None and len(txt_infos) > 0:
-            downloadZipPath = os.path.join(output_dir, "images-tagger-" + datetime.now().strftime("%Y%m%d-%H%M%S") + ".zip")
-            with zipfile.ZipFile(downloadZipPath, 'w', zipfile.ZIP_DEFLATED) as taggers_zip:
-                for info in txt_infos:
-                    # Get file name from lookup
-                    taggers_zip.write(info["path"], arcname=info["name"])
-            download.append(downloadZipPath)
+                    progress(current_progress, desc=f"image{idx:02d}, predict finished")
+                    timer.checkpoint(f"image{idx:02d}, predict finished")
+                    
+                    if llama3_reorganize_model_repo:
+                        print(f"Starting reorganize with llama3...")
+                        reorganize_strings = llama3_reorganize.reorganize(sorted_general_strings)
+                        reorganize_strings = re.sub(r" *Title: *", "", reorganize_strings)
+                        reorganize_strings = re.sub(r"\n+", ",", reorganize_strings)
+                        reorganize_strings = re.sub(r",,+", ",", reorganize_strings)
+                        sorted_general_strings += "," + reorganize_strings
+
+                        current_progress += progressRatio/progressTotal;
+                        progress(current_progress, desc=f"image{idx:02d}, llama3 reorganize finished")
+                        timer.checkpoint(f"image{idx:02d}, llama3 reorganize finished")
+
+                    txt_file = self.create_file(sorted_general_strings, output_dir, image_name + ".txt")
+                    txt_infos.append({"path":txt_file, "name": image_name + ".txt"})
+                    
+                    # Store the result in tag_results using image_path as the key
+                    tag_results[image_path] = { 
+                        "strings": sorted_general_strings, 
+                        "strings2": categorized_output_string,  # Store the categorized output string here
+                        "classified_tags": classified_tags, 
+                        "rating": rating, 
+                        "character_res": character_res, 
+                        "general_res": general_res, 
+                        "unclassified_tags": unclassified_tags 
+                    }
+
+                    timer.report()
+                except Exception as e:
+                    print(traceback.format_exc())
+                    print("Error predict: " + str(e))
+            # Result
+            download = []
+            if txt_infos is not None and len(txt_infos) > 0:
+                downloadZipPath = os.path.join(output_dir, "images-tagger-" + datetime.now().strftime("%Y%m%d-%H%M%S") + ".zip")
+                with zipfile.ZipFile(downloadZipPath, 'w', zipfile.ZIP_DEFLATED) as taggers_zip:
+                    for info in txt_infos:
+                        # Get file name from lookup
+                        taggers_zip.write(info["path"], arcname=info["name"])
+                download.append(downloadZipPath)
+                
+            if llama3_reorganize_model_repo:
+                llama3_reorganize.release_vram()
+                del llama3_reorganize
+                 
+            progress(1, desc=f"Predict completed")
+            timer.report_all()  # Print all recorded times
+            print("Predict is complete.")
             
-        if llama3_reorganize_model_repo:
-            llama3_reorganize.release_vram()
-            del llama3_reorganize
-         
-        progress(1, desc=f"Predict completed")
-        timer.report_all()  # Print all recorded times
-        print("Predict is complete.")
-        
-        # Collect all categorized output strings into a single string
-        final_categorized_output = ', '.join(categorized_output_strings)
-        
-        return download, sorted_general_strings, classified_tags, rating, character_res, general_res, unclassified_tags, tag_results, final_categorized_output
-# END
+            return download, sorted_general_strings, final_categorized_output, classified_tags, rating, character_res, general_res, unclassified_tags, tag_results
 
 def get_selection_from_gallery(gallery: list, tag_results: dict, selected_state: gr.SelectData):
     if not selected_state:
         return selected_state
 
-    tag_result = { "strings": "", "classified_tags": "{}", "rating": "", "character_res": "", "general_res": "", "unclassified_tags": "{}" }
+    tag_result = { 
+        "strings": "", 
+        "strings2": "", 
+        "classified_tags": "{}", 
+        "rating": "", 
+        "character_res": "", 
+        "general_res": "", 
+        "unclassified_tags": "{}" 
+    }
     if selected_state.value["image"]["path"] in tag_results:
         tag_result = tag_results[selected_state.value["image"]["path"]]
 
-    return (selected_state.value["image"]["path"], selected_state.value["caption"]), tag_result["strings"], tag_result["classified_tags"], tag_result["rating"], tag_result["character_res"], tag_result["general_res"], tag_result["unclassified_tags"]
+    return (selected_state.value["image"]["path"], selected_state.value["caption"]), tag_result["strings"], tag_result["strings2"], tag_result["classified_tags"], tag_result["rating"], tag_result["character_res"], tag_result["general_res"], tag_result["unclassified_tags"]
 
 def append_gallery(gallery: list, image: str):
     if gallery is None:
@@ -645,7 +665,7 @@ def remove_image_from_gallery(gallery: list, selected_image: str):
     if not gallery or not selected_image:
         return gallery
 
-    selected_image = ast.literal_eval(selected_image) #Use ast.literal_eval to parse text into a tuple.
+    selected_image = ast.literal_eval(selected_image) # Use ast.literal_eval to parse text into a tuple.
     # Remove the selected image from the gallery
     if selected_image in gallery:
         gallery.remove(selected_image)
@@ -860,7 +880,7 @@ def process_image(image, task_prompt, text_input=None):
         return results, output_image
     else:
         return "", None  # Return empty string and None for unknown task prompts
-##############
+
 # Custom CSS to set the height of the gr.Dropdown menu
 css = """
 div.progress-level div.progress-level-inner {
@@ -927,7 +947,6 @@ scheduler.start()
 next_run_time_utc = restart_space_job.next_run_time.astimezone(timezone.utc)
 NEXT_RESTART = f"Next Restart: {next_run_time_utc.strftime('%Y-%m-%d %H:%M:%S')} (UTC) - The space will restart every 2 days to ensure stability and performance. It uses a background scheduler to handle the restart process."
 
-# Using "JohnSmith9982/small_and_pretty" theme
 with gr.Blocks(title=TITLE, css=css, theme="Werli/Multi-Tagger", fill_width=True) as demo:
     gr.Markdown(value=f"<h1 style='text-align: center; margin-bottom: 1rem'>{TITLE}</h1>")
     gr.Markdown(value=DESCRIPTION)
@@ -1014,7 +1033,7 @@ with gr.Blocks(title=TITLE, css=css, theme="Werli/Multi-Tagger", fill_width=True
                 download_file = gr.File(label="Output (Download)") # 0
                 character_res = gr.Label(label="Output (characters)") # 1
                 sorted_general_strings = gr.Textbox(label="Output (string)", show_label=True, show_copy_button=True) # 2
-                categorized_output = gr.Textbox(label="Categorized Output (string)", show_label=True, show_copy_button=True) # 3
+                final_categorized_output = gr.Textbox(label="Categorized Output (string)", show_label=True, show_copy_button=True) # 3
                 categorized = gr.JSON(label="Categorized (tags)") # 4
                 rating = gr.Label(label="Rating") # 5
                 general_res = gr.Label(label="Output (tags)") # 6
@@ -1023,6 +1042,7 @@ with gr.Blocks(title=TITLE, css=css, theme="Werli/Multi-Tagger", fill_width=True
                     [
                         download_file,
                         sorted_general_strings,
+                        final_categorized_output,
                         categorized,
                         rating,
                         character_res,
@@ -1037,7 +1057,7 @@ with gr.Blocks(title=TITLE, css=css, theme="Werli/Multi-Tagger", fill_width=True
             upload_button.upload(extend_gallery, inputs=[gallery, upload_button], outputs=gallery)
             # Event to update the selected image when an image is clicked in the gallery
             selected_image = gr.Textbox(label="Selected Image", visible=False)
-            gallery.select(get_selection_from_gallery, inputs=[gallery, tag_results], outputs=[selected_image, sorted_general_strings, categorized, rating, character_res, general_res, unclassified])
+            gallery.select(get_selection_from_gallery, inputs=[gallery, tag_results], outputs=[selected_image, sorted_general_strings, final_categorized_output, categorized, rating, character_res, general_res, unclassified])
             # Event to remove a selected image from the gallery
             remove_button.click(remove_image_from_gallery, inputs=[gallery, selected_image], outputs=gallery)
         submit.click(
@@ -1055,8 +1075,8 @@ with gr.Blocks(title=TITLE, css=css, theme="Werli/Multi-Tagger", fill_width=True
                 additional_tags_append,
                 tag_results,
             ],
-            outputs=[download_file, sorted_general_strings, categorized, rating, character_res, general_res, unclassified, tag_results,  categorized_output,],
-        )  
+            outputs=[download_file, sorted_general_strings, final_categorized_output, categorized, rating, character_res, general_res, unclassified, tag_results,],
+        )
         gr.Examples(
             [["images/1girl.png", VIT_LARGE_MODEL_DSV3_REPO, 0.35, False, 0.85, False]], 
             inputs=[
@@ -1078,8 +1098,7 @@ with gr.Blocks(title=TITLE, css=css, theme="Werli/Multi-Tagger", fill_width=True
                 text_input = gr.Textbox(label="Text Input (optional)")
                 submit_btn = gr.Button(value="Submit")
             with gr.Column(variant="panel"):
-            #OUTPUT
-                output_text = gr.Textbox(label="Output Text", show_label=True, show_copy_button=True, lines=8) # Here is the problem!
+                output_text = gr.Textbox(label="Output Text", show_label=True, show_copy_button=True, lines=8)
                 output_img = gr.Image(label="Output Image")
         gr.Examples(
             examples=[
@@ -1094,4 +1113,4 @@ with gr.Blocks(title=TITLE, css=css, theme="Werli/Multi-Tagger", fill_width=True
         )
         submit_btn.click(process_image, [input_img, task_prompt, text_input], [output_text, output_img])
         
-demo.queue(max_size=2).launch() # test
+demo.queue(max_size=2).launch()