From 4ec8d234a28b32e691b49ca5913415d6c616fa84 Mon Sep 17 00:00:00 2001 From: toshiaki1729 <116595002+toshiaki1729@users.noreply.github.com> Date: Sun, 29 Jan 2023 01:07:41 +0900 Subject: [PATCH] use glob instead of my redundant code --- .../dataset_tag_editor/dataset_tag_editor.py | 23 ++----------------- 1 file changed, 2 insertions(+), 21 deletions(-) diff --git a/scripts/dataset_tag_editor/dataset_tag_editor.py b/scripts/dataset_tag_editor/dataset_tag_editor.py index 59c3858..169e12c 100644 --- a/scripts/dataset_tag_editor/dataset_tag_editor.py +++ b/scripts/dataset_tag_editor/dataset_tag_editor.py @@ -1,5 +1,6 @@ import os import re +import glob from typing import List, Set, Optional from modules import shared from modules.textual_inversion.dataset import re_numbers_at_start @@ -90,26 +91,6 @@ def interrogate_image_waifu_v2(path, threshold): return ', '.join(tagger.get_arranged_tags(res)) -def get_filepath_set(dir: str, recursive: bool): - if recursive: - dirs_to_see = [dir] - result = set() - while len(dirs_to_see) > 0: - current_dir = dirs_to_see.pop() - basenames = os.listdir(current_dir) - paths = {os.path.join(current_dir, basename) for basename in basenames} - for path in paths: - if os.path.isdir(path): - dirs_to_see.append(path) - elif os.path.isfile(path): - result.add(path) - return result - else: - basenames = os.listdir(dir) - paths = {os.path.join(dir, basename) for basename in basenames} - return {path for path in paths if os.path.isfile(path)} - - class DatasetTagEditor: def __init__(self): # from modules.textual_inversion.dataset @@ -430,7 +411,7 @@ class DatasetTagEditor: print(f'[tag-editor] Also loading from subdirectories.') try: - filepath_set = get_filepath_set(dir=img_dir, recursive=recursive) + filepath_set = [p for p in glob.glob(os.path.join(img_dir, '**'), recursive=recursive) if os.path.isfile(p)] except Exception as e: print(e) print('[tag-editor] Loading Aborted.')