autotag: fix bug when reading file from stdin.

Another attempt at fixing #2. b216057 had a bug where you couldn't pass
a file through stdin because we tried to read the file twice, which we
can't do if the input is a pipe.
This commit is contained in:
evazion
2022-06-29 20:02:18 -05:00
parent 069af1c0c4
commit aaaae88e77
3 changed files with 13 additions and 12 deletions
+9 -9
View File
@@ -5,6 +5,7 @@ import click
import itertools
import logging
import PIL
from fastai.vision.core import PILImage
from autotagger import Autotagger
from pathlib import Path
from more_itertools import ichunked
@@ -31,10 +32,13 @@ def main(files, threshold, limit, bs, csv, input_file, group_tags, name_only, mo
click.get_current_context().exit()
for filepaths in ichunked(paths, bs):
files = list(filter(None, [open_image(filepath) for filepath in filepaths]))
predictions = autotagger.predict(files, threshold=threshold, limit=limit, bs=bs)
paths_with_images = list(filter(None, [open_image(filepath) for filepath in filepaths]))
filepaths = [x[0] for x in paths_with_images]
images = [x[1] for x in paths_with_images]
for file, tags in zip(files, predictions):
predictions = autotagger.predict(images, threshold=threshold, limit=limit, bs=bs)
for file, tags in zip(filepaths, predictions):
output_result(Path(file.name), tags, csv, group_tags, name_only)
def output_result(filepath, tags, csv, group_tags, name_only):
@@ -63,12 +67,8 @@ def recurse_dir(directory):
def open_image(filepath):
try:
# Load the image to check that it's a valid file. Open the file twice because PIL closes the underlying file.
file = click.open_file(filepath, "rb")
image = PIL.Image.open(file)
image.load()
image.close()
return click.open_file(filepath, "rb")
with click.open_file(filepath, "rb") as file:
return (filepath, PILImage.create(file))
except PIL.UnidentifiedImageError as err:
logging.warning(f"Skipped {filepath} (not an image)")
return None