Update app.py

app.py

@@ -10,6 +10,7 @@ from transformers import AutoProcessor, AutoModelForVision2Seq
 import cv2
 import ast
 
+
 colors = [
     (0, 255, 0),
     (0, 0, 255),
@@ -30,18 +31,21 @@ colors = [
 ]
 
 color_map = {
-    f"{color_id}": f"#{hex(color[2])[2:].zfill(2)}{hex(color[1])[2:].zfill(2)}{hex(color[0])[2:].zfill(2)}" for color_id, color in enumerate(colors)
+    f"{color_id}": f"#{hex(color[2])[2:].zfill(2)}{hex(color[1])[2:].zfill(2)}{hex(color[0])[2:].zfill(2)}"
+    for color_id, color in enumerate(colors)
 }
 
 
 def is_overlapping(rect1, rect2):
     x1, y1, x2, y2 = rect1
     x3, y3, x4, y4 = rect2
+
     return not (x2 < x3 or x1 > x4 or y2 < y3 or y1 > y4)
 
 
 def draw_entity_boxes_on_image(image, entities, show=False, save_path=None, entity_index=-1):
     """_summary_
+
     Args:
         image (_type_): image or image path
         collect_entity_location (_type_): _description_
@@ -59,10 +63,13 @@ def draw_entity_boxes_on_image(image, entities, show=False, save_path=None, enti
         else:
             raise ValueError(f"invaild image path, {image}")
     elif isinstance(image, torch.Tensor):
-        # pdb.set_trace()
         image_tensor = image.cpu()
-        reverse_norm_mean = torch.tensor([0.48145466, 0.4578275, 0.40821073])[:, None, None]
-        reverse_norm_std = torch.tensor([0.26862954, 0.26130258, 0.27577711])[:, None, None]
+        reverse_norm_mean = torch.tensor([0.48145466, 0.4578275, 0.40821073])[
+            :, None, None
+        ]
+        reverse_norm_std = torch.tensor([0.26862954, 0.26130258, 0.27577711])[
+            :, None, None
+        ]
         image_tensor = image_tensor * reverse_norm_std + reverse_norm_mean
         pil_img = T.ToPILImage()(image_tensor)
         image_h = pil_img.height
@@ -79,59 +86,85 @@ def draw_entity_boxes_on_image(image, entities, show=False, save_path=None, enti
         indices = [entity_index]
 
     # Not to show too many bboxes
-    entities = entities[:len(color_map)]
-
+    entities = entities[: len(color_map)]
     new_image = image.copy()
     previous_bboxes = []
-    # size of text
+
     text_size = 1
-    # thickness of text
-    text_line = 1  # int(max(1 * min(image_h, image_w) / 512, 1))
+    text_line = 1
     box_line = 3
-    (c_width, text_height), _ = cv2.getTextSize("F", cv2.FONT_HERSHEY_COMPLEX, text_size, text_line)
+    (c_width, text_height), _ = cv2.getTextSize(
+        "F", cv2.FONT_HERSHEY_COMPLEX, text_size, text_line
+    )
     base_height = int(text_height * 0.675)
     text_offset_original = text_height - base_height
     text_spaces = 3
-
-    # num_bboxes = sum(len(x[-1]) for x in entities)
-    used_colors = colors  # random.sample(colors, k=num_bboxes)
-
+    used_colors = colors
     color_id = -1
+
     for entity_idx, (entity_name, (start, end), bboxes) in enumerate(entities):
         color_id += 1
         if entity_idx not in indices:
             continue
-        for bbox_id, (x1_norm, y1_norm, x2_norm, y2_norm) in enumerate(bboxes):
-            # if start is None and bbox_id > 0:
-            #     color_id += 1
-            orig_x1, orig_y1, orig_x2, orig_y2 = int(x1_norm * image_w), int(y1_norm * image_h), int(x2_norm * image_w), int(y2_norm * image_h)
 
-            # draw bbox
-            # random color
-            color = used_colors[color_id]  # tuple(np.random.randint(0, 255, size=3).tolist())
-            new_image = cv2.rectangle(new_image, (orig_x1, orig_y1), (orig_x2, orig_y2), color, box_line)
+        for bbox_id, (x1_norm, y1_norm, x2_norm, y2_norm) in enumerate(bboxes):
+            orig_x1, orig_y1, orig_x2, orig_y2 = (
+                int(x1_norm * image_w),
+                int(y1_norm * image_h),
+                int(x2_norm * image_w),
+                int(y2_norm * image_h),
+            )
 
-            l_o, r_o = box_line // 2 + box_line % 2, box_line // 2 + box_line % 2 + 1
+            color = used_colors[color_id]
+            new_image = cv2.rectangle(
+                new_image,
+                (orig_x1, orig_y1),
+                (orig_x2, orig_y2),
+                color,
+                box_line,
+            )
 
+            l_o, r_o = (
+                box_line // 2 + box_line % 2,
+                box_line // 2 + box_line % 2 + 1,
+            )
             x1 = orig_x1 - l_o
             y1 = orig_y1 - l_o
 
             if y1 < text_height + text_offset_original + 2 * text_spaces:
-                y1 = orig_y1 + r_o + text_height + text_offset_original + 2 * text_spaces
+                y1 = (
+                    orig_y1
+                    + r_o
+                    + text_height
+                    + text_offset_original
+                    + 2 * text_spaces
+                )
                 x1 = orig_x1 + r_o
 
-            # add text background
-            (text_width, text_height), _ = cv2.getTextSize(f"  {entity_name}", cv2.FONT_HERSHEY_COMPLEX, text_size, text_line)
-            text_bg_x1, text_bg_y1, text_bg_x2, text_bg_y2 = x1, y1 - (text_height + text_offset_original + 2 * text_spaces), x1 + text_width, y1
+            (text_width, text_height), _ = cv2.getTextSize(
+                f" {entity_name}", cv2.FONT_HERSHEY_COMPLEX, text_size, text_line
+            )
+            text_bg_x1 = x1
+            text_bg_y1 = y1 - (
+                text_height + text_offset_original + 2 * text_spaces
+            )
+            text_bg_x2 = x1 + text_width
+            text_bg_y2 = y1
 
             for prev_bbox in previous_bboxes:
-                while is_overlapping((text_bg_x1, text_bg_y1, text_bg_x2, text_bg_y2), prev_bbox):
-                    text_bg_y1 += (text_height + text_offset_original + 2 * text_spaces)
-                    text_bg_y2 += (text_height + text_offset_original + 2 * text_spaces)
-                    y1 += (text_height + text_offset_original + 2 * text_spaces)
+                while is_overlapping(
+                    (text_bg_x1, text_bg_y1, text_bg_x2, text_bg_y2), prev_bbox
+                ):
+                    text_bg_y1 += text_height + text_offset_original + 2 * text_spaces
+                    text_bg_y2 += text_height + text_offset_original + 2 * text_spaces
+                    y1 += text_height + text_offset_original + 2 * text_spaces
 
                     if text_bg_y2 >= image_h:
-                        text_bg_y1 = max(0, image_h - (text_height + text_offset_original + 2 * text_spaces))
+                        text_bg_y1 = max(
+                            0,
+                            image_h
+                            - (text_height + text_offset_original + 2 * text_spaces),
+                        )
                         text_bg_y2 = image_h
                         y1 = image_h
                         break
@@ -141,22 +174,34 @@ def draw_entity_boxes_on_image(image, entities, show=False, save_path=None, enti
                 for j in range(text_bg_x1, text_bg_x2):
                     if i < image_h and j < image_w:
                         if j < text_bg_x1 + 1.35 * c_width:
-                            # original color
                             bg_color = color
                         else:
-                            # white
                             bg_color = [255, 255, 255]
-                        new_image[i, j] = (alpha * new_image[i, j] + (1 - alpha) * np.array(bg_color)).astype(np.uint8)
+                        new_image[i, j] = (
+                            alpha * new_image[i, j]
+                            + (1 - alpha) * np.array(bg_color)
+                        ).astype(np.uint8)
 
             cv2.putText(
-                new_image, f"  {entity_name}", (x1, y1 - text_offset_original - 1 * text_spaces), cv2.FONT_HERSHEY_COMPLEX, text_size, (0, 0, 0), text_line, cv2.LINE_AA
+                new_image,
+                f" {entity_name}",
+                (x1, y1 - text_offset_original - 1 * text_spaces),
+                cv2.FONT_HERSHEY_COMPLEX,
+                text_size,
+                (0, 0, 0),
+                text_line,
+                cv2.LINE_AA,
+            )
+
+            previous_bboxes.append(
+                (text_bg_x1, text_bg_y1, text_bg_x2, text_bg_y2)
             )
-            # previous_locations.append((x1, y1))
-            previous_bboxes.append((text_bg_x1, text_bg_y1, text_bg_x2, text_bg_y2))
 
     pil_image = Image.fromarray(new_image[:, :, [2, 1, 0]])
+
     if save_path:
         pil_image.save(save_path)
+
     if show:
         pil_image.show()
 
@@ -164,29 +209,41 @@ def draw_entity_boxes_on_image(image, entities, show=False, save_path=None, enti
 
 
 def main():
-
     ckpt = "microsoft/kosmos-2-patch14-224"
-
     model = AutoModelForVision2Seq.from_pretrained(ckpt).to("cuda")
     processor = AutoProcessor.from_pretrained(ckpt)
 
     def generate_predictions(image_input, text_input):
+        """
+        Generate a grounded image description and annotated entity boxes with Kosmos-2.
 
+        Use this tool when you need to describe an image and identify grounded visual entities.
+
+        Args:
+            image_input (PIL.Image.Image): Input image to describe and ground.
+            text_input (str): Description mode, either "Brief" or "Detailed".
+
+        Returns:
+            tuple: Annotated image, highlighted generated description, and serialized entity data.
+        """
         # Save the image and load it again to match the original Kosmos-2 demo.
         # (https://github.com/microsoft/unilm/blob/f4695ed0244a275201fff00bee495f76670fbe70/kosmos-2/demo/gradio_app.py#L345-L346)
         user_image_path = "/tmp/user_input_test_image.jpg"
         image_input.save(user_image_path)
+
         # This might give different results from the original argument `image_input`
         image_input = Image.open(user_image_path)
 
         if text_input == "Brief":
-            text_input = "<grounding>An image of"
+            text_input = "An image of"
         elif text_input == "Detailed":
-            text_input = "<grounding>Describe this image in detail:"
+            text_input = "Describe this image in detail:"
         else:
-            text_input = f"<grounding>{text_input}"
+            text_input = f"{text_input}"
 
-        inputs = processor(text=text_input, images=image_input, return_tensors="pt").to("cuda")
+        inputs = processor(text=text_input, images=image_input, return_tensors="pt").to(
+            "cuda"
+        )
 
         generated_ids = model.generate(
             pixel_values=inputs["pixel_values"],
@@ -198,9 +255,11 @@ def main():
             max_new_tokens=128,
         )
 
-        generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+        generated_text = processor.batch_decode(
+            generated_ids, skip_special_tokens=True
+        )[0]
 
-        # By default, the generated  text is cleanup and the entities are extracted.
+        # By default, the generated text is cleanup and the entities are extracted.
         processed_text, entities = processor.post_process_generation(generated_text)
 
         annotated_image = draw_entity_boxes_on_image(image_input, entities, show=False)
@@ -208,21 +267,22 @@ def main():
         color_id = -1
         entity_info = []
         filtered_entities = []
+
         for entity in entities:
             entity_name, (start, end), bboxes = entity
+
             if start == end:
                 # skip bounding bbox without a `phrase` associated
                 continue
+
             color_id += 1
-            # for bbox_id, _ in enumerate(bboxes):
-                # if start is None and bbox_id > 0:
-                #     color_id += 1
             entity_info.append(((start, end), color_id))
             filtered_entities.append(entity)
 
         colored_text = []
         prev_start = 0
         end = 0
+
         for idx, ((start, end), color_id) in enumerate(entity_info):
             if start > prev_start:
                 colored_text.append((processed_text[prev_start:start], None))
@@ -230,55 +290,73 @@ def main():
             prev_start = end
 
         if end < len(processed_text):
-            colored_text.append((processed_text[end:len(processed_text)], None))
+            colored_text.append((processed_text[end : len(processed_text)], None))
 
         return annotated_image, colored_text, str(filtered_entities)
 
     term_of_use = """
-    ### Terms of use  
-    By using this model, users are required to agree to the following terms:  
-    The model is intended for academic and research purposes. 
-    The utilization of the model to create unsuitable material is strictly forbidden and not endorsed by this work. 
-    The accountability for any improper or unacceptable application of the model rests exclusively with the individuals who generated such content. 
-    
-    ### License
-    This project is licensed under the [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct).
-    """
+### Terms of use
+
+By using this model, users are required to agree to the following terms:
+
+The model is intended for academic and research purposes.
+The utilization of the model to create unsuitable material is strictly forbidden and not endorsed by this work.
+The accountability for any improper or unacceptable application of the model rests exclusively with the individuals who generated such content.
+
+### License
+
+This project is licensed under the [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct).
+"""
 
     with gr.Blocks(title="Kosmos-2", theme=gr.themes.Base()).queue() as demo:
-        gr.Markdown(("""
-            # Kosmos-2: Grounding Multimodal Large Language Models to the World
-            [[Paper]](https://arxiv.org/abs/2306.14824) [[Code]](https://github.com/microsoft/unilm/blob/master/kosmos-2)
-            """))
+        gr.Markdown(
+            (
+                """
+# Kosmos-2: Grounding Multimodal Large Language Models to the World
+
+[[Paper]](https://arxiv.org/abs/2306.14824) [[Code]](https://github.com/microsoft/unilm/blob/master/kosmos-2)
+"""
+            )
+        )
+
         with gr.Row():
             with gr.Column():
                 image_input = gr.Image(type="pil", label="Test Image")
-                text_input = gr.Radio(["Brief", "Detailed"], label="Description Type", value="Brief")
-
+                text_input = gr.Radio(
+                    ["Brief", "Detailed"], label="Description Type", value="Brief"
+                )
                 run_button = gr.Button("Run", visible=True)
 
             with gr.Column():
                 image_output = gr.Image(type="pil")
                 text_output1 = gr.HighlightedText(
-                                    label="Generated Description",
-                                    combine_adjacent=False,
-                                    show_legend=True,
-                                    color_map=color_map
-                                )
+                    label="Generated Description",
+                    combine_adjacent=False,
+                    show_legend=True,
+                    color_map=color_map,
+                )
 
         with gr.Row():
             with gr.Column():
