Upload 2 files

app.py

@@ -15,9 +15,9 @@ def run_vip(
     n_samples_init,
     n_samples_opt,
     n_iters,
-    n_recurssion,
+    n_parallel_trials,
     openai_api_key,
-    progress=gr.Progress(track_tqdm=True),
+    progress=gr.Progress(track_tqdm=False),
 ):
 
   if not openai_api_key:
@@ -53,7 +53,7 @@ def run_vip(
   }
 
   vlm = GPT4V(openai_api_key=openai_api_key)
-  ims, center, _ = vip_runner(
+  vip_gen = vip_runner(
       vlm,
       im,
       query,
@@ -62,9 +62,10 @@ def run_vip(
       n_samples_init=n_samples_init,
       n_samples_opt=n_samples_opt,
       n_iters=n_iters,
-      recursion_level=n_recurssion,
+      n_parallel_trials=n_parallel_trials,
   )
-  return ims, f'Final selected coordinate: {np.round(center, decimals=0)}'
+  for rst in vip_gen:
+    yield rst 
 
 
 examples = [
@@ -117,11 +118,11 @@ The Info textbox will show the final selected pixel coordinate that PIVOT conver
 """.strip())
 
   gr.Markdown(
-      '## Example Images and Queries\n Drag images into the image box below'
+      '## Example Images and Queries\n Drag images into the image box below (Try safari on Mac if dragging does not work)'
   )
   with gr.Row(equal_height=True):
     for example in examples:
-      gr.Image(value=example['im_path'], label=example['desc'])
+      gr.Image(value=example['im_path'], type='numpy', label=example['desc'])
 
   gr.Markdown('## New Query')
   with gr.Row():
@@ -160,8 +161,8 @@ The Info textbox will show the final selected pixel coordinate that PIVOT conver
         inp_n_iters = gr.Slider(
             label='N Iterations', minimum=1, maximum=5, value=3, step=1
         )
-        inp_n_recurssions = gr.Slider(
-            label='N Ensemble Recursions', minimum=0, maximum=3, value=0, step=1
+        inp_n_parallel_trials = gr.Slider(
+            label='N Parallel Trials', minimum=1, maximum=3, value=1, step=1
         )
       btn_run = gr.Button('Run')
 
@@ -171,6 +172,7 @@ The Info textbox will show the final selected pixel coordinate that PIVOT conver
         columns=4,
         rows=1,
         interactive=False,
+        object_fit="contain", height="auto"
     )
     out_info = gr.Textbox(label='Info', lines=1)
 
@@ -182,7 +184,7 @@ The Info textbox will show the final selected pixel coordinate that PIVOT conver
           inp_n_samples_init,
           inp_n_samples_opt,
           inp_n_iters,
-          inp_n_recurssions,
+          inp_n_parallel_trials,
           inp_openai_api_key,
       ],
       outputs=[out_ims, out_info],

vip_runner.py

@@ -5,6 +5,7 @@ import re
 
 import cv2
 from tqdm import trange
+import numpy as np
 import vip
 
 
@@ -48,7 +49,11 @@ def vip_perform_selection(prompter, vlm, im, desc, arm_coord, samples, top_n):
   prompt_seq = [make_prompt(desc, top_n=top_n), encoded_image_circles]
   response = vlm.query(prompt_seq)
 
-  arrow_ids = extract_json(response, "points")
+  try:
+    arrow_ids = extract_json(response, "points")
+  except Exception as e:
+    print(e)
+    arrow_ids = []
   return arrow_ids, image_circles_np
 
 
@@ -61,7 +66,7 @@ def vip_runner(
     n_samples_init=25,
     n_samples_opt=10,
     n_iters=3,
-    recursion_level=0,
+    n_parallel_trials=1,
 ):
   """VIP."""
 
@@ -72,10 +77,11 @@ def vip_runner(
   output_ims = []
   arm_coord = (int(im.shape[1] / 2), int(im.shape[0] / 2))
 
-  if recursion_level == 0:
+  new_samples = []
+  center_mean = action_spec["loc"]
+  for i in range(n_parallel_trials):
     center_mean = action_spec["loc"]
     center_std = action_spec["scale"]
-    selected_samples = []
     for itr in trange(n_iters):
       if itr == 0:
         style["num_samples"] = n_samples_init
@@ -96,6 +102,7 @@ def vip_runner(
           image_circles_np, selected_samples, arm_coord
       )
       output_ims.append(image_circles_marked_np)
+      yield output_ims, f"Image generated for parallel sample {i+1}/{n_parallel_trials} iteration {itr+1}/{n_iters}. Still working..."
 
       # if at last iteration, pick one answer out of the selected ones
       if itr == n_iters - 1:
@@ -112,30 +119,11 @@ def vip_runner(
             im, selected_samples, arm_coord
         )
         output_ims.append(image_circles_marked_np)
+        new_samples += selected_samples
+        yield output_ims, f"Image generated for parallel sample {i+1}/{n_parallel_trials} last iteration. Still working..."
       center_mean, center_std = prompter.fit(arrow_ids, samples)
 
-    if output_ims:
-      return (
-          output_ims,
-          prompter.action_to_coord(center_mean, im, arm_coord).xy,
-          selected_samples,
-      )
-  else:
-    new_samples = []
-    for i in range(3):
-      out_ims, _, cur_samples = vip_runner(
-          vlm=vlm,
-          im=im,
-          desc=desc,
-          style=style,
-          action_spec=action_spec,
-          n_samples_init=n_samples_init,
-          n_samples_opt=n_samples_opt,
-          n_iters=n_iters,
-          recursion_level=recursion_level - 1,
-      )
-      output_ims += out_ims
-      new_samples += cur_samples
+  if n_parallel_trials > 1:
     # adjust sample label to avoid duplications
     for sample_id in range(len(new_samples)):
       new_samples[sample_id].label = str(sample_id)
@@ -154,10 +142,12 @@ def vip_runner(
     output_ims.append(image_circles_marked_np)
     center_mean, _ = prompter.fit(arrow_ids, new_samples)
 
-    if output_ims:
-      return (
-          output_ims,
-          prompter.action_to_coord(center_mean, im, arm_coord).xy,
-          selected_samples,
-      )
+  if output_ims:
+    yield (
+        output_ims,
+        (
+            "Final selected coordinate:"
+            f" {np.round(prompter.action_to_coord(center_mean, im, arm_coord).xy, decimals=0)}"
+        ),
+    )
   return [], "Unable to understand query"