Update app.py

app.py

@@ -9,11 +9,15 @@ from utils import check_ocr_box, get_yolo_model, get_caption_model_processor, ge
 from PIL import Image
 
 from ultralytics import YOLO
-yolo_model = YOLO('weights/icon_detect/best.pt')  # Removed .to('cuda')
+yolo_model = YOLO('weights/icon_detect/best.pt')
 
 from transformers import AutoProcessor, AutoModelForCausalLM 
 processor = AutoProcessor.from_pretrained("microsoft/Florence-2-base", trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained("weights/icon_caption_florence", torch_dtype=torch.float32, trust_remote_code=True)  # Changed dtype to float32 and removed .to('cuda')
+model = AutoModelForCausalLM.from_pretrained(
+    "weights/icon_caption_florence",
+    torch_dtype=torch.float32,
+    trust_remote_code=True
+)
 
 caption_model_processor = {'processor': processor, 'model': model}
 print('Finished loading model.')
@@ -57,7 +61,7 @@ def process(
     )
     text, ocr_bbox = ocr_bbox_rslt
 
-    dino_labled_img, label_coordinates, parsed_content_list = get_som_labeled_img(
+    dino_labeled_img, label_coordinates, parsed_content_list = get_som_labeled_img(
         image_save_path,
         yolo_model,
         BOX_TRESHOLD=box_threshold,
@@ -68,10 +72,12 @@ def process(
         ocr_text=text,
         iou_threshold=iou_threshold
     )
-    image = Image.open(io.BytesIO(base64.b64decode(dino_labled_img)))
+    image = Image.open(io.BytesIO(base64.b64decode(dino_labeled_img)))
     print('Finished processing.')
-    parsed_content_list = '\n'.join(parsed_content_list)
-    return image, str(parsed_content_list)
+    parsed_content_list_str = '\n'.join(parsed_content_list)
+    label_coordinates_str = '\n'.join([str(coord) for coord in label_coordinates])
+
+    return image, parsed_content_list_str, label_coordinates_str
 
 with gr.Blocks() as demo:
     gr.Markdown(MARKDOWN)
@@ -88,6 +94,8 @@ with gr.Blocks() as demo:
             image_output_component = gr.Image(type='pil', label='Image Output')
             text_output_component = gr.Textbox(
                 label='Parsed Screen Elements', placeholder='Text Output')
+            coordinates_output_component = gr.Textbox(
+                label='Coordinates', placeholder='Coordinates Output')
 
     submit_button_component.click(
         fn=process,
@@ -96,7 +104,11 @@ with gr.Blocks() as demo:
             box_threshold_component,
             iou_threshold_component
         ],
-        outputs=[image_output_component, text_output_component]
+        outputs=[
+            image_output_component,
+            text_output_component,
+            coordinates_output_component
+        ]
     )
 
-demo.queue().launch(share=False)
+demo.queue().launch(share=False)