-                gr.Examples(examples=[
-                            ["images/two_dogs.jpg", "Detailed"],
-                            ["images/snowman.png", "Brief"],
-                            ["images/man_ball.png", "Detailed"],
-                        ], inputs=[image_input, text_input])
+                gr.Examples(
+                    examples=[
+                        ["images/two_dogs.jpg", "Detailed"],
+                        ["images/snowman.png", "Brief"],
+                        ["images/man_ball.png", "Detailed"],
+                    ],
+                    inputs=[image_input, text_input],
+                )
+
             with gr.Column():
-                gr.Examples(examples=[
-                            ["images/six_planes.png", "Brief"],
-                            ["images/quadrocopter.jpg", "Brief"],
-                            ["images/carnaby_street.jpg", "Brief"],
-                        ], inputs=[image_input, text_input])
+                gr.Examples(
+                    examples=[
+                        ["images/six_planes.png", "Brief"],
+                        ["images/quadrocopter.jpg", "Brief"],
+                        ["images/carnaby_street.jpg", "Brief"],
+                    ],
+                    inputs=[image_input, text_input],
+                )
+
         gr.Markdown(term_of_use)
 
         # record which text span (label) is selected
@@ -291,25 +369,43 @@ def main():
         def get_text_span_label(evt: gr.SelectData):
             if evt.value[-1] is None:
                 return -1
