Merge pull request #3 from dokutan/master

Skip files that aren't images
2022-06-28 03:28:26 -05:00
parent d43ce9c4e7 8e4690b1b3
commit 38beec98e4
1 changed files with 6 additions and 1 deletions
@@ -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)