Update app.py

app.py

@@ -1,11 +1,11 @@
-from typing import Optional, Tuple
+from typing import Optional
 import spaces
 import gradio as gr
 import numpy as np
 import torch
 from PIL import Image
 import io
-import json
+
 
 import base64, os
 from util.utils import check_ocr_box, get_yolo_model, get_caption_model_processor, get_som_labeled_img
@@ -23,6 +23,7 @@ snapshot_download(repo_id=repo_id, local_dir=local_dir)
 
 print(f"Repository downloaded to: {local_dir}")
 
+
 yolo_model = get_yolo_model(model_path='weights/icon_detect/model.pt')
 caption_model_processor = get_caption_model_processor(model_name="florence2", model_name_or_path="weights/icon_caption")
 # caption_model_processor = get_caption_model_processor(model_name="blip2", model_name_or_path="weights/icon_caption_blip2")
@@ -42,14 +43,18 @@ DEVICE = torch.device('cuda')
 
 @spaces.GPU
 @torch.inference_mode()
+# @torch.autocast(device_type="cuda", dtype=torch.bfloat16)
 def process(
     image_input,
     box_threshold,
     iou_threshold,
     use_paddleocr,
     imgsz
-) -> Tuple[str, str]:
+) -> Optional[Image.Image]:
 
+    # image_save_path = 'imgs/saved_image_demo.png'
+    # image_input.save(image_save_path)
+    # image = Image.open(image_save_path)
     box_overlay_ratio = image_input.size[0] / 3200
     draw_bbox_config = {
         'text_scale': 0.8 * box_overlay_ratio,
@@ -57,93 +62,38 @@ def process(
         'text_padding': max(int(3 * box_overlay_ratio), 1),
         'thickness': max(int(3 * box_overlay_ratio), 1),
     }
+    # import pdb; pdb.set_trace()
 
-    ocr_bbox_rslt, is_goal_filtered = check_ocr_box(
-        image_input, 
-        display_img=False, 
-        output_bb_format='xyxy', 
-        goal_filtering=None, 
-        easyocr_args={'paragraph': False, 'text_threshold': 0.9}, 
-        use_paddleocr=use_paddleocr
-    )
+    ocr_bbox_rslt, is_goal_filtered = check_ocr_box(image_input, display_img = False, output_bb_format='xyxy', goal_filtering=None, easyocr_args={'paragraph': False, 'text_threshold':0.9}, use_paddleocr=use_paddleocr)
     text, ocr_bbox = ocr_bbox_rslt
-    
-    dino_labled_img, label_coordinates, parsed_content_list = get_som_labeled_img(
-        image_input, 
-        yolo_model, 
-        BOX_TRESHOLD=box_threshold, 
-        output_coord_in_ratio=True, 
-        ocr_bbox=ocr_bbox,
-        draw_bbox_config=draw_bbox_config, 
-        caption_model_processor=caption_model_processor, 
-        ocr_text=text,
-        iou_threshold=iou_threshold, 
-        imgsz=imgsz
-    )
-    
+    dino_labled_img, label_coordinates, parsed_content_list = get_som_labeled_img(image_input, yolo_model, BOX_TRESHOLD = box_threshold, output_coord_in_ratio=True, ocr_bbox=ocr_bbox,draw_bbox_config=draw_bbox_config, caption_model_processor=caption_model_processor, ocr_text=text,iou_threshold=iou_threshold, imgsz=imgsz,)  
+    image = Image.open(io.BytesIO(base64.b64decode(dino_labled_img)))
     print('finish processing')
-    parsed_content_list = '\n'.join([f'icon {i}: ' + str(v) for i, v in enumerate(parsed_content_list)])
-    
-    # Convert label_coordinates to JSON string for API consumption
-    label_coordinates_json = json.dumps(label_coordinates)
-    
-    return str(parsed_content_list), label_coordinates_json
-
-# Create interface with simplified component definitions
+    parsed_content_list = '\n'.join([f'icon {i}: ' + str(v) for i,v in enumerate(parsed_content_list)])
+    # parsed_content_list = str(parsed_content_list)
+    return image, str(parsed_content_list)
+
 with gr.Blocks() as demo:
     gr.Markdown(MARKDOWN)
-    
     with gr.Row():
         with gr.Column():
             image_input_component = gr.Image(
-                type='pil', 
-                label='Upload image'
-            )
-            
-            # Simplified slider definitions
+                type='pil', label='Upload image')
+            # set the threshold for removing the bounding boxes with low confidence, default is 0.05
             box_threshold_component = gr.Slider(
-                minimum=0.01, 
-                maximum=1.0, 
-                value=0.05, 
-                step=0.01,
-                label='Box Threshold'
-            )
-            
+                label='Box Threshold', minimum=0.01, maximum=1.0, step=0.01, value=0.05)
+            # set the threshold for removing the bounding boxes with large overlap, default is 0.1
             iou_threshold_component = gr.Slider(
-                minimum=0.01, 
-                maximum=1.0, 
-                value=0.1, 
-                step=0.01,
-                label='IOU Threshold'
-            )
-            
+                label='IOU Threshold', minimum=0.01, maximum=1.0, step=0.01, value=0.1)
             use_paddleocr_component = gr.Checkbox(
-                value=True,
-                label='Use PaddleOCR'
-            )
-            
+                label='Use PaddleOCR', value=True)
             imgsz_component = gr.Slider(
-                minimum=640, 
-                maximum=1920, 
-                value=640, 
-                step=32,
-                label='Icon Detect Image Size'
-            )
-            
+                label='Icon Detect Image Size', minimum=640, maximum=1920, step=32, value=640)
             submit_button_component = gr.Button(
-                value='Submit', 
-                variant='primary'
-            )
-            
+                value='Submit', variant='primary')
         with gr.Column():
-            text_output_component = gr.Textbox(
-                label='Parsed screen elements', 
-                placeholder='Text Output'
-            )
-            coordinates_output_component = gr.Textbox(
-                label='Label Coordinates (JSON)', 
-                placeholder='Coordinates JSON Output'
-            )
+            image_output_component = gr.Image(type='pil', label='Image Output')
+            text_output_component = gr.Textbox(label='Parsed screen elements', placeholder='Text Output')
 
     submit_button_component.click(
         fn=process,
@@ -154,18 +104,10 @@ with gr.Blocks() as demo:
             use_paddleocr_component,
             imgsz_component
         ],
-        outputs=[text_output_component, coordinates_output_component]
+        outputs=[image_output_component, text_output_component]
     )
 
-# Try launching with different configurations
-try:
-    demo.queue().launch(share=True)
-except Exception as e:
-    print(f"Error launching with queue: {e}")
-    # Fallback: try without queue
-    try:
-        demo.launch(share=True)
-    except Exception as e2:
-        print(f"Error launching without queue: {e2}")
-        # Final fallback: basic launch
-        demo.launch(debug=True, show_error=True, share=True)
+# demo.launch(debug=False, show_error=True, share=True)
+# demo.launch(share=True, server_port=7861, server_name='0.0.0.0')
+demo.queue().launch(share=True)
+