Add return type.
This commit is contained in:
@@ -1,3 +1,4 @@
|
||||
from typing import Generator
|
||||
from tagger.interrogator import Interrogator
|
||||
from PIL import Image
|
||||
from pathlib import Path
|
||||
@@ -60,7 +61,7 @@ def image_interrogate(image_path: Path):
|
||||
return Interrogator.postprocess_tags(result[1], threshold=args.threshold)
|
||||
return Interrogator.postprocess_tags(result[1], threshold=args.threshold, escape_tag=True, replace_underscore=True)
|
||||
|
||||
def explore_files(folder_path: Path):
|
||||
def explore_image_files(folder_path: Path) -> Generator[Path, None, None]:
|
||||
"""
|
||||
Explore files by folder path
|
||||
"""
|
||||
@@ -68,13 +69,12 @@ def explore_files(folder_path: Path):
|
||||
if path.is_file() and path.suffix in ['.png', '.jpg', '.jpeg', '.webp']:
|
||||
yield path
|
||||
elif args.recursive and path.is_dir():
|
||||
yield from explore_files(path)
|
||||
yield from explore_image_files(path)
|
||||
|
||||
if args.dir:
|
||||
root_path = Path(args.dir)
|
||||
for file_path in explore_files(root_path):
|
||||
image_path = Path(file_path)
|
||||
caption_path = image_path.parent / f'{file_path.stem}{args.ext}'
|
||||
for image_path in explore_image_files(root_path):
|
||||
caption_path = image_path.parent / f'{image_path.stem}{args.ext}'
|
||||
|
||||
if caption_path.is_file() and not args.overwrite:
|
||||
# skip if caption exists
|
||||
|
||||
Reference in New Issue
Block a user