Create app.py

app.py

@@ -0,0 +1,94 @@
+import gradio as gr
+import os
+import sys
+from uuid import uuid1
+from PIL import Image
+from zipfile import ZipFile
+import pathlib
+import shutil
+import pandas as pd
+import deepsparse
+
+rn50_embedding_pipeline = deepsparse.Pipeline.create(
+    task="embedding-extraction",
+    base_task="image-classification", # tells the pipeline to expect images and normalize input with ImageNet means/stds
+    model_path="zoo:cv/classification/resnet_v1-50/pytorch/sparseml/imagenet/channel20_pruned75_quant-none-vnni",
+    emb_extraction_layer=-3, # extracts last layer before projection head and softmax
+)
+
+def zip_ims(g):
+    from uuid import uuid1
+    if g is None:
+        return None
+    l = list(map(lambda x: x["name"], g))
+    if not l:
+        return None
+    zip_file_name ="tmp.zip"
+    with ZipFile(zip_file_name ,"w") as zipObj:
+        for ele in l:
+            zipObj.write(ele, "{}.png".format(uuid1()))
+        #zipObj.write(file2.name, "file2")
+    return zip_file_name
+
+def unzip_ims(zip_file_name):
+    print("call file")
+    if zip_file_name is None:
+        return {}
+    unzip_path = "img_dir"
+    if os.path.exists(unzip_path):
+        shutil.rmtree(unzip_path)
+    with ZipFile(zip_file_name) as archive:
+        archive.extractall(unzip_path)
+    im_name_l = pd.Series(
+    list(pathlib.Path(unzip_path).rglob("*.png")) + \
+    list(pathlib.Path(unzip_path).rglob("*.jpg")) + \
+    list(pathlib.Path(unzip_path).rglob("*.jpeg"))
+    ).map(str).values.tolist()
+    embeddings = rn50_embedding_pipeline(images=im_name_l)
+    if os.path.exists(unzip_path):
+        shutil.rmtree(unzip_path)
+    return {
+        "names": im_name_l,
+        "embs": embeddings
+    }
+
+
+def emb_img_func(im):
+    print("call im :")
+    if im is None:
+        return {}
+    im_obj = Image.fromarray(im)
+    im_name = "{}.png".format(uuid1())
+    im_obj.save(im_name)
+    embeddings = rn50_embedding_pipeline(images=[im_name])
+    os.remove(im_name)
+    return {
+        "names": [im_name],
+        "embs": embeddings
+    }
+
+'''
+def emb_gallery_func(gallery):
+    print("call ga :")
+    if gallery is None:
+        return []
+    im_name_l = list(map(lambda x: x["name"], images))
+    embeddings = rn50_embedding_pipeline(images=im_name_l)
+    return embeddings
+'''
+
+with gr.Blocks() as demo:
+    with gr.Row():
+        with gr.Column():
+            inputs_0 = gr.Image(label = "Input Image for embed")
+            button_0 = gr.Button("Image button")
+        with gr.Column():
+            inputs_1 = gr.File(label = "Input Images zip file for embed")
+            button_1 = gr.Button("Image File button")
+    with gr.Row():
+        outputs = gr.JSON(label = "Output Embeddings")
+
+    button_0.click(fn = emb_img_func, inputs = inputs_0, outputs = outputs)
+    button_1.click(fn = unzip_ims, inputs = inputs_1, outputs = outputs)
+
+demo.launch("0.0.0.0")