From cf0b86be808cfe52e7f17959ba61b83e2dbc78fa Mon Sep 17 00:00:00 2001
From: corkborg <62742115+corkborg@users.noreply.github.com>
Date: Sun, 1 Oct 2023 19:17:25 +0900
Subject: [PATCH] Improve code.

---
 run.py                 | 8 ++++++--
 tagger/interrogator.py | 6 ------
 2 files changed, 6 insertions(+), 8 deletions(-)

diff --git a/run.py b/run.py
index deb40c7..d8d5bed 100644
--- a/run.py
+++ b/run.py
@@ -1,7 +1,9 @@
 
 
+import sys
 from tagger.interrogator import Interrogator, WaifuDiffusionInterrogator
 from PIL import Image
+from pathlib import Path
 
 interrogator = WaifuDiffusionInterrogator(
     'wd14-convnextv2-v2',
@@ -9,9 +11,11 @@ interrogator = WaifuDiffusionInterrogator(
     revision='v2.0'
 )
 
-im = Image.open("image.jpg")
+image_path = Path(sys.argv[1])
+
+im = Image.open(image_path)
 result = interrogator.interrogate(im)
-tags = Interrogator.postprocess_tags(result[1])
+tags = Interrogator.postprocess_tags(result[1], threshold=0.35)
 print()
 
 for k, v in tags.items():
diff --git a/tagger/interrogator.py b/tagger/interrogator.py
index cb7dd48..138972a 100644
--- a/tagger/interrogator.py
+++ b/tagger/interrogator.py
@@ -10,16 +10,11 @@ from pathlib import Path
 from huggingface_hub import hf_hub_download
 import re
 
-# from tagger.deepbooru import re_special as tag_escape_pattern
-
 tag_escape_pattern = re.compile(r'([\\()])')
 
-# i'm not sure if it's okay to add this file to the repository
 import tagger.dbimutils as dbimutils
 
-# select a device to process
 use_cpu = True
-
 if use_cpu:
     tf_device_name = '/cpu:0'
 else:
@@ -29,7 +24,6 @@ class Interrogator:
     @staticmethod
     def postprocess_tags(
         tags: Dict[str, float],
-
         threshold=0.35,
         additional_tags: List[str] = [],
         exclude_tags: List[str] = [],