kapitan/wdtagger
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

🩹 fix(color): rgb -> bgr

Browse Source
This commit is contained in:
Jannchie 2024-06-11 16:15:40 +09:00
parent ae2838f58e
commit 5e4629b8ea
1 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
wdtagger/__init__.py
View File

@ -183,7 +183,7 @@ class Tagger:
self.model_target_size = None
self.cache_dir = cache_dir
self.hf_token = hf_token
self.load_model(model_repo, cache_dir, hf_token)
self.load_model(model_repo, cache_dir, hf_token, num_threads=8)
def load_model(
self, model_repo, cache_dir=None, hf_token=None, num_threads: int = None
@ -252,9 +252,9 @@ class Tagger:
(target_size, target_size),
Image.BICUBIC,
)
return np.asarray(padded_image, dtype=np.float32)
array = np.asarray(padded_image, dtype=np.float32)
array = array[:, :, [2, 1, 0]]
return array
def tag(
self,
image: Image.Image | list[Image.Image],