Fix bug with auto tags (#1), upgrade to wd-v1-4-convnext-tagger.v3

README.md

@@ -8,8 +8,10 @@ Candy Machine is a nascent image tagger for creating Dreambooth models or LORAs
 
 - Automatic file conversion to .png from .jpg, .gif, .webp and more
 - Automatic tag analysis using wd14-convnextv2.v1
+- Placeholder tag templates, i.e. `{type} clothes`, where `{type}` can be specified when adding a tag
 - Customizable tag layout, for consistent tagging
 - Basic image editing (crop, rotate and flip horizontal)
+- Keyboard shortcuts (press '?' for a list)
 - More features coming soon!
 
 Candy Machine runs locally on your machine, and no information is transmitted to a third party server.

TODO.md

@@ -1,5 +1,6 @@
 My immediate TODOS:
 
+- Double "orange shirt" 2303181c8fdfcf81_1024x1360_0.png
 - Add new tags to the front of the list
 - Placeholder isn't highlighted properly when adding a tag
 - Show suggestions when adding a tag to a category

requirements.txt

@@ -7,4 +7,4 @@ Pillow==10.0.1
 deepdanbooru==1.0.2
 onnxruntime==1.17.3
 opencv-python==4.9.0.80
-waitress
+waitress==3.0.0

server/main.py

@@ -1,9 +1,10 @@
+import argparse
 import io
 import os
 import sys
 import threading
-from turtle import listen
 import webbrowser
+from turtle import listen
 from urllib.parse import unquote
 
 from consts import LOWERCASE_IS_TRUE, PROJECTS_DIR
@@ -20,8 +21,7 @@
 from image import Crop, valid_import_directory
 from PIL import Image
 from project import Project
-from waitress import create_server
-import argparse
+from waitress.server import create_server
 
 parser = argparse.ArgumentParser()
 parser.add_argument("--prod", action="store_true")
@@ -216,7 +216,9 @@ def serve(path):
         os.makedirs(PROJECTS_DIR, exist_ok=True)
 
     if IS_PROD:
-        threading.Timer(1.25, lambda: webbrowser.open_new(f"http://127.0.0.1:{PORT}")).start()
+        threading.Timer(
+            1.25, lambda: webbrowser.open_new(f"http://127.0.0.1:{PORT}")
+        ).start()
         server = create_server(app, listen=f"*:{PORT}")
         server.run()
     else:

server/project.py

@@ -494,9 +494,10 @@ def _analyze_auto_tags(self) -> list[TagInfo]:
                     fill_in_examples[fill_in] = fill_in_examples.get(fill_in, []) + [
                         tag
                     ]
-                    del tags[tag]
+                    if tag in tags:
+                        del tags[tag]
                 # Remove the extact match too, because it's covered by the fill in.
-                if tag == ngram:
+                if tag == ngram and tag in tags:
                     del tags[tag]
 
         # Remove any with 1 count.

server/project_test.py

@@ -214,14 +214,70 @@ def test_crop(self):
         output_img_name = self.project.edit_image(
             self.fname, 0, False, Crop(0, 0, 50, 50)
         )
-        self.assertEqual(output_img_name, f"{expected_prefix}_0.png")
+        self.assertEqual(output_img_name, f"{expected_prefix}_1.png")
 
         # Check that the image size is as expected.
         output_img = Image.open(self.project.img_path(output_img_name))
         self.assertEqual(output_img.size, (50, 50))
         output_img.close()
 
-        self.new_kw_fname = self.project.img_path(f"{self.hash}_50x50_0.txt")
+        self.new_kw_fname = self.project.img_path(f"{self.hash}_50x50_1.txt")
+
+
+class TestAnalyzeAutoTags(unittest.TestCase):
+
+    def make_auto_tag_file(self, filename, tags):
+        auto_tags_file = os.path.join(self.temp_project_dir_path, "auto_tags", filename)
+        os.makedirs(os.path.dirname(auto_tags_file), exist_ok=True)
+        with open(auto_tags_file, "w") as f:
+            f.write(tags)
+
+    def setUp(self):
+        # Temp project directory.
+        self.temp_project_dir = tempfile.TemporaryDirectory()
+        self.temp_project_dir_path = self.temp_project_dir.name
+
+        # Temp auto tags.
+        self.make_auto_tag_file("image01.txt", "test, big image, scarf")
+        self.make_auto_tag_file("image02.txt", "test, really big image, red")
+        self.make_auto_tag_file(
+            "image03.txt", "test, very-really big image, red, scarf"
+        )
+
+        self.maxDiff = None
+
+    def tearDown(self):
+        self.temp_project_dir.cleanup()
+
+    def test_analyze_auto_tags(self):
+
+        p = Project(self.temp_project_dir_path)
+        auto_tags = p._analyze_auto_tags()
+
+        results = [atag.to_dict() for atag in auto_tags]
+
+        self.assertEqual(
+            results,
+            [
+                {"tag": "test", "count": 3, "examples": []},
+                {
+                    "tag": "{type} image",
+                    "count": 3,
+                    "examples": [
+                        "big image",
+                        "really big image",
+                        "very-really big image",
+                    ],
+                },
+                {"tag": "scarf", "count": 2, "examples": []},
+                {"tag": "red", "count": 2, "examples": []},
+                {
+                    "tag": "{type} big image",
+                    "count": 2,
+                    "examples": ["really big image", "very-really big image"],
+                },
+            ],
+        )
 
 
 if __name__ == "__main__":

server/thirdparty/tagger/run.py

@@ -1,4 +1,4 @@
-MODEL = "wd14-convnextv2.v1"
+MODEL = "wd-v1-4-convnext-tagger.v3"
 THRESHOLD = 0.35
 EXTENSION = ".txt"
 USE_CPU = False

server/thirdparty/tagger/run_test.py

@@ -4,8 +4,8 @@
 from .run import interrogate_directory
 
 EXPECTED_TAGS = {
-    "apple.txt": "fruit, food, apple, food focus, no humans, realistic, still life, blurry, simple background, close-up",
-    "teeth.txt": "solo, open mouth, teeth, close-up, tongue, realistic, lips, nose, male focus",
+    "apple.txt": "fruit, food, no humans, apple, food focus, realistic, still life, black background, cherry, photorealistic, blurry",
+    "teeth.txt": "realistic, solo, open mouth, teeth, close-up, smile, 1boy, lips, horror \\(theme\\), photorealistic, male focus",
 }