add json input

app.py

@@ -1,57 +1,102 @@
 from transformers import AutoFeatureExtractor, AutoModelForImageClassification, pipeline
-import torch 
+import torch
 from PIL import Image
 import gradio as gr
 
+import aiohttp
+import asyncio
+from io import BytesIO
+
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 dtype = torch.float16
-nsfw_pipe = pipeline("image-classification", 
-                          model= AutoModelForImageClassification.from_pretrained("carbon225/vit-base-patch16-224-hentai"),
-                          feature_extractor=AutoFeatureExtractor.from_pretrained("carbon225/vit-base-patch16-224-hentai"),
-                          device=device,
-                          torch_dtype=dtype)
+nsfw_pipe = pipeline("image-classification",
+                     model=AutoModelForImageClassification.from_pretrained(
+                         "carbon225/vit-base-patch16-224-hentai"),
+                     feature_extractor=AutoFeatureExtractor.from_pretrained(
+                         "carbon225/vit-base-patch16-224-hentai"),
+                     device=device,
+                     torch_dtype=dtype)
 
 
-style_pipe = pipeline("image-classification", 
-                          model= AutoModelForImageClassification.from_pretrained("cafeai/cafe_style"),
-                          feature_extractor=AutoFeatureExtractor.from_pretrained("cafeai/cafe_style"),
-                          device=device,
-                          torch_dtype=dtype)
+style_pipe = pipeline("image-classification",
+                      model=AutoModelForImageClassification.from_pretrained(
+                          "cafeai/cafe_style"),
+                      feature_extractor=AutoFeatureExtractor.from_pretrained(
+                          "cafeai/cafe_style"),
+                      device=device,
+                      torch_dtype=dtype)
 
-aesthetic_pipe = pipeline("image-classification", 
-                          model= AutoModelForImageClassification.from_pretrained("cafeai/cafe_aesthetic"),
-                          feature_extractor=AutoFeatureExtractor.from_pretrained("cafeai/cafe_aesthetic"),
+aesthetic_pipe = pipeline("image-classification",
+                          model=AutoModelForImageClassification.from_pretrained(
+                              "cafeai/cafe_aesthetic"),
+                          feature_extractor=AutoFeatureExtractor.from_pretrained(
+                              "cafeai/cafe_aesthetic"),
                           device=device,
                           torch_dtype=dtype)
 
-def predict(image, files=None):
-    images_paths = [image]
-    if not files == None: 
-        images_paths = list(map(lambda x: x.name, files))
-    pil_images = [Image.open(image_path).convert("RGB") for image_path in images_paths]
-    
+
+async def fetch_image(session, image_url):
+    print(f"fetching image {image_url}")
+    async with session.get(image_url) as response:
+        if response.status == 200 and response.headers['content-type'].startswith('image'):
+            pil_image = Image.open(BytesIO(await response.read())).convert('RGB')
+            # resize image proportional
+            # image = ImageOps.fit(image, (400, 400), Image.LANCZOS)
+
+            return pil_image
+    return None
+
+
+async def fetch_images(image_urls):
+    async with aiohttp.ClientSession() as session:
+        tasks = [asyncio.ensure_future(fetch_image(
+            session, image_url)) for image_url in image_urls]
+        return await asyncio.gather(*tasks)
+
+
+async def predict(json=None, enable_gallery=True, image=None, files=None):
+    print(json)
+
+    if image or files:
+        if image is not None:
+            images_paths = [image]
+        elif files is not None:
+            images_paths = list(map(lambda x: x.name, files))
+        pil_images = [Image.open(image_path).convert("RGB")
+                      for image_path in images_paths]
+    elif json is not None:
+        pil_images = await fetch_images(json["urls"])
+
     style = style_pipe(pil_images)
     aesthetic = aesthetic_pipe(pil_images)
     nsfw = nsfw_pipe(pil_images)
-    results = [ a + b + c  for (a,b,c) in zip(style, aesthetic, nsfw)]
+    results = [a + b + c for (a, b, c) in zip(style, aesthetic, nsfw)]
     label_data = {}
     if image is not None:
-        label_data = { row["label"]:row["score"] for row in results[0] }
-    
-    return label_data, results
+        label_data = {row["label"]: row["score"] for row in results[0]}
+
+    return results, label_data, pil_images if enable_gallery else None
+
 
 with gr.Blocks() as blocks:
     with gr.Row():
         with gr.Column():
             image = gr.Image(label="Image to test", type="filepath")
-            files = gr.File(label="Multipls Images", file_types=["image"], file_count="multiple")
+            files = gr.File(label="Multipls Images", file_types=[
+                            "image"], file_count="multiple")
+            enable_gallery = gr.Checkbox(label="Enable Gallery", value=True)
+            json = gr.JSON(label="Results", value={"urls": [
+                'https://d26smi9133w0oo.cloudfront.net/diffusers-gallery/b9fb3257-6a54-455e-b636-9d61cf261676.jpg',
+                'https://d26smi9133w0oo.cloudfront.net/diffusers-gallery/062eb9be-76eb-4d7e-9299-d1ebea14b46f.jpg',
+                'https://d26smi9133w0oo.cloudfront.net/diffusers-gallery/8ff6d4f6-08d0-4a31-818c-4d32ab146f81.jpg']})
         with gr.Column():
             label = gr.Label(label="style")
             results = gr.JSON(label="Results")
-            # gallery = gr.Gallery().style(grid=[2], height="auto")
+            gallery = gr.Gallery().style(grid=[2], height="auto")
     btn = gr.Button("Run")
-    
-    btn.click(fn=predict, inputs=[image, files], outputs=[label, results], api_name="inference")
+
+    btn.click(fn=predict, inputs=[json, enable_gallery, image, files],
+              outputs=[results, label, gallery], api_name="inference")
 
 blocks.queue()
-blocks.launch(debug=True,inline=True)
+blocks.launch(debug=True, inline=True)

requirements.txt

@@ -1,4 +1,5 @@
 transformers
 gradio
 --extra-index-url https://download.pytorch.org/whl/cu113
-torch
+torch
+aiohttp