draft pipeline

app.py

@@ -0,0 +1,82 @@
+import gradio as gr
+import json
+from transformers import pipeline, AutoModelForImageClassification, AutoFeatureExtractor
+import requests
+import asyncio
+import httpx
+import time
+import io
+from PIL import Image
+import PIL
+
+HF_MODEL_PATH = (
+    "ImageIN/levit-192_finetuned_on_unlabelled_IA_with_snorkel_labels"
+)
+
+classif_model = AutoModelForImageClassification.from_pretrained(HF_MODEL_PATH)
+feature_extractor = AutoFeatureExtractor.from_pretrained(HF_MODEL_PATH)
+
+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()
+
+
+def get_image_urls_from_manifest(data):
+    image_urls = []
+    for sequences in data['sequences']:
+        for canvases in sequences['canvases']:
+            image_urls.extend(image['resource']['@id'] for image in canvases['images'])
+    return image_urls
+
+
+def resize_iiif_urls(im_url, size='224'):
+    parts = im_url.split("/")
+    parts[6] = size, size
+    return "/".join(parts)
+
+
+async def get_image(client, url):
+    try:
+        resp = await client.get(url, timeout=30)
+        return Image.open(io.BytesIO(resp.content))
+    except (PIL.UnidentifiedImageError, httpx.ReadTimeout):
+        return None
+
+
+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]
+
+
+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))
+    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']))
+            if len(predicted_images):
+                return predicted_images
+
+
+demo = gr.Interface(
+    fn=predict,
+    inputs=gr.Text(),
+    outputs=gr.Gallery(),
+    title="ImageIN",
+    description="Identify illustrations in pages of historical books!",
+)
+demo.launch(debug=True, share=True)