onnx/rknnrun.py

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()