From f907013aa99e91d24ed4004bd121196f5736c259 Mon Sep 17 00:00:00 2001 From: Billy Cao Date: Mon, 4 Nov 2024 09:41:06 +0800 Subject: [PATCH] Add icon detect image size option --- demo.ipynb | 2 +- gradio_demo.py | 10 +++++++--- utils.py | 7 ++++--- 3 files changed, 12 insertions(+), 7 deletions(-) diff --git a/demo.ipynb b/demo.ipynb index ec502bd..73ba6cc 100644 --- a/demo.ipynb +++ b/demo.ipynb @@ -104,7 +104,7 @@ "ocr_bbox_rslt, is_goal_filtered = check_ocr_box(image_path, display_img = False, output_bb_format='xyxy', goal_filtering=None, easyocr_args={'paragraph': False, 'text_threshold':0.9}, use_paddleocr=True)\n", "text, ocr_bbox = ocr_bbox_rslt\n", "\n", - "dino_labled_img, label_coordinates, parsed_content_list = get_som_labeled_img(image_path, som_model, BOX_TRESHOLD = BOX_TRESHOLD, output_coord_in_ratio=False, ocr_bbox=ocr_bbox,draw_bbox_config=draw_bbox_config, caption_model_processor=caption_model_processor, ocr_text=text,use_local_semantics=True, iou_threshold=0.1)\n", + "dino_labled_img, label_coordinates, parsed_content_list = get_som_labeled_img(image_path, som_model, BOX_TRESHOLD = BOX_TRESHOLD, output_coord_in_ratio=False, ocr_bbox=ocr_bbox,draw_bbox_config=draw_bbox_config, caption_model_processor=caption_model_processor, ocr_text=text,use_local_semantics=True, iou_threshold=0.1, imgsz=640)\n", "\n", "\n" ] diff --git a/gradio_demo.py b/gradio_demo.py index f420b7a..825f7fd 100644 --- a/gradio_demo.py +++ b/gradio_demo.py @@ -61,7 +61,8 @@ def process( image_input, box_threshold, iou_threshold, - use_paddleocr + use_paddleocr, + imgsz ) -> Optional[Image.Image]: image_save_path = 'imgs/saved_image_demo.png' @@ -71,7 +72,7 @@ def process( ocr_bbox_rslt, is_goal_filtered = check_ocr_box(image_save_path, 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 # print('prompt:', prompt) - dino_labled_img, label_coordinates, parsed_content_list = get_som_labeled_img(image_save_path, 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) + dino_labled_img, label_coordinates, parsed_content_list = get_som_labeled_img(image_save_path, 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(parsed_content_list) @@ -93,6 +94,8 @@ with gr.Blocks() as demo: label='IOU Threshold', minimum=0.01, maximum=1.0, step=0.01, value=0.1) use_paddleocr_component = gr.Checkbox( label='Use PaddleOCR', value=True) + imgsz_component = gr.Slider( + label='Icon Detect Image Size', minimum=640, maximum=1920, step=32, value=640) submit_button_component = gr.Button( value='Submit', variant='primary') with gr.Column(): @@ -105,7 +108,8 @@ with gr.Blocks() as demo: image_input_component, box_threshold_component, iou_threshold_component, - use_paddleocr_component + use_paddleocr_component, + imgsz_component ], outputs=[image_output_component, text_output_component] ) diff --git a/utils.py b/utils.py index 8fc9728..25dab4b 100755 --- a/utils.py +++ b/utils.py @@ -280,7 +280,7 @@ def predict(model, image, caption, box_threshold, text_threshold): return boxes, logits, phrases -def predict_yolo(model, image_path, box_threshold): +def predict_yolo(model, image_path, box_threshold, imgsz): """ Use huggingface model to replace the original model """ # model = model['model'] @@ -288,6 +288,7 @@ def predict_yolo(model, image_path, box_threshold): result = model.predict( source=image_path, conf=box_threshold, + imgsz=imgsz # iou=0.5, # default 0.7 ) boxes = result[0].boxes.xyxy#.tolist() # in pixel space @@ -297,7 +298,7 @@ def predict_yolo(model, image_path, box_threshold): return boxes, conf, phrases -def get_som_labeled_img(img_path, model=None, BOX_TRESHOLD = 0.01, output_coord_in_ratio=False, ocr_bbox=None, text_scale=0.4, text_padding=5, draw_bbox_config=None, caption_model_processor=None, ocr_text=[], use_local_semantics=True, iou_threshold=0.9,prompt=None): +def get_som_labeled_img(img_path, model=None, BOX_TRESHOLD = 0.01, output_coord_in_ratio=False, ocr_bbox=None, text_scale=0.4, text_padding=5, draw_bbox_config=None, caption_model_processor=None, ocr_text=[], use_local_semantics=True, iou_threshold=0.9,prompt=None,imgsz=640): """ ocr_bbox: list of xyxy format bbox """ TEXT_PROMPT = "clickable buttons on the screen" @@ -309,7 +310,7 @@ def get_som_labeled_img(img_path, model=None, BOX_TRESHOLD = 0.01, output_coord_ if False: # TODO xyxy, logits, phrases = predict(model=model, image=image_source, caption=TEXT_PROMPT, box_threshold=BOX_TRESHOLD, text_threshold=TEXT_TRESHOLD) else: - xyxy, logits, phrases = predict_yolo(model=model, image_path=img_path, box_threshold=BOX_TRESHOLD) + xyxy, logits, phrases = predict_yolo(model=model, image_path=img_path, box_threshold=BOX_TRESHOLD, imgsz=imgsz) xyxy = xyxy / torch.Tensor([w, h, w, h]).to(xyxy.device) image_source = np.asarray(image_source) phrases = [str(i) for i in range(len(phrases))]