+
             return int(evt.value[-1])
+
         # and set this information to `selected`
-        text_output1.select(get_text_span_label, None, selected)
+        text_output1.select(
+            get_text_span_label,
+            None,
+            selected,
+            api_visibility="private",
+        )
 
         # update output image when we change the span (enity) selection
         def update_output_image(img_input, image_output, entities, idx):
             entities = ast.literal_eval(entities)
-            updated_image = draw_entity_boxes_on_image(img_input, entities, entity_index=idx)
+            updated_image = draw_entity_boxes_on_image(
+                img_input, entities, entity_index=idx
+            )
             return updated_image
-        selected.change(update_output_image, [image_input, image_output, entity_output, selected], [image_output])
 
-        run_button.click(fn=generate_predictions,
-                         inputs=[image_input, text_input],
-                         outputs=[image_output, text_output1, entity_output],
-                         show_progress=True, queue=True)
+        selected.change(
+            update_output_image,
+            [image_input, image_output, entity_output, selected],
+            [image_output],
+            api_visibility="private",
+        )
+
+        run_button.click(
+            fn=generate_predictions,
+            inputs=[image_input, text_input],
+            outputs=[image_output, text_output1, entity_output],
+            show_progress=True,
+            queue=True,
+        )
 
-    demo.launch(share=False, ssr_mode=False)
+    demo.launch(share=False, ssr_mode=False, mcp_server=True)
 
 
 if __name__ == "__main__":
     main()
-    # trigger
+# trigger