autotag: accept files from stdin with '-'.

Fix `autotag` so you can pass a filename of '-' to read a file from
stdin. This way you can do this:

  docker run --rm -i ghcr.io/danbooru/autotagger autotag - < image.jpg

...to perform prediction on a single file outside of Docker.
This commit is contained in:
evazion
2022-06-22 15:21:30 -05:00
parent e4344ab2d4
commit 96b6a12924
3 changed files with 11 additions and 11 deletions
+2 -1
View File
@@ -26,8 +26,9 @@ class Autotagger:
return learn
def predict(self, images, threshold=0.01, limit=50, bs=64):
def predict(self, files, threshold=0.01, limit=50, bs=64):
with self.learn.no_bar(), self.learn.no_logging():
images = [PILImage.create(file) for file in files]
dl = self.learn.dls.test_dl(images, bs=bs)
batch, _ = self.learn.get_preds(dl=dl)