Update app.py
Browse files
app.py
CHANGED
@@ -8,18 +8,18 @@ import torch
|
|
8 |
import logging
|
9 |
from os import environ
|
10 |
from transformers import OwlViTProcessor, OwlViTForObjectDetection
|
11 |
-
|
12 |
from myscaledb import Client
|
13 |
from classifier import Classifier, prompt2vec, tune, SplitLayer
|
14 |
from query_model import simple_query, topk_obj_query, rev_query
|
15 |
from card_model import card, obj_card, style
|
16 |
from box_utils import postprocess
|
17 |
|
18 |
-
environ[
|
19 |
|
20 |
OBJ_DB_NAME = "mqdb_demo.coco_owl_vit_b_32_objects"
|
21 |
IMG_DB_NAME = "mqdb_demo.coco_owl_vit_b_32_images"
|
22 |
-
MODEL_ID =
|
23 |
DIMS = 512
|
24 |
|
25 |
qtime = 0
|
@@ -34,9 +34,9 @@ def build_model(name="google/owlvit-base-patch32"):
|
|
34 |
Returns:
|
35 |
(model, processor): OwlViT model and its processor for both image and text
|
36 |
"""
|
37 |
-
device =
|
38 |
if torch.cuda.is_available():
|
39 |
-
device =
|
40 |
model = OwlViTForObjectDetection.from_pretrained(name).to(device)
|
41 |
processor = OwlViTProcessor.from_pretrained(name)
|
42 |
return model, processor
|
@@ -44,7 +44,7 @@ def build_model(name="google/owlvit-base-patch32"):
|
|
44 |
|
45 |
@st.experimental_singleton(show_spinner=False)
|
46 |
def init_owlvit():
|
47 |
-
"""
|
48 |
|
49 |
Returns:
|
50 |
model, processor
|
@@ -55,7 +55,7 @@ def init_owlvit():
|
|
55 |
|
56 |
@st.experimental_singleton(show_spinner=False)
|
57 |
def init_db():
|
58 |
-
"""
|
59 |
|
60 |
Returns:
|
61 |
meta_field: Meta field that records if an image is viewed or not
|
@@ -63,15 +63,15 @@ def init_db():
|
|
63 |
"""
|
64 |
meta = []
|
65 |
client = Client(
|
66 |
-
url=st.secrets["DB_URL"], user=st.secrets["USER"], password=st.secrets["PASSWD"]
|
|
|
67 |
# We can check if the connection is alive
|
68 |
assert client.is_alive()
|
69 |
return meta, client
|
70 |
|
71 |
|
72 |
def refresh_index():
|
73 |
-
"""
|
74 |
-
"""
|
75 |
del st.session_state["meta"]
|
76 |
st.session_state.meta = []
|
77 |
st.session_state.query_num = 0
|
@@ -80,16 +80,16 @@ def refresh_index():
|
|
80 |
init_db.clear()
|
81 |
# refresh session states
|
82 |
st.session_state.meta, st.session_state.index = init_db()
|
83 |
-
if
|
84 |
del st.session_state.clf
|
85 |
-
if
|
86 |
del st.session_state.xq
|
87 |
-
if
|
88 |
del st.session_state.topk_img_id
|
89 |
|
90 |
|
91 |
def query(xq, exclude_list=None):
|
92 |
-
"""
|
93 |
|
94 |
In this part, we will retrieve A LOT OF data from the server,
|
95 |
including TopK boxes and their embeddings, the counterpart of non-TopK boxes in TopK images.
|
@@ -98,7 +98,7 @@ def query(xq, exclude_list=None):
|
|
98 |
xq (numpy.ndarray or list of floats): Query vector
|
99 |
|
100 |
Returns:
|
101 |
-
matches: list of Records object. Keys referrring to selected columns group by images.
|
102 |
Exclude the user's viewlist.
|
103 |
img_matches: list of Records object. Containing other non-TopK but hit objects among TopK images.
|
104 |
side_matches: list of Records object. Containing REAL TopK objects disregard the user's view history
|
@@ -112,27 +112,47 @@ def query(xq, exclude_list=None):
|
|
112 |
while attempt < 3:
|
113 |
try:
|
114 |
matches = topk_obj_query(
|
115 |
-
st.session_state.index,
|
116 |
-
|
117 |
-
|
118 |
-
|
|
|
|
|
|
|
|
|
|
|
119 |
st.session_state.topk_img_id = img_ids
|
120 |
status_bar[0].write("Retrieving TopK Images...")
|
121 |
pbar.progress(25)
|
122 |
o_matches = rev_query(
|
123 |
-
st.session_state.index,
|
124 |
-
|
|
|
|
|
|
|
|
|
|
|
125 |
status_bar[0].write("Retrieving TopKs Objects...")
|
126 |
pbar.progress(50)
|
127 |
-
side_matches = simple_query(
|
128 |
-
|
129 |
-
|
130 |
-
|
|
|
|
|
|
|
|
|
|
|
131 |
pbar.progress(75)
|
132 |
if len(img_ids) > 0:
|
133 |
img_matches = rev_query(
|
134 |
-
st.session_state.index,
|
135 |
-
|
|
|
|
|
|
|
|
|
|
|
136 |
else:
|
137 |
img_matches = []
|
138 |
status_bar[0].write("DONE!")
|
@@ -163,22 +183,31 @@ def init_random_query():
|
|
163 |
|
164 |
|
165 |
def submit(meta):
|
166 |
-
"""
|
167 |
-
"""
|
168 |
# Only updating the meta if the train button is pressed
|
169 |
st.session_state.meta.extend(meta)
|
170 |
st.session_state.step += 1
|
171 |
matches = st.session_state.matched_boxes
|
172 |
-
X, y = list(
|
173 |
-
|
174 |
-
|
175 |
-
|
176 |
-
|
177 |
-
|
178 |
-
|
179 |
-
|
180 |
-
|
181 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
182 |
|
183 |
|
184 |
# st.set_page_config(layout="wide")
|
@@ -186,210 +215,262 @@ def submit(meta):
|
|
186 |
# Boxes are drawn in SVGs.
|
187 |
st.write(style(), unsafe_allow_html=True)
|
188 |
|
189 |
-
|
190 |
-
|
191 |
-
|
192 |
-
|
193 |
-
|
194 |
-
|
195 |
-
|
196 |
-
|
197 |
-
|
198 |
-
|
199 |
-
st.
|
200 |
-
|
201 |
-
|
202 |
-
|
203 |
-
|
204 |
-
|
205 |
-
|
206 |
-
|
207 |
-
|
208 |
-
|
209 |
-
|
210 |
-
|
211 |
-
|
212 |
-
|
213 |
-
|
214 |
-
|
215 |
-
|
216 |
-
|
217 |
-
|
218 |
-
|
219 |
-
|
220 |
-
|
221 |
-
|
222 |
-
|
223 |
-
|
224 |
-
|
225 |
-
|
226 |
-
|
227 |
-
|
228 |
-
|
229 |
-
|
230 |
-
|
231 |
-
|
232 |
-
|
233 |
-
|
234 |
-
|
235 |
-
|
236 |
-
|
237 |
-
|
238 |
-
|
239 |
-
|
240 |
-
|
241 |
-
|
242 |
-
|
243 |
-
|
244 |
-
|
245 |
-
|
246 |
-
|
247 |
-
if upld_model is not None:
|
248 |
-
import onnx
|
249 |
-
from onnx import numpy_helper
|
250 |
-
_model = onnx.load(upld_model)
|
251 |
-
st.session_state.text_prompts = [
|
252 |
-
node.name for node in _model.graph.output] + ['none']
|
253 |
-
weights = _model.graph.initializer
|
254 |
-
xq = numpy_helper.to_array(weights[0]).T
|
255 |
-
assert xq.shape[0] == len(
|
256 |
-
st.session_state.text_prompts)-1 and xq.shape[1] == DIMS
|
257 |
-
st.session_state.xq = xq
|
258 |
-
_ = [elem.empty() for elem in start]
|
259 |
-
else:
|
260 |
-
logging.info(f"Input prompt is {prompt}")
|
261 |
-
st.session_state.text_prompts = prompt.split(',') + ['none']
|
262 |
-
input_ids, xq = prompt2vec(
|
263 |
-
st.session_state.text_prompts[:-1], model, tokenizer)
|
264 |
st.session_state.xq = xq
|
|
|
|
|
265 |
_ = [elem.empty() for elem in start]
|
266 |
-
|
267 |
-
|
268 |
-
|
269 |
-
|
270 |
-
|
271 |
-
st.session_state.
|
272 |
-
|
273 |
-
|
274 |
-
|
275 |
-
|
276 |
-
if
|
277 |
-
|
278 |
-
|
279 |
-
|
280 |
-
|
281 |
-
|
282 |
-
|
283 |
-
|
284 |
-
|
285 |
-
|
286 |
-
|
287 |
-
|
288 |
-
|
289 |
-
|
290 |
-
|
291 |
-
|
292 |
-
|
293 |
-
|
294 |
-
|
295 |
-
|
296 |
-
|
297 |
-
|
298 |
-
|
299 |
-
|
300 |
-
|
301 |
-
|
302 |
-
|
303 |
-
|
304 |
-
|
305 |
-
|
306 |
-
|
307 |
-
|
308 |
-
|
309 |
-
|
310 |
-
|
311 |
-
|
312 |
-
|
313 |
-
|
314 |
-
|
315 |
-
|
316 |
-
|
317 |
-
|
318 |
-
|
319 |
-
|
320 |
-
|
321 |
-
|
322 |
-
|
323 |
-
|
324 |
-
|
325 |
-
|
326 |
-
|
327 |
-
|
328 |
-
|
329 |
-
|
330 |
-
|
331 |
-
|
332 |
-
|
333 |
-
|
334 |
-
|
335 |
-
|
336 |
-
|
337 |
-
|
338 |
-
|
339 |
-
|
340 |
-
|
341 |
-
|
342 |
-
|
343 |
-
|
344 |
-
|
345 |
-
|
346 |
-
|
347 |
-
|
348 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
349 |
with st.container():
|
350 |
-
|
351 |
-
|
352 |
-
|
353 |
-
boxes_w_img
|
354 |
-
img_matches)
|
355 |
-
|
356 |
-
# Sort the result according to their relavancy
|
357 |
-
boxes_w_img = sorted(
|
358 |
-
boxes_w_img, key=lambda x: x[4], reverse=True)
|
359 |
-
|
360 |
-
st.session_state.matched_boxes = {}
|
361 |
-
# For each images in the retrieved images, DISPLAY
|
362 |
for img_id, img_url, img_w, img_h, img_score, boxes in boxes_w_img:
|
363 |
-
|
364 |
-
# prepare inputs for training
|
365 |
-
st.session_state.matched_boxes.update(
|
366 |
-
{b[0]: b for b in boxes})
|
367 |
args = img_url, img_w, img_h, boxes
|
368 |
-
|
369 |
-
|
370 |
-
|
371 |
-
|
372 |
-
|
373 |
-
|
374 |
-
|
375 |
-
|
376 |
-
|
377 |
-
|
378 |
-
|
379 |
-
|
380 |
-
|
381 |
-
|
382 |
-
|
383 |
-
|
384 |
-
|
385 |
-
|
386 |
-
|
387 |
-
|
388 |
-
|
389 |
-
|
390 |
-
|
391 |
-
|
392 |
-
|
393 |
-
|
394 |
-
|
395 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
8 |
import logging
|
9 |
from os import environ
|
10 |
from transformers import OwlViTProcessor, OwlViTForObjectDetection
|
11 |
+
from bot import Bot, Message
|
12 |
from myscaledb import Client
|
13 |
from classifier import Classifier, prompt2vec, tune, SplitLayer
|
14 |
from query_model import simple_query, topk_obj_query, rev_query
|
15 |
from card_model import card, obj_card, style
|
16 |
from box_utils import postprocess
|
17 |
|
18 |
+
environ["TOKENIZERS_PARALLELISM"] = "true"
|
19 |
|
20 |
OBJ_DB_NAME = "mqdb_demo.coco_owl_vit_b_32_objects"
|
21 |
IMG_DB_NAME = "mqdb_demo.coco_owl_vit_b_32_images"
|
22 |
+
MODEL_ID = "google/owlvit-base-patch32"
|
23 |
DIMS = 512
|
24 |
|
25 |
qtime = 0
|
|
|
34 |
Returns:
|
35 |
(model, processor): OwlViT model and its processor for both image and text
|
36 |
"""
|
37 |
+
device = "cpu"
|
38 |
if torch.cuda.is_available():
|
39 |
+
device = "cuda"
|
40 |
model = OwlViTForObjectDetection.from_pretrained(name).to(device)
|
41 |
processor = OwlViTProcessor.from_pretrained(name)
|
42 |
return model, processor
|
|
|
44 |
|
45 |
@st.experimental_singleton(show_spinner=False)
|
46 |
def init_owlvit():
|
47 |
+
"""Initialize OwlViT Model
|
48 |
|
49 |
Returns:
|
50 |
model, processor
|
|
|
55 |
|
56 |
@st.experimental_singleton(show_spinner=False)
|
57 |
def init_db():
|
58 |
+
"""Initialize the Database Connection
|
59 |
|
60 |
Returns:
|
61 |
meta_field: Meta field that records if an image is viewed or not
|
|
|
63 |
"""
|
64 |
meta = []
|
65 |
client = Client(
|
66 |
+
url=st.secrets["DB_URL"], user=st.secrets["USER"], password=st.secrets["PASSWD"]
|
67 |
+
)
|
68 |
# We can check if the connection is alive
|
69 |
assert client.is_alive()
|
70 |
return meta, client
|
71 |
|
72 |
|
73 |
def refresh_index():
|
74 |
+
"""Clean the session"""
|
|
|
75 |
del st.session_state["meta"]
|
76 |
st.session_state.meta = []
|
77 |
st.session_state.query_num = 0
|
|
|
80 |
init_db.clear()
|
81 |
# refresh session states
|
82 |
st.session_state.meta, st.session_state.index = init_db()
|
83 |
+
if "clf" in st.session_state:
|
84 |
del st.session_state.clf
|
85 |
+
if "xq" in st.session_state:
|
86 |
del st.session_state.xq
|
87 |
+
if "topk_img_id" in st.session_state:
|
88 |
del st.session_state.topk_img_id
|
89 |
|
90 |
|
91 |
def query(xq, exclude_list=None):
|
92 |
+
"""Query matched w.r.t a given vector
|
93 |
|
94 |
In this part, we will retrieve A LOT OF data from the server,
|
95 |
including TopK boxes and their embeddings, the counterpart of non-TopK boxes in TopK images.
|
|
|
98 |
xq (numpy.ndarray or list of floats): Query vector
|
99 |
|
100 |
Returns:
|
101 |
+
matches: list of Records object. Keys referrring to selected columns group by images.
|
102 |
Exclude the user's viewlist.
|
103 |
img_matches: list of Records object. Containing other non-TopK but hit objects among TopK images.
|
104 |
side_matches: list of Records object. Containing REAL TopK objects disregard the user's view history
|
|
|
112 |
while attempt < 3:
|
113 |
try:
|
114 |
matches = topk_obj_query(
|
115 |
+
st.session_state.index,
|
116 |
+
xq,
|
117 |
+
IMG_DB_NAME,
|
118 |
+
OBJ_DB_NAME,
|
119 |
+
exclude_list=exclude_list,
|
120 |
+
topk=5000,
|
121 |
+
)
|
122 |
+
img_ids = [r["img_id"] for r in matches]
|
123 |
+
if "topk_img_id" not in st.session_state:
|
124 |
st.session_state.topk_img_id = img_ids
|
125 |
status_bar[0].write("Retrieving TopK Images...")
|
126 |
pbar.progress(25)
|
127 |
o_matches = rev_query(
|
128 |
+
st.session_state.index,
|
129 |
+
xq,
|
130 |
+
st.session_state.topk_img_id,
|
131 |
+
IMG_DB_NAME,
|
132 |
+
OBJ_DB_NAME,
|
133 |
+
thresh=0.1,
|
134 |
+
)
|
135 |
status_bar[0].write("Retrieving TopKs Objects...")
|
136 |
pbar.progress(50)
|
137 |
+
side_matches = simple_query(
|
138 |
+
st.session_state.index,
|
139 |
+
xq,
|
140 |
+
IMG_DB_NAME,
|
141 |
+
OBJ_DB_NAME,
|
142 |
+
thresh=-1,
|
143 |
+
topk=5000,
|
144 |
+
)
|
145 |
+
status_bar[0].write("Retrieving Non-TopK in Another TopK Images...")
|
146 |
pbar.progress(75)
|
147 |
if len(img_ids) > 0:
|
148 |
img_matches = rev_query(
|
149 |
+
st.session_state.index,
|
150 |
+
xq,
|
151 |
+
img_ids,
|
152 |
+
IMG_DB_NAME,
|
153 |
+
OBJ_DB_NAME,
|
154 |
+
thresh=0.1,
|
155 |
+
)
|
156 |
else:
|
157 |
img_matches = []
|
158 |
status_bar[0].write("DONE!")
|
|
|
183 |
|
184 |
|
185 |
def submit(meta):
|
186 |
+
"""Tune the model w.r.t given score from user."""
|
|
|
187 |
# Only updating the meta if the train button is pressed
|
188 |
st.session_state.meta.extend(meta)
|
189 |
st.session_state.step += 1
|
190 |
matches = st.session_state.matched_boxes
|
191 |
+
X, y = list(
|
192 |
+
zip(
|
193 |
+
*(
|
194 |
+
(
|
195 |
+
v[-1],
|
196 |
+
st.session_state.text_prompts.index(st.session_state[f"label-{i}"]),
|
197 |
+
)
|
198 |
+
for i, v in matches.items()
|
199 |
+
)
|
200 |
+
)
|
201 |
+
)
|
202 |
+
st.session_state.xq = tune(
|
203 |
+
st.session_state.clf, X, y, iters=int(st.session_state.iters)
|
204 |
+
)
|
205 |
+
(
|
206 |
+
st.session_state.matches,
|
207 |
+
st.session_state.img_matches,
|
208 |
+
st.session_state.side_matches,
|
209 |
+
st.session_state.o_matches,
|
210 |
+
) = query(st.session_state.xq, st.session_state.meta)
|
211 |
|
212 |
|
213 |
# st.set_page_config(layout="wide")
|
|
|
215 |
# Boxes are drawn in SVGs.
|
216 |
st.write(style(), unsafe_allow_html=True)
|
217 |
|
218 |
+
bot = Bot(app_name="HF OwlViT", enabled=True, bot_key=st.secrets['BOT_KEY'])
|
219 |
+
try:
|
220 |
+
with st.spinner("Connecting DB..."):
|
221 |
+
st.session_state.meta, st.session_state.index = init_db()
|
222 |
+
|
223 |
+
with st.spinner("Loading Models..."):
|
224 |
+
# Initialize model
|
225 |
+
model, tokenizer = init_owlvit()
|
226 |
+
# If its a fresh start... (query not set)
|
227 |
+
if "xq" not in st.session_state:
|
228 |
+
with st.container():
|
229 |
+
st.title("Object Detection Safari")
|
230 |
+
start = [st.empty() for _ in range(8)]
|
231 |
+
start[0].info(
|
232 |
+
"""
|
233 |
+
We extracted boxes from **287,104** images in COCO Dataset, including its train / val / test /
|
234 |
+
unlabeled images, collecting **165,371,904 boxes** which are then filtered with common prompts.
|
235 |
+
You can search with almost any words or phrases you can think of. Please enjoy your journey of
|
236 |
+
an adventure to COCO.
|
237 |
+
"""
|
238 |
+
)
|
239 |
+
prompt = start[1].text_input(
|
240 |
+
"Prompt:",
|
241 |
+
value="",
|
242 |
+
placeholder="Examples: football, billboard, stop sign, watermark ...",
|
243 |
+
)
|
244 |
+
with start[2].container():
|
245 |
+
st.write(
|
246 |
+
"You can search with multiple keywords. Plese separate with commas but with no space."
|
247 |
+
)
|
248 |
+
st.write("For example: `cat,dog,tree`")
|
249 |
+
st.markdown(
|
250 |
+
"""
|
251 |
+
<p style="color:gray;"> Don\'t know what to search? Try <b>Random</b>!</p>
|
252 |
+
""",
|
253 |
+
unsafe_allow_html=True,
|
254 |
+
)
|
255 |
+
|
256 |
+
upld_model = start[4].file_uploader(
|
257 |
+
"Or you can upload your previous run!", type="onnx"
|
258 |
+
)
|
259 |
+
upld_btn = start[5].button(
|
260 |
+
"Use Loaded Weights", disabled=upld_model is None, on_click=refresh_index
|
261 |
+
)
|
262 |
+
|
263 |
+
with start[3]:
|
264 |
+
col = st.columns(8)
|
265 |
+
has_no_prompt = len(prompt) == 0 and upld_model is None
|
266 |
+
prompt_xq = col[6].button(
|
267 |
+
"Prompt", disabled=len(prompt) == 0, on_click=refresh_index
|
268 |
+
)
|
269 |
+
random_xq = col[7].button(
|
270 |
+
"Random", disabled=not has_no_prompt, on_click=refresh_index
|
271 |
+
)
|
272 |
+
matches = []
|
273 |
+
img_matches = []
|
274 |
+
if random_xq:
|
275 |
+
xq = init_random_query()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
276 |
st.session_state.xq = xq
|
277 |
+
prompt = "unknown"
|
278 |
+
st.session_state.text_prompts = prompt.split(",") + ["none"]
|
279 |
_ = [elem.empty() for elem in start]
|
280 |
+
t0 = time()
|
281 |
+
(
|
282 |
+
st.session_state.matches,
|
283 |
+
st.session_state.img_matches,
|
284 |
+
st.session_state.side_matches,
|
285 |
+
st.session_state.o_matches,
|
286 |
+
) = query(st.session_state.xq, st.session_state.meta)
|
287 |
+
t1 = time()
|
288 |
+
qtime = (t1 - t0) * 1000
|
289 |
+
elif prompt_xq or upld_btn:
|
290 |
+
if upld_model is not None:
|
291 |
+
import onnx
|
292 |
+
from onnx import numpy_helper
|
293 |
+
|
294 |
+
_model = onnx.load(upld_model)
|
295 |
+
st.session_state.text_prompts = [
|
296 |
+
node.name for node in _model.graph.output
|
297 |
+
] + ["none"]
|
298 |
+
weights = _model.graph.initializer
|
299 |
+
xq = numpy_helper.to_array(weights[0]).T
|
300 |
+
assert (
|
301 |
+
xq.shape[0] == len(st.session_state.text_prompts) - 1
|
302 |
+
and xq.shape[1] == DIMS
|
303 |
+
)
|
304 |
+
st.session_state.xq = xq
|
305 |
+
_ = [elem.empty() for elem in start]
|
306 |
+
else:
|
307 |
+
logging.info(f"Input prompt is {prompt}")
|
308 |
+
st.session_state.text_prompts = prompt.split(",") + ["none"]
|
309 |
+
input_ids, xq = prompt2vec(
|
310 |
+
st.session_state.text_prompts[:-1], model, tokenizer
|
311 |
+
)
|
312 |
+
st.session_state.xq = xq
|
313 |
+
_ = [elem.empty() for elem in start]
|
314 |
+
t0 = time()
|
315 |
+
(
|
316 |
+
st.session_state.matches,
|
317 |
+
st.session_state.img_matches,
|
318 |
+
st.session_state.side_matches,
|
319 |
+
st.session_state.o_matches,
|
320 |
+
) = query(st.session_state.xq, st.session_state.meta)
|
321 |
+
t1 = time()
|
322 |
+
qtime = (t1 - t0) * 1000
|
323 |
+
|
324 |
+
# If its not a fresh start (query is set)
|
325 |
+
if "xq" in st.session_state:
|
326 |
+
o_matches = st.session_state.o_matches
|
327 |
+
side_matches = st.session_state.side_matches
|
328 |
+
img_matches = st.session_state.img_matches
|
329 |
+
matches = st.session_state.matches
|
330 |
+
# initialize classifier
|
331 |
+
if "clf" not in st.session_state:
|
332 |
+
st.session_state.clf = Classifier(st.session_state.xq)
|
333 |
+
st.session_state.step = 0
|
334 |
+
if qtime > 0:
|
335 |
+
st.info(
|
336 |
+
"Query done in {0:.2f} ms and returned {1:d} images with {2:d} boxes".format(
|
337 |
+
qtime,
|
338 |
+
len(matches),
|
339 |
+
sum(
|
340 |
+
[
|
341 |
+
len(m["box_id"]) + len(im["box_id"])
|
342 |
+
for m, im in zip(matches, img_matches)
|
343 |
+
]
|
344 |
+
),
|
345 |
+
)
|
346 |
+
)
|
347 |
+
|
348 |
+
# export the model into executable ONNX
|
349 |
+
st.session_state.dnld_model = BytesIO()
|
350 |
+
torch.onnx.export(
|
351 |
+
torch.nn.Sequential(st.session_state.clf.model, SplitLayer()),
|
352 |
+
torch.zeros([1, len(st.session_state.xq[0])]),
|
353 |
+
st.session_state.dnld_model,
|
354 |
+
input_names=["input"],
|
355 |
+
output_names=st.session_state.text_prompts[:-1],
|
356 |
+
)
|
357 |
+
|
358 |
+
dnld_nam = st.text_input(
|
359 |
+
"Download Name:",
|
360 |
+
f'{("_".join([i.replace(" ", "-") for i in st.session_state.text_prompts[:-1]]) if "text_prompts" in st.session_state else "model")}.onnx',
|
361 |
+
max_chars=50,
|
362 |
+
)
|
363 |
+
dnld_btn = st.download_button(
|
364 |
+
"Download your classifier!", st.session_state.dnld_model, dnld_nam
|
365 |
+
)
|
366 |
+
# build up a sidebar to display REAL TopK in DB
|
367 |
+
# this will change during user's finetune. But sometime it would lead to bad results
|
368 |
+
side_bar_len = min(240 // len(st.session_state.text_prompts), 120)
|
369 |
+
with st.sidebar:
|
370 |
+
with st.expander("Top-K Images"):
|
371 |
with st.container():
|
372 |
+
boxes_w_img, _ = postprocess(
|
373 |
+
o_matches, st.session_state.text_prompts, None
|
374 |
+
)
|
375 |
+
boxes_w_img = sorted(boxes_w_img, key=lambda x: x[4], reverse=True)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
376 |
for img_id, img_url, img_w, img_h, img_score, boxes in boxes_w_img:
|
|
|
|
|
|
|
|
|
377 |
args = img_url, img_w, img_h, boxes
|
378 |
+
st.write(card(*args), unsafe_allow_html=True)
|
379 |
+
|
380 |
+
with st.expander("Top-K Objects", expanded=True):
|
381 |
+
side_cols = st.columns(len(st.session_state.text_prompts[:-1]))
|
382 |
+
for _cols, m in zip(side_cols, side_matches):
|
383 |
+
with _cols.container():
|
384 |
+
for cx, cy, w, h, logit, img_url, img_w, img_h in zip(
|
385 |
+
m["cx"],
|
386 |
+
m["cy"],
|
387 |
+
m["w"],
|
388 |
+
m["h"],
|
389 |
+
m["logit"],
|
390 |
+
m["img_url"],
|
391 |
+
m["img_w"],
|
392 |
+
m["img_h"],
|
393 |
+
):
|
394 |
+
st.write(
|
395 |
+
"{:s}: {:.4f}".format(
|
396 |
+
st.session_state.text_prompts[m["label"]], logit
|
397 |
+
)
|
398 |
+
)
|
399 |
+
_html = obj_card(
|
400 |
+
img_url, img_w, img_h, cx, cy, w, h, dst_len=side_bar_len
|
401 |
+
)
|
402 |
+
components.html(_html, side_bar_len, side_bar_len)
|
403 |
+
with st.container():
|
404 |
+
# Here let the user interact with batch labeling
|
405 |
+
with st.form("batch", clear_on_submit=False):
|
406 |
+
col = st.columns([1, 9])
|
407 |
+
|
408 |
+
# If there is nothing to show about
|
409 |
+
if len(matches) <= 0:
|
410 |
+
st.warning(
|
411 |
+
"Oops! We didn't find anything relevant to your query! Pleas try another one :/"
|
412 |
+
)
|
413 |
+
else:
|
414 |
+
st.session_state.iters = st.slider(
|
415 |
+
"Number of Iterations to Update",
|
416 |
+
min_value=0,
|
417 |
+
max_value=10,
|
418 |
+
step=1,
|
419 |
+
value=2,
|
420 |
+
)
|
421 |
+
# No matter what happened the user wants a way back
|
422 |
+
col[1].form_submit_button("Choose a new prompt", on_click=refresh_index)
|
423 |
+
|
424 |
+
# If there are things to show
|
425 |
+
if len(matches) > 0:
|
426 |
+
with st.container():
|
427 |
+
prompt_labels = st.session_state.text_prompts
|
428 |
+
|
429 |
+
# Post processing boxes regarding to their score, intersection
|
430 |
+
boxes_w_img, meta = postprocess(
|
431 |
+
matches, st.session_state.text_prompts, img_matches
|
432 |
+
)
|
433 |
+
|
434 |
+
# Sort the result according to their relavancy
|
435 |
+
boxes_w_img = sorted(boxes_w_img, key=lambda x: x[4], reverse=True)
|
436 |
+
|
437 |
+
st.session_state.matched_boxes = {}
|
438 |
+
# For each images in the retrieved images, DISPLAY
|
439 |
+
for img_id, img_url, img_w, img_h, img_score, boxes in boxes_w_img:
|
440 |
+
|
441 |
+
# prepare inputs for training
|
442 |
+
st.session_state.matched_boxes.update({b[0]: b for b in boxes})
|
443 |
+
args = img_url, img_w, img_h, boxes
|
444 |
+
|
445 |
+
# display boxes
|
446 |
+
with st.expander(
|
447 |
+
"{:s}: {:.4f}".format(img_id, img_score), expanded=True
|
448 |
+
):
|
449 |
+
ind_b = 0
|
450 |
+
# 4 columns: (img, obj, obj, obj)
|
451 |
+
img_row = st.columns([4, 2, 2, 2])
|
452 |
+
img_row[0].write(card(*args), unsafe_allow_html=True)
|
453 |
+
# crop objects out of the original image
|
454 |
+
for b in boxes:
|
455 |
+
_id, cx, cy, w, h, label, logit, is_selected, _ = b
|
456 |
+
with img_row[1 + ind_b % 3].container():
|
457 |
+
st.write("{:s}: {:.4f}".format(label, logit))
|
458 |
+
# quite hacky: with streamlit components API
|
459 |
+
_html = obj_card(
|
460 |
+
img_url, img_w, img_h, *b[1:5], dst_len=120
|
461 |
+
)
|
462 |
+
components.html(_html, 120, 120)
|
463 |
+
# the user will choose the right label of the given object
|
464 |
+
st.selectbox(
|
465 |
+
"Class",
|
466 |
+
prompt_labels,
|
467 |
+
index=prompt_labels.index(label),
|
468 |
+
key=f"label-{_id}",
|
469 |
+
)
|
470 |
+
ind_b += 1
|
471 |
+
col[0].form_submit_button("Train!", on_click=lambda: submit(meta))
|
472 |
+
except Exception as e:
|
473 |
+
msg = Message()
|
474 |
+
msg.content = str(e.with_traceback(None))
|
475 |
+
msg.type_hint = str(type(e).__name__)
|
476 |
+
bot.incident(msg)
|