Split part of vision encoder to CPU and optimize Transpose ops. (Reupload to correct path)

onnx/rknnrun.py

@@ -1,9 +1,12 @@
+import random
 from rknnlite.api.rknn_lite import RKNNLite
 from transformers import AutoProcessor
-from PIL import Image
+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__)))
@@ -17,7 +20,7 @@ rknn_encoder = RKNNLite(verbose=False)
 rknn_decoder_prefill = RKNNLite(verbose=False)
 
 # Load RKNN models
-ret = rknn_vision_encoder.load_rknn('./vision_encoder.rknn')
+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')
 
@@ -26,27 +29,44 @@ ret = rknn_vision_encoder.init_runtime()
 ret = rknn_encoder.init_runtime()
 ret = rknn_decoder_prefill.init_runtime()
 
-text_embed = ort.InferenceSession("embed_tokens.onnx", providers=['CPUExecutionProvider'])
+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.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("./lena.png")
-
-# resize image to 512x512
-image = image.resize((512, 512))
+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>"
-inputs = processor(text=prompt, images=image, return_tensors="np", do_resize=False)
+
+## 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_features = rknn_vision_encoder.inference(inputs=[inputs["pixel_values"]], data_format='nchw')[0]
+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
@@ -87,6 +107,10 @@ 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
@@ -99,7 +123,7 @@ encoder_kv = decoder_outs[1:]
 
 # 8. run decoder decode stage(autoregressive) (using onnxruntime)
 generated_tokens = []
-max_new_tokens = 32
+max_new_tokens = 512
 decoder_decode_total_time = 0
 while generated_tokens.__len__() < max_new_tokens:
     # 获取上一步的输出
@@ -111,7 +135,7 @@ while generated_tokens.__len__() < max_new_tokens:
     
     # 使用argmax选择下一个token (贪心算法)
     next_token = np.argmax(next_token_logits, axis=-1)[0]
-    # print("next_token: ", next_token)
+    print("next_token: ", next_token)
     # 将新生成的token添加到结果中
     generated_tokens.append(next_token)
 
@@ -119,7 +143,7 @@ while generated_tokens.__len__() < max_new_tokens:
     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)
@@ -171,11 +195,130 @@ print(f"Decoder decode total time: {decoder_decode_total_time:.2f} ms")
 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, image_size=(image.width, image.height))
+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()