build and run success

Author: ZHEQIUSHUI

.gitignore

@@ -0,0 +1,6 @@
+.vscode
+weights
+__pycache__
+axpi_pro_bsp_sdk
+build*
+third_party

CMakeLists.txt

@@ -0,0 +1,81 @@
+# set cmake_install_prefix path
+if(NOT DEFINED CMAKE_INSTALL_PREFIX)
+    set(CMAKE_INSTALL_PREFIX "${CMAKE_BINARY_DIR}/install" CACHE PATH "Installation Directory")
+endif()
+
+message(STATUS "CMAKE_INSTALL_PREFIX = ${CMAKE_INSTALL_PREFIX}")
+
+# check if building type is not set
+if(NOT CMAKE_BUILD_TYPE)
+    set(CMAKE_BUILD_TYPE Release CACHE STRING "Choose the type of build" FORCE)
+endif()
+
+# bsp
+if(NOT BSP_MSP_DIR)
+    set(BSP_MSP_DIR ${CMAKE_SOURCE_DIR}/axpi_pro_bsp_sdk/msp/out)
+endif()
+
+message(STATUS "BSP_MSP_DIR = ${BSP_MSP_DIR}")
+
+# set(CMAKE_TOOLCHAIN_FILE ${CMAKE_SOURCE_DIR}/toolchains/aarch64-none-linux-gnu.toolchain.cmake)
+cmake_minimum_required(VERSION 3.0)
+project(OWLVIT-ONNX-AX650-CPP)
+
+# 添加 FLAGS 检查代码是否有明显 bug
+# include(cmake/overlook.cmake)
+
+add_compile_options(-std=c++17)
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+
+
+
+# link third-party library onnxruntime-win-x64-1.14.1 and opencv4
+if(NOT ONNXRUNTIME_DIR)
+    set(ONNXRUNTIME_DIR ${CMAKE_SOURCE_DIR}/third_party/onnxruntime)
+endif()
+
+message(STATUS "ONNXRUNTIME_DIR Path: ${ONNXRUNTIME_DIR}")
+
+include_directories(${ONNXRUNTIME_DIR}/include)
+link_directories("${ONNXRUNTIME_DIR}/lib")
+set(ONNXRUNTIME_LIB onnxruntime)
+
+if(NOT OpenCV_DIR)
+    set(OpenCV_DIR ${CMAKE_SOURCE_DIR}/third_party/opencv-mobile-4.6.0-ubuntu-2004/lib/cmake/opencv4)
+endif()
+
+message(STATUS "OPENCV_DIR Path: ${OpenCV_DIR}")
+find_package(OpenCV REQUIRED)
+
+include_directories(${OpenCV_INCLUDE_DIRS})
+include_directories(${BSP_MSP_DIR}/include)
+link_directories(${BSP_MSP_DIR}/lib)
+
+option(BUILD_WITH_AX650 "build with ax650" OFF)
+
+if(BUILD_WITH_AX650)
+    set(AXERA_TARGET_CHIP "AX650")
+    add_compile_definitions(BUILD_WITH_AX650)
+    add_compile_definitions(AXERA_TARGET_CHIP_AX650)
+endif()
+
+message(STATUS "BUILD_WITH_AX650 : ${BUILD_WITH_AX650}")
+
+add_library(owlvit STATIC
+    src/Runner/BaseRunner.cpp
+    src/Runner/ax_model_runner_ax650.cpp
+)
+
+add_executable(main src/main.cpp)
+
+target_link_libraries(main owlvit)
+
+target_link_libraries(main ${OpenCV_LIBS})
+target_link_libraries(main gomp ${ONNXRUNTIME_LIB})
+
+if(BUILD_WITH_AX650)
+    target_link_libraries(main ax_engine ax_interpreter ax_sys)
+endif()
+
+install(TARGETS main DESTINATION bin)

README.md

