Update app.py

app.py

@@ -171,6 +171,30 @@ def upload_callback(image_input, state, visual_chatgpt=None):
     return state, state, image_input, click_state, image_input, image_input, image_embedding, \
         original_size, input_size
 
+def store_click(image_input, point_prompt, click_mode, state, click_state, evt: gr.SelectData):
+    click_index = evt.index
+    if point_prompt == 'Positive':
+        coordinate = [click_index[0], click_index[1], 1]
+    else:
+        coordinate = [click_index[0], click_index[1], 0]
+    
+    if click_mode == 'Continuous':
+        click_state[0].append(coordinate)
+    elif click_mode == 'Single':
+        click_state[0] = [coordinate]  # Overwrite with latest click
+
+    return state, click_state
+
+def generate_caption(image_input, enable_wiki, language, sentiment, factuality, length, image_embedding, state, click_state, original_size, input_size, text_refiner, visual_chatgpt):
+    last_click = click_state[0][-1]
+    point_prompt = 'Positive' if last_click[2] == 1 else 'Negative'
+    
+    return inference_click(image_input, point_prompt, click_mode, enable_wiki, language, sentiment, factuality, length, image_embedding, state, click_state, original_size, input_size, text_refiner, visual_chatgpt, gr.SelectData(index=(last_click[0], last_click[1])))
+    
+
+
+
+
 def inference_click(image_input, point_prompt, click_mode, enable_wiki, language, sentiment, factuality,
                     length, image_embedding, state, click_state, original_size, input_size, text_refiner, visual_chatgpt,
                     evt: gr.SelectData):
