update

app.py

@@ -6,7 +6,8 @@ import httpx
 import io
 from PIL import Image
 import PIL
-
+from toolz import pluck
+from piffle.image import IIIFImageClient
 
 HF_MODEL_PATH = (
     "ImageIN/levit-192_finetuned_on_unlabelled_IA_with_snorkel_labels"
@@ -19,9 +20,6 @@ classif_pipeline = pipeline(
     "image-classification", model=classif_model, feature_extractor=feature_extractor
 )
 
-OUTPUT_SENTENCE = "This image is {result}."
-
-
 def load_manifest(inputs):
     with requests.get(inputs) as r:
         return r.json()
@@ -53,26 +51,44 @@ async def get_images(urls):
     async with httpx.AsyncClient() as client:
         tasks = [asyncio.ensure_future(get_image(client, url)) for url in urls]
         images = await asyncio.gather(*tasks)
-        return [image for image in images if image is not None]
+        assert len(images) == len(urls)
+        image_url_tuples = []
+        for url, image in zip(urls, images):
+            if image is not None:
+                image_url_tuples.append((url, image))
+        return image_url_tuples
+        # return [image for image in images if image is not None]
 
 
 def predict(inputs):
     data = load_manifest(inputs)
     urls = get_image_urls_from_manifest(data)
     resized_urls = [resize_iiif_urls(url) for url in urls]
-    images = asyncio.run(get_images(resized_urls))
+    images_urls = asyncio.run(get_images(resized_urls))
     predicted_images = []
-    for image in images:
-        top_pred = classif_pipeline(image, top_k=1)[0]
+    images = list(pluck(1, images_urls))
+    urls = list(pluck(0, images_urls))
+    predictions = classif_pipeline(images, top_k=1)
+    for url, pred in zip(urls,predictions):
+        top_pred = pred[0]
         if top_pred['label'] == 'illustrated':
-            predicted_images.append((image, top_pred['score']))
-        return predicted_images
-
+            image_url = IIIFImageClient.init_from_url(url)
+            image_url = image_url.canonicalize()
+            predicted_images.append((image_url.__str__(), f"Confidence: {top_pred['score']}, page: {10}"))
+    return predicted_images
+    # for image in images:
+    #     top_pred = classif_pipeline(image, top_k=1)[0]
+    #     if top_pred['label'] == 'illustrated':
+    #         predicted_images.append((image, top_pred['score']))
+    #     return predicted_images
+
+gallery = gr.Gallery()
+gallery.style(grid=3)
 
 demo = gr.Interface(
     fn=predict,
     inputs=gr.Text(),
-    outputs=gr.Gallery(),
+    outputs=gallery,
     title="ImageIN",
     description="Identify illustrations in pages of historical books!",
 )

requirements.txt

@@ -1,3 +1,5 @@
 transformers
 httpx
-torch
+torch
+toolz
+piffle