svjack commited on
Commit
66e03d0
·
1 Parent(s): bd8578f

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +94 -0
app.py ADDED
@@ -0,0 +1,94 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import os
3
+ import sys
4
+ from uuid import uuid1
5
+ from PIL import Image
6
+ from zipfile import ZipFile
7
+ import pathlib
8
+ import shutil
9
+ import pandas as pd
10
+ import deepsparse
11
+
12
+ rn50_embedding_pipeline = deepsparse.Pipeline.create(
13
+ task="embedding-extraction",
14
+ base_task="image-classification", # tells the pipeline to expect images and normalize input with ImageNet means/stds
15
+ model_path="zoo:cv/classification/resnet_v1-50/pytorch/sparseml/imagenet/channel20_pruned75_quant-none-vnni",
16
+ emb_extraction_layer=-3, # extracts last layer before projection head and softmax
17
+ )
18
+
19
+ def zip_ims(g):
20
+ from uuid import uuid1
21
+ if g is None:
22
+ return None
23
+ l = list(map(lambda x: x["name"], g))
24
+ if not l:
25
+ return None
26
+ zip_file_name ="tmp.zip"
27
+ with ZipFile(zip_file_name ,"w") as zipObj:
28
+ for ele in l:
29
+ zipObj.write(ele, "{}.png".format(uuid1()))
30
+ #zipObj.write(file2.name, "file2")
31
+ return zip_file_name
32
+
33
+ def unzip_ims(zip_file_name):
34
+ print("call file")
35
+ if zip_file_name is None:
36
+ return {}
37
+ unzip_path = "img_dir"
38
+ if os.path.exists(unzip_path):
39
+ shutil.rmtree(unzip_path)
40
+ with ZipFile(zip_file_name) as archive:
41
+ archive.extractall(unzip_path)
42
+ im_name_l = pd.Series(
43
+ list(pathlib.Path(unzip_path).rglob("*.png")) + \
44
+ list(pathlib.Path(unzip_path).rglob("*.jpg")) + \
45
+ list(pathlib.Path(unzip_path).rglob("*.jpeg"))
46
+ ).map(str).values.tolist()
47
+ embeddings = rn50_embedding_pipeline(images=im_name_l)
48
+ if os.path.exists(unzip_path):
49
+ shutil.rmtree(unzip_path)
50
+ return {
51
+ "names": im_name_l,
52
+ "embs": embeddings
53
+ }
54
+
55
+
56
+ def emb_img_func(im):
57
+ print("call im :")
58
+ if im is None:
59
+ return {}
60
+ im_obj = Image.fromarray(im)
61
+ im_name = "{}.png".format(uuid1())
62
+ im_obj.save(im_name)
63
+ embeddings = rn50_embedding_pipeline(images=[im_name])
64
+ os.remove(im_name)
65
+ return {
66
+ "names": [im_name],
67
+ "embs": embeddings
68
+ }
69
+
70
+ '''
71
+ def emb_gallery_func(gallery):
72
+ print("call ga :")
73
+ if gallery is None:
74
+ return []
75
+ im_name_l = list(map(lambda x: x["name"], images))
76
+ embeddings = rn50_embedding_pipeline(images=im_name_l)
77
+ return embeddings
78
+ '''
79
+
80
+ with gr.Blocks() as demo:
81
+ with gr.Row():
82
+ with gr.Column():
83
+ inputs_0 = gr.Image(label = "Input Image for embed")
84
+ button_0 = gr.Button("Image button")
85
+ with gr.Column():
86
+ inputs_1 = gr.File(label = "Input Images zip file for embed")
87
+ button_1 = gr.Button("Image File button")
88
+ with gr.Row():
89
+ outputs = gr.JSON(label = "Output Embeddings")
90
+
91
+ button_0.click(fn = emb_img_func, inputs = inputs_0, outputs = outputs)
92
+ button_1.click(fn = unzip_ims, inputs = inputs_1, outputs = outputs)
93
+
94
+ demo.launch("0.0.0.0")