@@ -185,18 +209,6 @@ def inference_click(image_input, point_prompt, click_mode, enable_wiki, language
     input_points = prompt['input_point']
     input_labels = prompt['input_label']
 
-    click_state[0] = input_points
-    click_state[1] = input_labels
-    state = state + [("Image point: {}, Input label: {}".format(click_state[0], click_state[1]), None)]
-
-    return state, click_state
-
-def submit_caption(image_input, enable_wiki, language, sentiment, factuality, length, image_embedding, state, click_state, original_size, input_size, text_refiner, visual_chatgpt):
-    coordinates = json.dumps(list(zip(click_state[0], click_state[1])))
-    prompt = get_click_prompt(coordinates, click_state, 'Single')
-    input_points = prompt['input_point']
-    input_labels = prompt['input_label']
-
     controls = {'length': length,
                 'sentiment': sentiment,
                 'factuality': factuality,
@@ -217,39 +229,37 @@ def submit_caption(image_input, enable_wiki, language, sentiment, factuality, le
     enable_wiki = True if enable_wiki in ['True', 'TRUE', 'true', True, 'Yes', 'YES', 'yes'] else False
     out = model.inference(image_input, prompt, controls, disable_gpt=True, enable_wiki=enable_wiki, verbose=True, args={'clip_filter': False})[0]
 
+    state = state + [("Image point: {}, Input label: {}".format(prompt["input_point"], prompt["input_label"]), None)]
     state = state + [(None, "raw_caption: {}".format(out['generated_captions']['raw_caption']))]
-
-    update_click_state(click_state, out['generated_captions']['raw_caption'], 'Single')
+    update_click_state(click_state, out['generated_captions']['raw_caption'], click_mode)
     text = out['generated_captions']['raw_caption']
     input_mask = np.array(out['mask'].convert('P'))
     image_input = mask_painter(np.array(image_input), input_mask)
     origin_image_input = image_input
-    image_input = create_bubble_frame(image_input, text, (input_points[-1][0], input_points[-1][1]), input_mask,
+    image_input = create_bubble_frame(image_input, text, (click_index[0], click_index[1]), input_mask,
                                       input_points=input_points, input_labels=input_labels)
-
+    x, y = input_points[-1]
+    
     if visual_chatgpt is not None:
-        new_image_path = get_new_image_name('chat_image', func_name='upload')
-        image_input.save(new_image_path)
-        visual_chatgpt.current_image = new_image_path
-        img_caption = model.captioner.inference(image_input, filter=False, args={'text_prompt':''})['caption']
-        Human_prompt = f'\nHuman: The description of the image with path {new_image_path} is: {img_caption}. This information helps you to understand this image, but you should use tools to finish following tasks, rather than directly imagine from my description. If you understand, say \"Received\". \n'
-        AI_prompt = "Received."
-        visual_chatgpt.global_prompt = Human_prompt + 'AI: ' + AI_prompt
-        visual_chatgpt.agent.memory.buffer = visual_chatgpt.agent.memory.buffer + visual_chatgpt.global_prompt
+        print('inference_click: add caption to chatGPT memory')
+        new_crop_save_path = get_new_image_name('chat_image', func_name='crop')
+        Image.open(out["crop_save_path"]).save(new_crop_save_path)
+        point_prompt = f'You should primarly use tools on the selected regional image (description: {text}, path: {new_crop_save_path}), which is a part of the whole image (path: {visual_chatgpt.current_image}). If human mentioned some objects not in the selected region, you can use tools on the whole image.'
+        visual_chatgpt.point_prompt = point_prompt
 
+    yield state, state, click_state, image_input
     if not args.disable_gpt and model.text_refiner:
         refined_caption = model.text_refiner.inference(query=text, controls=controls, context=out['context_captions'],
                                                        enable_wiki=enable_wiki)
+        # new_cap = 'Original: ' + text + '. Refined: ' + refined_caption['caption']
         new_cap = refined_caption['caption']
         if refined_caption['wiki']:
             state = state + [(None, "Wiki: {}".format(refined_caption['wiki']))]
         state = state + [(None, f"caption: {new_cap}")]
-        refined_image_input = create_bubble_frame(origin_image_input, new_cap, (input_points[-1][0], input_points[-1][1]),
+        refined_image_input = create_bubble_frame(origin_image_input, new_cap, (click_index[0], click_index[1]),
                                                   input_mask,
                                                   input_points=input_points, input_labels=input_labels)
-        return state, state, click_state, refined_image_input
-
-    return state, state, click_state, image_input
+        yield state, state, click_state, refined_image_input
 
 
 def get_sketch_prompt(mask: Image.Image):
@@ -556,10 +566,8 @@ def create_ui():
         )
         clear_button_text.click(clear_chat_memory, inputs=[visual_chatgpt])
         
-        submit_button_click.click(submit_caption,
-                                  [origin_image, enable_wiki, language, sentiment, factuality, length, image_embedding, state, click_state, original_size, input_size, text_refiner, visual_chatgpt],
-                                  [chatbot, state, click_state, image_input])
-        
+        submit_button_click.click(generate_caption, inputs=[origin_image, enable_wiki, language, sentiment, factuality, length, image_embedding, state, click_state, original_size, input_size, text_refiner, visual_chatgpt], outputs=[chatbot, state, click_state, image_input])
+
         image_input.clear(
             lambda: (None, [], [], [[], [], []], "", "", ""),
             [],
@@ -588,15 +596,7 @@ def create_ui():
                               image_embedding, original_size, input_size])
         example_image.change(clear_chat_memory, inputs=[visual_chatgpt])
         # select coordinate
-        image_input.select(
-            inference_click,
-            inputs=[
-                origin_image, point_prompt, click_mode, enable_wiki, language, sentiment, factuality, length,
-                image_embedding, state, click_state, original_size, input_size, text_refiner, visual_chatgpt
-            ],
-            outputs=[chatbot, state, click_state, image_input],
-            show_progress=False, queue=True
-        )
+        image_input.select(store_click, inputs=[origin_image, point_prompt, click_mode, state, click_state], outputs=[state, click_state])
 
         submit_button_sketcher.click(
             inference_traject,