From 8e4690b1b3b05a03adbae835a18821de6e6df1c8 Mon Sep 17 00:00:00 2001
From: dokutan <54861821+dokutan@users.noreply.github.com>
Date: Sun, 26 Jun 2022 21:05:10 +0200
Subject: [PATCH] Skip files that aren't images

---
 autotagger/autotagger.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/autotagger/autotagger.py b/autotagger/autotagger.py
index d195e3a..c2384ba 100644
--- a/autotagger/autotagger.py
+++ b/autotagger/autotagger.py
@@ -28,7 +28,12 @@ class Autotagger:
 
     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]
+            def create_image(file):
+                try:
+                    return PILImage.create(file)
+                except:
+                    return None
+            images = list(filter(lambda i: i != None, [create_image(file) for file in files]))
             dl = self.learn.dls.test_dl(images, bs=bs)
             batch, _ = self.learn.get_preds(dl=dl)