Split part of vision encoder to CPU and optimize Transpose ops.

convert.py

@@ -0,0 +1,344 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+from rknn.api import RKNN
+from math import exp
+from sys import exit
+
+import onnx
+import onnxscript
+
+batch_size = 1
+# embed_seq_len = 590
+
+prompt_tokens_list = [15, 17, 21, 25]
+
+encoder_seq_len_list = [577 + p for p in prompt_tokens_list]
+
+decoder_seq_len = 1
+
+# set current directory to the directory of this file
+import os
+os.chdir(os.path.dirname(os.path.abspath(__file__)))
+
+import subprocess
+import select
+
+def run_python_code(code):
+    # 启动子进程并执行代码
+    process = subprocess.Popen(
+        ['python', '-c', code],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True
+    )
+
+    # 实时读取子进程的输出和错误输出
+    while True:
+        reads = [process.stdout.fileno(), process.stderr.fileno()]
+        ret = select.select(reads, [], [])
+        
+        for fd in ret[0]:
+            if fd == process.stdout.fileno():
+                output = process.stdout.readline()
+                if output:
+                    print(output.strip())
+            if fd == process.stderr.fileno():
+                err = process.stderr.readline()
+                if err:
+                    print(f"Error: {err.strip()}")
+        
+        if process.poll() is not None:
+            break
+
+def convert_decoder():
+    rknn = RKNN(verbose=True)
+
+    ONNX_MODEL="decoder_model.onnx"
+    RKNN_MODEL=ONNX_MODEL.replace(".onnx",".rknn")
+    DATASET="dataset.txt"
+    QUANTIZE=False
+
+    # [[batch_size, encoder_seq_len], 
+    # [batch_size, encoder_seq_len, 768], 
+    # [batch_size, decoder_seq_len, 768]]
+    input_shapes =[[[batch_size, encoder_seq_len],
+                            [batch_size, encoder_seq_len, 768],
+                            [batch_size, decoder_seq_len, 768]] for encoder_seq_len in encoder_seq_len_list]
+    # pre-process config
+    print('--> Config model')
+    rknn.config(quantized_algorithm='normal', quantized_method='channel', target_platform='rk3588', optimization_level=3, single_core_mode=True,
+                dynamic_input=input_shapes)
+    print('done')
+
+    # Load ONNX model
+    print('--> Loading model')
+    ret = rknn.load_onnx(model=ONNX_MODEL,
+                         )
+    if ret != 0:
+        print('Load model failed!')
+        exit(ret)
+    print('done')
+
+    # Build model
+    print('--> Building model')
+    ret = rknn.build(do_quantization=QUANTIZE, dataset=DATASET, rknn_batch_size=None)
+    if ret != 0:
+        print('Build model failed!')
+        exit(ret)
+    print('done')
+
+    #export
+    print('--> Export RKNN model')
+    ret = rknn.export_rknn(RKNN_MODEL)
+    if ret != 0:
+        print('Export RKNN model failed!')
+        exit(ret)
+    print('done')
+
+def convert_encoder():
+    rknn = RKNN(verbose=True)
+
+    ONNX_MODEL="encoder_model.onnx"
+    RKNN_MODEL=ONNX_MODEL.replace(".onnx",".rknn")
+    DATASET="dataset.txt"
+    QUANTIZE=False
+
+    #[[batch_size, encoder_seq_len], [batch_size, encoder_seq_len, 768]]
+    input_shapes = [[[batch_size, encoder_seq_len], [batch_size, encoder_seq_len, 768]] for encoder_seq_len in encoder_seq_len_list]
+    # pre-process config
+    print('--> Config model')
+    rknn.config(quantized_algorithm='normal', quantized_method='channel', target_platform='rk3588', optimization_level=3, single_core_mode=True, dynamic_input=input_shapes)
+    print('done')
+
+    # Load ONNX model
+    print('--> Loading model')
+    ret = rknn.load_onnx(model=ONNX_MODEL
+                         )
+    if ret != 0:
+        print('Load model failed!')
+        exit(ret)
+    print('done')
+
+    # Build model
+    print('--> Building model')
+    ret = rknn.build(do_quantization=QUANTIZE, dataset=DATASET, rknn_batch_size=None)
+    if ret != 0:
+        print('Build model failed!')
+        exit(ret)
+    print('done')
+
+    # Export RKNN model
+    print('--> Export RKNN model')
+    ret = rknn.export_rknn(RKNN_MODEL)
+    if ret != 0:
+        print('Export RKNN model failed!')
+        exit(ret)
+    print('done')
+
+def convert_vision():
+    rknn = RKNN(verbose=True)
+
+    ONNX_MODEL="vision_encoder.onnx"
+    DATASET="dataset.txt"
+    QUANTIZE=False
+
+    # split the first Transformers block into a separate model because it's too large to fit in the rknn
+    onnx.utils.extract_model(ONNX_MODEL, "vision_encoder_part1.onnx", ['pixel_values'], ['/blocks.0/blocks.0.0/channel_block/channel_attn/Add_output_0'])
+    
+    ##### Build stage 1, this will crash the python process, so we need to run it in a separate process
+    code = f"""
+from rknn.api import RKNN
+rknn = RKNN(verbose=True)
+ONNX_MODEL="vision_encoder.onnx"
+RKNN_MODEL=ONNX_MODEL.replace(".onnx",".rknn")
+DATASET="dataset.txt"
+QUANTIZE=False
+batch_size = {batch_size}
+# pre-process config
+print('--> Config model')
+rknn.config(quantized_algorithm='normal', quantized_method='channel', target_platform='rk3588', optimization_level=3, single_core_mode=True)
+print('done')
+
+# Load ONNX model
+print('--> Loading model')
+ret = rknn.load_onnx(model=ONNX_MODEL,
+                        inputs=["pixel_values"],
+                        input_size_list=[[batch_size, 3, 768, 768]],
+                        )
+if ret != 0:
+    print('Load model failed!')
+    exit(ret)
+print('done')
+
+print('--> Building model stage 1')
+ret = rknn.build(do_quantization=QUANTIZE, dataset=DATASET, rknn_batch_size=None)
+if ret != 0:
+    print('Build model failed!')
+    exit(ret)
+print('done')
+    """
+    run_python_code(code)
+    print("Build stage 1 done")
+
+    intermidiate_model = onnx.load("check3_fuse_ops.onnx")
+
+    # fuse ops
+    from onnxscript.rewriter import pattern
+    import onnx.numpy_helper as onh
+    import numpy as np
+    def tp_rs_tp_rs_tp_pattern(op, input1, perm1, shape2, perm3, shape4, perm5):
+        i1 = op.Transpose(input1, perm=perm1)
+        i2 = op.Reshape(i1, shape2)
+        i3 = op.Transpose(i2, perm=perm3)
+        i4 = op.Reshape(i3, shape4)
+        i5 = op.Transpose(i4, perm=perm5)
+        return i5
+
+    def fused_pattern(op, input1, perm1, shape2, perm3, shape4, perm5):
+        rs1_shape = op.Constant(value=onh.from_array(np.array([input1.shape[0]* 3, input1.shape[1]//3, input1.shape[2], input1.shape[3]], dtype=np.int64)))
+        fi1 = op.Reshape(input1, rs1_shape)
+        fi2 = op.Transpose(fi1, perm=[0, 2, 1, 3])
+        elems = input1.shape[0] * input1.shape[1] * input1.shape[2] * input1.shape[3]
+        rs4_shape = op.Constant(value=onh.from_array(np.array([elems / 32 / 144, 32, 1, 144], dtype=np.int64)))
+        fi3 = op.Reshape(fi2, rs4_shape)
+        return fi3
+
+    rewrite_rule = pattern.RewriteRule(tp_rs_tp_rs_tp_pattern, fused_pattern)
+    rewrite_rule_set = pattern.RewriteRuleSet([rewrite_rule],commute=True)
+    fused_model = onnxscript.rewriter.rewrite(
+        intermidiate_model,
+        pattern_rewrite_rules=rewrite_rule_set
+    )
+    onnx.save(fused_model, "vision_encoder_part2.onnx")
+    ONNX_MODEL = "vision_encoder_part2.onnx"
+    RKNN_MODEL=ONNX_MODEL.replace(".onnx",".rknn")
+    del intermidiate_model
+    del fused_model
+
+
+    rknn = RKNN(verbose=True)
+
+    # pre-process config
+    print('--> Config model')
+    rknn.config(quantized_algorithm='normal', quantized_method='channel', target_platform='rk3588', optimization_level=3, single_core_mode=True)
+    print('done')
+
+    # Load ONNX model
+    print('--> Loading model')
+    ret = rknn.load_onnx(model="check3_fuse_ops.onnx",
+                         inputs=["/blocks.0/blocks.0.0/channel_block/channel_attn/Add_output_0-rs"],
+                         input_size_list=[[batch_size, 128, 1, 36864]],)
+    if ret != 0:
+        print('Load model failed!')
+        exit(ret)
+    print('done')
+
+    # Build model
+    print('--> Building model stage 2')
+    ret = rknn.build(do_quantization=QUANTIZE, dataset=DATASET, rknn_batch_size=None)
+    if ret != 0:
+        print('Build model failed!')
+        exit(ret)
+    print('done')
+
+    # Export RKNN model
+    print('--> Export RKNN model')
+    ret = rknn.export_rknn(RKNN_MODEL)
+    if ret != 0:
+        print('Export RKNN model failed!')
+        exit(ret)
+    print('done')
+    
+    
+    
+    
+
+
+
+def check_vision_model():
+    rknn = RKNN(verbose=True)
+
+    ONNX_MODEL="vision_encoder.onnx"
+    RKNN_MODEL=ONNX_MODEL.replace(".onnx",".rknn")
+    DATASET="dataset.txt"
+    QUANTIZE=False
+
+    # pre-process config
+    print('--> Config model')
+    rknn.config(quantized_algorithm='normal', quantized_method='channel', target_platform='rk3588', optimization_level=3, single_core_mode=True )
+    print('done')
+
+    # Load ONNX model
+    print('--> Loading model')
+    ret = rknn.load_onnx(model=ONNX_MODEL,
+                         inputs=["pixel_values"],
+                         input_size_list=[[batch_size, 3, vision_size[0], vision_size[1]]],
+                         )
+    if ret != 0:
+        print('Load model failed!')
+        exit(ret)
+    print('done')
+    
+    # Build model
+    print('--> Building model')
+    ret = rknn.build(do_quantization=QUANTIZE, dataset=DATASET, rknn_batch_size=None)
+    if ret != 0:
+        print('Build model failed!')
+        exit(ret)
+    print('done')
+    
+    # Export RKNN model
+    print('--> Export RKNN model')
+    ret = rknn.export_rknn(RKNN_MODEL)
+    if ret != 0:
+        print('Export RKNN model failed!')
+        exit(ret)
+    print('done')
+
+    #init runtime
+    print('--> Init runtime environment')
+    ret = rknn.init_runtime(target='rk3588')
+    if ret != 0:
+        print('Init runtime environment failed!')
+        exit(ret)
+    print('done')
+
+    #precision check
+    print('--> Precision check')
+    ret = rknn.accuracy_analysis(inputs=["lena.png"], target='rk3588')
+    if ret != 0:
+        print('Precision check failed!')
+        exit(ret)
+    print('done')
+    
+    
+    
+
+
+import argparse
+# python convert.py <decoder|encoder|vision|all>
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("model", type=str, help="Model to convert")
+    parser.add_argument("--check", action="store_true", help="Check model")
+    args = parser.parse_args()
+    if args.model == "decoder":
+        convert_decoder()
+    elif args.model == "encoder":
+        convert_encoder()
+    # elif args.model == "embed":   # embed is faster with cpu
+    #     convert_embed()
+    elif args.model == "vision":
+        if args.check:
+            check_vision_model()
+        else:
+            convert_vision()
+    elif args.model == "all":
+        convert_decoder()
+        convert_encoder()
+        # convert_embed()
+        convert_vision()
+    else:
+        print("Invalid model")
+        exit(1)

rknnrun.py

@@ -0,0 +1,325 @@
+import random
+from rknnlite.api.rknn_lite import RKNNLite
+from transformers import AutoProcessor
+from PIL import Image, ImageDraw
+import numpy as np
+import onnxruntime as ort
+import time
+import matplotlib.pyplot as plt  
+import matplotlib.patches as patches  
+# set current working directory to the directory of this file
+import os
+os.chdir(os.path.dirname(os.path.abspath(__file__)))
+
+# 初始化总时间计数器
+total_time = 0
+
+# Initialize RKNNLite instances
+rknn_vision_encoder = RKNNLite(verbose=False)
+rknn_encoder = RKNNLite(verbose=False)
+rknn_decoder_prefill = RKNNLite(verbose=False)
+
+# Load RKNN models
+ret = rknn_vision_encoder.load_rknn('./vision_encoder_part2.rknn')
+ret = rknn_encoder.load_rknn('./encoder_model.rknn')
+ret = rknn_decoder_prefill.load_rknn('./decoder_model.rknn')
+
+# Init runtime environment for each model
+ret = rknn_vision_encoder.init_runtime()
+ret = rknn_encoder.init_runtime()
+ret = rknn_decoder_prefill.init_runtime()
+
+text_embed = ort.InferenceSession("embed_tokens_fp16.onnx", providers=['CPUExecutionProvider'])
+decoder_decode = ort.InferenceSession("decoder_model_merged_q4.onnx", providers=['CPUExecutionProvider'])
+vision_encoder = ort.InferenceSession("vision_encoder_part1.onnx", providers=['CPUExecutionProvider'])
+prompt_tokens_list = [15, 17, 21, 25]
+
+# 1. prepare inputs
+processor = AutoProcessor.from_pretrained("/home/firefly/mnt/zt-rk3588-nn/expr/Florence-2-base-ft", trust_remote_code=True)
+
+# 2. prepare image
+image = Image.open("./test.jpg")
+original_image = image.copy()
+original_size = image.size
+# resize image to 768x768
+image = image.resize((768, 768))
+# 3. prepare text
+prompt = "<MORE_DETAILED_CAPTION>"
+
+## try tokenize first
+input_tokens_len = processor.tokenizer(prompt, return_tensors="np")["input_ids"].shape[1]
+print("input_tokens_len: ", input_tokens_len)
+## select the closest greater value
+pad_to = 0
+for i in prompt_tokens_list:
+    if i >= input_tokens_len:
+        pad_to = i
+        break
+print("pad_to: ", pad_to)
+inputs = processor(text=prompt, images=image, return_tensors="np", do_resize=False, padding="max_length", max_length=pad_to + 577, truncation=True)
+for k, v in inputs.items():
+    print(k, v.shape)
+
+# 4. run vision encoder using RKNN
+start_time = time.time()
+image_features0 = vision_encoder.run(None, {
+    "pixel_values": inputs["pixel_values"]
+})[0]
+image_features = rknn_vision_encoder.inference(inputs=[image_features0.reshape(1, 128, 1, 36864)])[0]
+
+end_time = time.time()
+vision_encoder_time = (end_time - start_time) * 1000
+total_time += vision_encoder_time
+print(f"Vision encoder time: {vision_encoder_time:.2f} ms")
+print(image_features.shape)
+np.save("image_features.npy", image_features)
+
+# 5. run text embed using RKNN
+start_time = time.time()
+inputs_embeds = text_embed.run(None, {
+    "input_ids": inputs["input_ids"]
+})[0]
+end_time = time.time()
+text_embed_time = (end_time - start_time) * 1000
+total_time += text_embed_time
+print(f"Text embed time: {text_embed_time:.2f} ms")
+print(inputs_embeds.shape)
+
+# 6. concat image features and text embed
+batch_size, image_token_length = image_features.shape[:-1]
+image_attention_mask = np.ones((batch_size, image_token_length))
+task_prefix_embeds = inputs_embeds
+task_prefix_attention_mask = np.ones((batch_size, task_prefix_embeds.shape[1]))
+if len(task_prefix_attention_mask.shape) == 3:
+    task_prefix_attention_mask = task_prefix_attention_mask[:, 0]
+inputs_embeds = np.concatenate([image_features, task_prefix_embeds], axis=1)
+attention_mask = np.concatenate([image_attention_mask, task_prefix_attention_mask], axis=1)
+
+# 6. run encoder using RKNN
+start_time = time.time()
+encoder_out = rknn_encoder.inference(inputs=[attention_mask.astype(np.int64),inputs_embeds])
+end_time = time.time()
+encoder_time = (end_time - start_time) * 1000
+total_time += encoder_time
+print(f"Encoder time: {encoder_time:.2f} ms")
+encoder_hidden_states = encoder_out[0]
+print(encoder_hidden_states.shape)
+
+# 7. run decoder prefill stage using RKNN
+start_time = time.time()
+next_token = processor.tokenizer.bos_token_id
+next_input_embeds = text_embed.run(None, {
+    "input_ids": np.array([[next_token]], dtype=np.int64)
+})[0]
+decoder_outs = rknn_decoder_prefill.inference(inputs=[attention_mask.astype(np.int64), encoder_hidden_states,inputs_embeds[:, -1:]])
+end_time = time.time()
+decoder_prefill_time = (end_time - start_time) * 1000
+total_time += decoder_prefill_time
+print(f"Decoder prefill time: {decoder_prefill_time:.2f} ms")
+# for output in decoder_outs:
+#     print(output.shape)
+
+encoder_kv = decoder_outs[1:]
+
+# 8. run decoder decode stage(autoregressive) (using onnxruntime)
+generated_tokens = []
+max_new_tokens = 512
+decoder_decode_total_time = 0
+while generated_tokens.__len__() < max_new_tokens:
+    # 获取上一步的输出
+    logits = decoder_outs[0]
+    decoder_kv = decoder_outs[1:]
+    
+    # 选择最后一个token的logits
+    next_token_logits = logits[:, -1, :]
+    
+    # 使用argmax选择下一个token (贪心算法)
+    next_token = np.argmax(next_token_logits, axis=-1)[0]
+    print("next_token: ", next_token)
+    # 将新生成的token添加到结果中
+    generated_tokens.append(next_token)
+
+    # 如果生成了结束符,则停止生成
+    if next_token == 2: # </s>
+        break
+    
+    #  准备下一步的输入
+    start_time = time.time()
+    next_input_embeds = text_embed.run(None, {
+        "input_ids": np.array([[next_token]], dtype=np.int64)
+    })[0]
+    end_time = time.time()
+    text_embed_time = (end_time - start_time) * 1000
+    decoder_decode_total_time += text_embed_time
+
+    # 运行decoder的decode阶段
+    start_time = time.time()
+    decoder_outs = decoder_decode.run(None, {
+        "use_cache_branch": np.array([True], dtype=np.bool_),
+        "inputs_embeds": next_input_embeds,
+        "encoder_hidden_states": encoder_hidden_states,
+        "encoder_attention_mask": attention_mask.astype(np.int64),
+        "past_key_values.0.decoder.key": decoder_kv[0],
+        "past_key_values.0.decoder.value": decoder_kv[1],
+        "past_key_values.0.encoder.key": encoder_kv[2],
+        "past_key_values.0.encoder.value": encoder_kv[3],
+        "past_key_values.1.decoder.key": decoder_kv[4],
+        "past_key_values.1.decoder.value": decoder_kv[5],
+        "past_key_values.1.encoder.key": encoder_kv[6],
+        "past_key_values.1.encoder.value": encoder_kv[7],
+        "past_key_values.2.decoder.key": decoder_kv[8],
+        "past_key_values.2.decoder.value": decoder_kv[9],
+        "past_key_values.2.encoder.key": encoder_kv[10],
+        "past_key_values.2.encoder.value": encoder_kv[11],
+        "past_key_values.3.decoder.key": decoder_kv[12],
+        "past_key_values.3.decoder.value": decoder_kv[13],
+        "past_key_values.3.encoder.key": encoder_kv[14],
+        "past_key_values.3.encoder.value": encoder_kv[15],
+        "past_key_values.4.decoder.key": decoder_kv[16],
+        "past_key_values.4.decoder.value": decoder_kv[17],
+        "past_key_values.4.encoder.key": encoder_kv[18],
+        "past_key_values.4.encoder.value": encoder_kv[19],
+        "past_key_values.5.decoder.key": decoder_kv[20],
+        "past_key_values.5.decoder.value": decoder_kv[21],
+        "past_key_values.5.encoder.key": encoder_kv[22],
+        "past_key_values.5.encoder.value": encoder_kv[23],
+    })
+    end_time = time.time()
+    decoder_decode_time = (end_time - start_time) * 1000
+    decoder_decode_total_time += decoder_decode_time
+
+total_time += decoder_decode_total_time
+print(f"Decoder decode total time: {decoder_decode_total_time:.2f} ms")
+
+# 将生成的tokens转换为文本
+print("generated_tokens: ", generated_tokens)
+generated_text = processor.batch_decode([generated_tokens], skip_special_tokens=False)[0]
+print("Generated Text:", generated_text)
+parsed_answer = processor.post_process_generation(generated_text, task=prompt.split(">")[0].strip() + ">", image_size=original_size)
+print("Parsed Answer:", parsed_answer)
+
+print(f"Total inference time: {total_time:.2f} ms")
+
+# postprocess
+from PIL import Image, ImageDraw, ImageFont
+
+from PIL import Image, ImageDraw, ImageFont
+
+def plot_bbox(image, data):
+    # Convert the image to a PIL Image if it's not already
+    if not isinstance(image, Image.Image):
+        image = Image.fromarray(image)
+    
+    # Create a drawing context
+    draw = ImageDraw.Draw(image)
+    
+    # Load a larger font
+    try:
+        font = ImageFont.truetype("arial.ttf", 20)  # 尝试加载Arial字体，大小为20
+    except IOError:
+        font = ImageFont.load_default().font_variant(size=20)  # 如果Arial不可用，使用默认字体并放大
+    
+    # Plot each bounding box
+    for bbox, label in zip(data['bboxes'], data['labels']):
+        # Unpack the bounding box coordinates
+        x1, y1, x2, y2 = bbox
+        # Draw the rectangle with thicker outline
+        draw.rectangle([x1, y1, x2, y2], outline="red", width=3)  # 增加线条宽度到3
+        
+        # Annotate the label
+        left, top, right, bottom = font.getbbox(label)
+        text_width = right - left
+        text_height = bottom - top
+        
+        # 增加文本背景框的大小
+        padding = 5
+        draw.rectangle([x1, y1 - text_height - padding*2, x1 + text_width + padding*2, y1], fill="red")
+        draw.text((x1 + padding, y1 - text_height - padding), label, fill="white", font=font)
+  
+    # Save the image
+    image.save("result_image.jpg")
+
+colormap = ['blue','orange','green','purple','brown','pink','gray','olive','cyan','red',
+            'lime','indigo','violet','aqua','magenta','coral','gold','tan','skyblue']
+
+def draw_polygons(image, prediction, fill_mask=False):  
+    """  
+    Draws segmentation masks with polygons on an image.  
+  
+    Parameters:  
+    - image_path: Path to the image file.  
+    - prediction: Dictionary containing 'polygons' and 'labels' keys.  
+                  'polygons' is a list of lists, each containing vertices of a polygon.  
+                  'labels' is a list of labels corresponding to each polygon.  
+    - fill_mask: Boolean indicating whether to fill the polygons with color.  
+    """  
+    # Load the image  
+   
+    draw = ImageDraw.Draw(image)  
+      
+   
+    # Set up scale factor if needed (use 1 if not scaling)  
+    scale = 1  
+      
+    # Iterate over polygons and labels  
+    for polygons, label in zip(prediction['polygons'], prediction['labels']):  
+        color = random.choice(colormap)  
+        fill_color = random.choice(colormap) if fill_mask else None  
+          
+        for _polygon in polygons:  
+            _polygon = np.array(_polygon).reshape(-1, 2)  
+            if len(_polygon) < 3:  
+                print('Invalid polygon:', _polygon)  
+                continue  
+              
+            _polygon = (_polygon * scale).reshape(-1).tolist()  
+              
+            # Draw the polygon  
+            if fill_mask:  
+                draw.polygon(_polygon, outline=color, fill=fill_color)  
+            else:  
+                draw.polygon(_polygon, outline=color)  
+              
+            # Draw the label text  
+            draw.text((_polygon[0] + 8, _polygon[1] + 2), label, fill=color)  
+  
+    # Save or display the image  
+    # image.show()  # Display the image  
+    # display(image)
+    image.save("result_image.jpg")
+
+
+
+def draw_ocr_bboxes(image, prediction, scale=1):
+    draw = ImageDraw.Draw(image)
+        
+    # Load a larger font
+    try:
+        font = ImageFont.truetype("arial.ttf", 18)  # 尝试加载Arial字体，大小为18
+    except IOError:
+        font = ImageFont.load_default().font_variant(size=18)  # 如果Arial不可用，使用默认字体并放大
+    bboxes, labels = prediction['quad_boxes'], prediction['labels']
+    for box, label in zip(bboxes, labels):
+        color = random.choice(colormap)
+        new_box = (np.array(box) * scale).tolist()
+        draw.polygon(new_box, width=3, outline=color)
+        draw.text((new_box[0]+8, new_box[1]+2),
+                    "{}".format(label),
+                    align="right",
+        
+                    fill=color)
+       
+    # display(image)
+    image.save("result_image.jpg")
+
+
+# draw_polygons(original_image, parsed_answer['<REFERRING_EXPRESSION_SEGMENTATION>'], fill_mask=True)
+# plot_bbox(original_image, parsed_answer[prompt.split(">")[0].strip() + ">"])
+# draw_ocr_bboxes(original_image, parsed_answer["<OCR_WITH_REGION>"], scale=1)
+
+
+
+# Release RKNNLite instances
+rknn_vision_encoder.release()
+rknn_encoder.release()
+rknn_decoder_prefill.release()