@@ -0,0 +1,55 @@
+# OWLVIT
+
+https://huggingface.co/docs/transformers/model_doc/owlvit
+
+
+
+<img src="ssd_horse.jpg" height="320" />  <img src="result.jpg" height="320" />
+
+## Build
+```
+mkdir build
+cd build
+```
+if x86 onnxruntime
+```
+cmake -DONNXRUNTIME_DIR=${onnxruntime_dir} -DOpenCV_DIR=${opencv_cmake_file_dir} ..
+```
+else if ax650
+```
+cmake -DONNXRUNTIME_DIR=${onnxruntime_dir} -DOpenCV_DIR=${opencv_cmake_file_dir} -DBSP_MSP_DIR=${msp_out_dir} -DBUILD_WITH_AX650=ON -DCMAKE_TOOLCHAIN_FILE=../toolchains/aarch64-none-linux-gnu.toolchain.cmake ..
+```
+```
+make -j4
+```
+aarch64-none-gnu library:\
+[onnxruntime](https://github.com/ZHEQIUSHUI/SAM-ONNX-AX650-CPP/releases/download/ax_models/onnxruntime-aarch64-none-gnu-1.16.0.zip)\
+[opencv](https://github.com/ZHEQIUSHUI/SAM-ONNX-AX650-CPP/releases/download/ax_models/libopencv-4.6-aarch64-none.zip)
+
+### run
+```
+/opt/test/owlvit # ./main --ienc owlvit-image.axmodel --tenc owlvit-text.onnx -d
+ owlvit-post.onnx -v vocab.txt -i ssd_horse.jpg -t text.txt
+Engine creating handle is done.
+Engine creating context is done.
+Engine get io info is done.
+Engine alloc io is done.
+[I][                            init][ 280]: BGR MODEL
+[I][              load_image_encoder][  17]: input size 768 768
+[I][              load_image_encoder][  29]: image feature len 442368
+[I][              load_image_encoder][  32]: pred box cnt  576
+[I][               load_text_encoder][ 141]: text feature len 512
+[I][                            main][ 116]: image_src [ssd_horse.jpg]
+[I][                            main][ 117]: text_src [text.txt]
+encode text Inference Cost time : 0.281856s
+[I][                          decode][ 239]: logits_size: 576
+[I][                          decode][ 239]: logits_size: 576
+[I][                          decode][ 239]: logits_size: 576
+[I][                          decode][ 239]: logits_size: 576
+post Inference Cost time : 0.0981112s
+a photo of person 268.899292 20.153463 88.163696 235.837906
+a photo of person 428.696014 123.745819 19.836823 55.102310
+horse 191.756058 55.418949 229.225601 318.581055
+a photo of car 0.000000 98.398750 145.470108 92.571877
+a photo of dog 145.470108 203.093140 57.306412 156.490570
+```

result.jpg

@@ -0,0 +1,32 @@
+import requests
+from PIL import Image
+import torch
+
+from transformers import OwlViTProcessor, OwlViTForObjectDetection
+
+processor = OwlViTProcessor.from_pretrained("google/owlvit-base-patch32")
+model = OwlViTForObjectDetection.from_pretrained("google/owlvit-base-patch32")
+
+url = "./test.jpg"
+image = Image.open(url)
+texts = [["a photo of people"]]
+inputs = processor(text=texts, images=image, return_tensors="pt")
+outputs = model(**inputs)
+
+# torch.onnx.export(model, (inputs["input_ids"],inputs["pixel_values"],inputs["attention_mask"]), "weights/owlvit.onnx", opset_version=14,
+#                   input_names=["input_ids","pixel_values","attention_mask"],
+#                   output_names=["logits","pred_boxes","text_embeds","image_embeds"])
+
+print(inputs)
+# print(processor)
+
+# Target image sizes (height, width) to rescale box predictions [batch_size, 2]
+target_sizes = torch.Tensor([image.size[::-1]])
+# Convert outputs (bounding boxes and class logits) to COCO API
+results = processor.post_process_object_detection(outputs=outputs, target_sizes=target_sizes, threshold=0.1)
+i = 0  # Retrieve predictions for the first image for the corresponding text queries
+text = texts[i]
+boxes, scores, labels = results[i]["boxes"], results[i]["scores"], results[i]["labels"]
+for box, score, label in zip(boxes, scores, labels):
+    box = [round(i, 2) for i in box.tolist()]
+    print(f"Detected {text[label]} with confidence {round(score.item(), 3)} at location {box}")

scripts/bpe_simple_vocab_16e6.txt.gz

@@ -0,0 +1,22 @@
+# optimum-cli export onnx --model google/owlvit-base-patch32 --task zero-shot-object-detection weights/ --opset 16
+
+import onnx
+
+input_path = "weights/owlvit.onnx"
+output_path = "weights/model-image.onnx"
+input_names = ["pixel_values"]
+output_names = ["image_embeds","pred_boxes"]
+
+onnx.utils.extract_model(input_path, output_path, input_names, output_names)
+
+# output_path = "weights/model-text.onnx"
+# input_names = ["input_ids","attention_mask"]
+# output_names = ["/owlvit/Div_output_0"]
+
+# onnx.utils.extract_model(input_path, output_path, input_names, output_names)
+
+output_path = "weights/model-post.onnx"
+input_names = ["image_embeds","input_ids","/owlvit/Div_output_0"]
+output_names = ["logits"]
+
+onnx.utils.extract_model(input_path, output_path, input_names, output_names)

scripts/hf_demo.py

@@ -0,0 +1,131 @@
+import onnxruntime
+
+class onnx_inferencer:
+
+    def __init__(self, model_path) -> None:
+        self.onnx_model_sess = onnxruntime.InferenceSession(model_path)
+        self.output_names = []
+        self.input_names = []
+        print(model_path)
+        for i in range(len(self.onnx_model_sess.get_inputs())):
+            self.input_names.append(self.onnx_model_sess.get_inputs()[i].name)
+            print("    input:", i,
+                  self.onnx_model_sess.get_inputs()[i].name,self.onnx_model_sess.get_inputs()[i].type,
+                  self.onnx_model_sess.get_inputs()[i].shape)
+
+        for i in range(len(self.onnx_model_sess.get_outputs())):
+            self.output_names.append(
+                self.onnx_model_sess.get_outputs()[i].name)
+            print("    output:", i,
+                  self.onnx_model_sess.get_outputs()[i].name,self.onnx_model_sess.get_outputs()[i].type,
+                  self.onnx_model_sess.get_outputs()[i].shape)
+        print("")
+
+    def get_input_count(self):
+        return len(self.input_names)
+
+    def get_input_shape(self, idx: int):
+        return self.onnx_model_sess.get_inputs()[idx].shape
+
+    def get_input_names(self):
+        return self.input_names
+
+    def get_output_count(self):
+        return len(self.output_names)
+
+    def get_output_shape(self, idx: int):
+        return self.onnx_model_sess.get_outputs()[idx].shape
+
+    def get_output_names(self):
+        return self.output_names
+
+    def inference(self, tensor):
+        return self.onnx_model_sess.run(
+            self.output_names, input_feed={self.input_names[0]: tensor})
+
+    def inference_multi_input(self, tensors: list):
+        inputs = dict()
+        for idx, tensor in enumerate(tensors):
+            inputs[self.input_names[idx]] = tensor
+        return self.onnx_model_sess.run(self.output_names, input_feed=inputs)
+
+backbone = onnx_inferencer("weights/owlvit-image.onnx")
+bert = onnx_inferencer("weights/owlvit-text.onnx")
+transformer = onnx_inferencer("weights/owlvit-post.onnx")
+
+import torchvision.transforms as T
+from tokenizer import build_tokenizer
+import torch
+import cv2
+import numpy as np
+from PIL import Image
+
+def load_image(image_path: str):
+    transform = T.Compose(
+        [
+            T.Resize([768,768]),
+            T.ToTensor(),
+            T.Normalize([0.48145466, 0.4578275, 0.40821073], [0.26862954,0.26130258,0.27577711]),
+        ]
+    )
+    image_source = Image.open(image_path).convert("RGB")
+    image = np.asarray(image_source)
+    image_transformed = transform(image_source)
+    return image, image_transformed
+
+
+tokenizer = build_tokenizer()
+
+BOX_TRESHOLD = 0.35
+TEXT_TRESHOLD = 0.25
+
+IMAGE_PATH = "./test.jpg"
+TEXT_PROMPT = ["football"]
+image_source, image = load_image(IMAGE_PATH)
+print(image.shape)
+image_embeds, pred_boxes = backbone.inference(image.unsqueeze(0).numpy())
+
+print(image_embeds[0].shape)
+
+input_ids = np.array([tokenizer.encode(t) for t in TEXT_PROMPT]).reshape(-1)
+print(input_ids)
+
+
+input_ids = np.pad([49406,*input_ids,49407],(0,16-len(input_ids)-2))
+print(input_ids)
+mask = (input_ids > 0).astype(np.int64)
+
+print(mask)
+
+text_embeds = bert.inference_multi_input([input_ids.reshape(1,16), mask.reshape(1,16)])[0].reshape(1,-1)
+print(text_embeds)
+logits = transformer.inference_multi_input([image_embeds[0].reshape(1,24,24,768),text_embeds,input_ids.reshape(1,16)])[0]
+
+logits = torch.Tensor(logits).sigmoid().numpy().reshape(-1)
+pred_boxes = pred_boxes.reshape(-1,4)
+
+print(logits.shape)
+print(pred_boxes.shape)
+
+# get idx of boxes with confidence > BOX_TRESHOLD
+idxs = np.where(logits > BOX_TRESHOLD)[0]
+if(len(idxs) == 0):
+    print("no boxes found")
+    exit()
+print(idxs)
+# print(logits[idx])
+# print(pred_boxes[idx][0])
+_h,_w,_ = image_source.shape
+image_source = cv2.cvtColor(image_source, cv2.COLOR_RGB2BGR)
+for idx in idxs:
+    print(idx,pred_boxes[idx])
+    xc,yc,w,h = pred_boxes[idx]
+    xc*=_w
+    yc*=_h
+    w*=_w
+    h*=_h
+
+    
+    cv2.rectangle(image_source,(int(xc-w/2),int(yc-h/2)),(int(xc+w/2),int(yc+h/2)),(0,0,255),2)
+cv2.imwrite("out.jpg", image_source)
+