Expose function

app.py

@@ -163,152 +163,146 @@ def draw_entity_boxes_on_image(image, entities, show=False, save_path=None, enti
     return pil_image
 
 
-def main():
 
-    ckpt = "microsoft/kosmos-2-patch14-224"
+ckpt = "microsoft/kosmos-2-patch14-224"
 
-    model = AutoModelForVision2Seq.from_pretrained(ckpt).to("cuda")
-    processor = AutoProcessor.from_pretrained(ckpt)
+model = AutoModelForVision2Seq.from_pretrained(ckpt).to("cuda")
+processor = AutoProcessor.from_pretrained(ckpt)
 
-    def generate_predictions(image_input, text_input):
+def generate_predictions(image_input, text_input):
 
-        # 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)
+    # 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"
-        elif text_input == "Detailed":
-            text_input = "<grounding>Describe this image in detail:"
-        else:
-            text_input = f"<grounding>{text_input}"
-
-        inputs = processor(text=text_input, images=image_input, return_tensors="pt").to("cuda")
-
-        generated_ids = model.generate(
-            pixel_values=inputs["pixel_values"],
-            input_ids=inputs["input_ids"],
-            attention_mask=inputs["attention_mask"],
-            image_embeds=None,
-            image_embeds_position_mask=inputs["image_embeds_position_mask"],
-            use_cache=True,
-            max_new_tokens=128,
-        )
-
-        generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-
-        # 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)
-
-        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))
-            colored_text.append((processed_text[start:end], f"{color_id}"))
-            prev_start = end
-
-        if end < len(processed_text):
-            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).
-    """
+    if text_input == "Brief":
+        text_input = "<grounding>An image of"
+    elif text_input == "Detailed":
+        text_input = "<grounding>Describe this image in detail:"
+    else:
+        text_input = f"<grounding>{text_input}"
+
+    inputs = processor(text=text_input, images=image_input, return_tensors="pt").to("cuda")
+
+    generated_ids = model.generate(
+        pixel_values=inputs["pixel_values"],
+        input_ids=inputs["input_ids"],
+        attention_mask=inputs["attention_mask"],
+        image_embeds=None,
+        image_embeds_position_mask=inputs["image_embeds_position_mask"],
+        use_cache=True,
+        max_new_tokens=128,
+    )
+
+    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+
+    # By default, the generated  text is cleanup and the entities are extracted.
+    processed_text, entities = processor.post_process_generation(generated_text)
 
-    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)
-            """))
-        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")
-
-                run_button = gr.Button(label="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,
-                                ).style(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])
-            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.Markdown(term_of_use)
-
-        # record which text span (label) is selected
-        selected = gr.Number(-1, show_label=False, placeholder="Selected", visible=False)
-
-        # record the current `entities`
-        entity_output = gr.Textbox(visible=False)
-
-        # get the current selected span label
-        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)
-
-        # 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)
-            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)
-
-    demo.launch(share=False)
-
-
-if __name__ == "__main__":
-    main()
-    # trigger
+    annotated_image = draw_entity_boxes_on_image(image_input, entities, show=False)
+
+    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))
+        colored_text.append((processed_text[start:end], f"{color_id}"))
+        prev_start = end
+
+    if end < len(processed_text):
+        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).
+"""
+
+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)
+        """))
+    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")
+
+            run_button = gr.Button(label="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,
+                            ).style(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])
+        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.Markdown(term_of_use)
+
+    # record which text span (label) is selected
+    selected = gr.Number(-1, show_label=False, placeholder="Selected", visible=False)
+
+    # record the current `entities`
+    entity_output = gr.Textbox(visible=False)
+
+    # get the current selected span label
+    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)
+
+    # 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)
+        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)
+
+demo.launch(share=False)