Function process

agents/ui_agent/app.py:148–248

(
    image_input,
    box_threshold,
    iou_threshold,
    use_paddleocr,
    imgsz, 
    instruction,
)

@torch.inference_mode()
# @torch.autocast(device_type="cuda", dtype=torch.bfloat16)
def process(
    image_input,
    box_threshold,
    iou_threshold,
    use_paddleocr,
    imgsz, 
    instruction,
) -> 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,
        'text_thickness': max(int(2 * box_overlay_ratio), 1),
        'text_padding': max(int(3 * box_overlay_ratio), 1),
        'thickness': max(int(3 * box_overlay_ratio), 1),
    }

    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=False, 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,)  
    parsed_content_list = '\n'.join([f'icon {i}: ' + str(v) for i,v in enumerate(parsed_content_list)])
    
    if len(instruction) == 0:
        print('finish processing')
        image = Image.open(io.BytesIO(base64.b64decode(dino_labled_img)))
        return image, str(parsed_content_list)

    elif instruction.startswith('Q:'):
        response = get_qa_response(instruction, image_input)
        return image_input, response

    # parsed_content_list = str(parsed_content_list)
    # convert xywh to yxhw
    label_coordinates_yxhw = {}
    for key, val in label_coordinates.items():
        if val[2] < 0 or val[3] < 0:
            continue
        label_coordinates_yxhw[key] = [val[1], val[0], val[3], val[2]]
    image_som = plot_boxes_with_marks(image_input.copy(), [val for key, val in label_coordinates_yxhw.items()], som_generator, edgecolor=(255,0,0), fn_save=None, normalized_to_pixel=False)

    # convert xywh to xyxy
    for key, val in label_coordinates.items():
        label_coordinates[key] = [val[0], val[1], val[0] + val[2], val[1] + val[3]]

    # normalize label_coordinates
    for key, val in label_coordinates.items():
        label_coordinates[key] = [val[0] / image_input.size[0], val[1] / image_input.size[1], val[2] / image_input.size[0], val[3] / image_input.size[1]]
    
    magma_response = get_som_response(instruction, image_som)
    print("magma repsonse: ", magma_response)

    # map magma_response into the mark id
    mark_id = extract_mark_id(magma_response)
    if mark_id is not None:
        if str(mark_id) in label_coordinates: