Update app.py

app.py

@@ -12,16 +12,43 @@ from zipfile import ZipFile
 import pathlib
 import shutil
 import pandas as pd
-import json
 import deepsparse
 
-rn50_embedding_pipeline = deepsparse.Pipeline.create(
+rn50_embedding_pipeline_default = 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
+    #emb_extraction_layer=-1, # extracts last layer before projection head and softmax
 )
 
+rn50_embedding_pipeline_last_1 = 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=-1, # extracts last layer before projection head and softmax
+)
+
+rn50_embedding_pipeline_last_2 = 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=-2, # extracts last layer before projection head and softmax
+)
+
+rn50_embedding_pipeline_last_3 = 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
+)
+
+rn50_embedding_pipeline_dict = {
+    "0": rn50_embedding_pipeline_default,
+    "1": rn50_embedding_pipeline_last_1,
+    "2": rn50_embedding_pipeline_last_2,
+    "3": rn50_embedding_pipeline_last_3
+}
+
 def zip_ims(g):
     from uuid import uuid1
     if g is None:
@@ -36,7 +63,8 @@ def zip_ims(g):
         #zipObj.write(file2.name, "file2")
     return zip_file_name
 
-def unzip_ims(zip_file_name):
+def unzip_ims_func(zip_file_name, choose_model,
+    rn50_embedding_pipeline_dict = rn50_embedding_pipeline_dict):
     print("call file")
     if zip_file_name is None:
         return json.dumps({})
@@ -50,6 +78,7 @@ def unzip_ims(zip_file_name):
     list(pathlib.Path(unzip_path).rglob("*.jpg")) + \
     list(pathlib.Path(unzip_path).rglob("*.jpeg"))
     ).map(str).values.tolist()
+    rn50_embedding_pipeline = rn50_embedding_pipeline_dict[choose_model]
     embeddings = rn50_embedding_pipeline(images=im_name_l)
     if os.path.exists(unzip_path):
         shutil.rmtree(unzip_path)
@@ -60,13 +89,15 @@ def unzip_ims(zip_file_name):
     })
 
 
-def emb_img_func(im):
+def emb_img_func(im, choose_model,
+    rn50_embedding_pipeline_dict = rn50_embedding_pipeline_dict):
     print("call im :")
     if im is None:
         return json.dumps({})
     im_obj = Image.fromarray(im)
     im_name = "{}.png".format(uuid1())
     im_obj.save(im_name)
+    rn50_embedding_pipeline = rn50_embedding_pipeline_dict[choose_model]
     embeddings = rn50_embedding_pipeline(images=[im_name])
     os.remove(im_name)
     return json.dumps({
@@ -85,6 +116,9 @@ def emb_gallery_func(gallery):
 '''
 
 with gr.Blocks() as demo:
+    with gr.Row():
+        choose_model = gr.Radio(choices=["0", "1", "2", "3"],
+        value="0", label="Choose embedding layer", elem_id="layer_radio")
     with gr.Row():
         with gr.Column():
             inputs_0 = gr.Image(label = "Input Image for embed")
@@ -96,6 +130,6 @@ with gr.Blocks() as demo:
         outputs = gr.Text(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)
+    button_1.click(fn = unzip_ims_func, inputs = inputs_1, outputs = outputs)
 
 demo.launch("0.0.0.0")