1.实现016新的导出rknn模型脚本\n2.017实现使用rknn模型的脚本
This commit is contained in:
parent
10761b756c
commit
1ad4413b9b
304
016新的导出rknn模型脚本.py
Normal file
304
016新的导出rknn模型脚本.py
Normal file
@ -0,0 +1,304 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
|
||||
"""
|
||||
PaddleOCR ONNX to RKNN Model Converter
|
||||
将PaddleOCR的文本检测和文本识别ONNX模型转换为RKNN模型
|
||||
"""
|
||||
|
||||
import os
|
||||
import numpy as np
|
||||
from rknn.api import RKNN
|
||||
import argparse
|
||||
|
||||
class PaddleOCRConverter:
|
||||
def __init__(self):
|
||||
self.rknn = RKNN(verbose=True)
|
||||
|
||||
def get_quantized_dtype(self, platform, quantize):
|
||||
"""
|
||||
根据平台和是否量化选择合适的quantized_dtype
|
||||
"""
|
||||
if platform.lower() in ['rk3588']:
|
||||
# RK3588支持的量化类型
|
||||
return 'w8a8' if quantize else None
|
||||
else:
|
||||
# 其他平台 (rk3566, rk3568等)
|
||||
return 'w8a8' if quantize else 'w8a16'
|
||||
|
||||
def convert_detection_model(self, onnx_path, rknn_path, platform='rk3588', quantize=False, dataset_path=None):
|
||||
"""
|
||||
转换文本检测模型
|
||||
Args:
|
||||
onnx_path: 输入的ONNX模型路径
|
||||
rknn_path: 输出的RKNN模型路径
|
||||
platform: 目标平台 (rk3588, rk3566, rk3568等)
|
||||
quantize: 是否进行量化
|
||||
dataset_path: 量化数据集路径
|
||||
"""
|
||||
print(f"开始转换文本检测模型: {onnx_path}")
|
||||
|
||||
# 获取适合的quantized_dtype
|
||||
quantized_dtype = self.get_quantized_dtype(platform, quantize)
|
||||
|
||||
# 配置RKNN
|
||||
config_params = {
|
||||
# 'mean_values': [[123.675, 116.28, 103.53]], # PaddleOCR标准化参数
|
||||
# 'std_values': [[58.395, 57.12, 57.375]],
|
||||
'target_platform': platform
|
||||
}
|
||||
|
||||
# 只有在需要时才添加quantized_dtype参数
|
||||
if quantized_dtype is not None:
|
||||
config_params['quantized_dtype'] = quantized_dtype
|
||||
|
||||
self.rknn.config(**config_params)
|
||||
|
||||
# 加载ONNX模型
|
||||
print("加载ONNX模型...")
|
||||
ret = self.rknn.load_onnx(model=onnx_path)
|
||||
if ret != 0:
|
||||
print(f"加载ONNX模型失败: {ret}")
|
||||
return False
|
||||
|
||||
# 构建模型
|
||||
print("构建RKNN模型...")
|
||||
ret = self.rknn.build(do_quantization=quantize, dataset=dataset_path)
|
||||
if ret != 0:
|
||||
print(f"构建RKNN模型失败: {ret}")
|
||||
return False
|
||||
|
||||
# 导出模型
|
||||
print(f"导出RKNN模型到: {rknn_path}")
|
||||
ret = self.rknn.export_rknn(rknn_path)
|
||||
if ret != 0:
|
||||
print(f"导出RKNN模型失败: {ret}")
|
||||
return False
|
||||
|
||||
print("文本检测模型转换完成!")
|
||||
return True
|
||||
|
||||
def convert_recognition_model(self, onnx_path, rknn_path, platform='rk3588', quantize=False, dataset_path=None):
|
||||
"""
|
||||
转换文本识别模型
|
||||
Args:
|
||||
onnx_path: 输入的ONNX模型路径
|
||||
rknn_path: 输出的RKNN模型路径
|
||||
platform: 目标平台
|
||||
quantize: 是否进行量化
|
||||
dataset_path: 量化数据集路径
|
||||
"""
|
||||
print(f"开始转换文本识别模型: {onnx_path}")
|
||||
|
||||
# 重新初始化RKNN实例
|
||||
self.rknn.release()
|
||||
self.rknn = RKNN(verbose=True)
|
||||
|
||||
# 获取适合的quantized_dtype
|
||||
quantized_dtype = self.get_quantized_dtype(platform, quantize)
|
||||
|
||||
# 配置RKNN (文本识别模型的预处理参数)
|
||||
config_params = {
|
||||
# 'mean_values': [[127.5, 127.5, 127.5]], # 文本识别模型标准化参数
|
||||
# 'std_values': [[127.5, 127.5, 127.5]],
|
||||
'target_platform': platform
|
||||
}
|
||||
|
||||
# 只有在需要时才添加quantized_dtype参数
|
||||
if quantized_dtype is not None:
|
||||
config_params['quantized_dtype'] = quantized_dtype
|
||||
|
||||
self.rknn.config(**config_params)
|
||||
|
||||
# 加载ONNX模型
|
||||
print("加载ONNX模型...")
|
||||
ret = self.rknn.load_onnx(model=onnx_path)
|
||||
if ret != 0:
|
||||
print(f"加载ONNX模型失败: {ret}")
|
||||
return False
|
||||
|
||||
# 构建模型
|
||||
print("构建RKNN模型...")
|
||||
ret = self.rknn.build(do_quantization=quantize, dataset=dataset_path)
|
||||
if ret != 0:
|
||||
print(f"构建RKNN模型失败: {ret}")
|
||||
return False
|
||||
|
||||
# 导出模型
|
||||
print(f"导出RKNN模型到: {rknn_path}")
|
||||
ret = self.rknn.export_rknn(rknn_path)
|
||||
if ret != 0:
|
||||
print(f"导出RKNN模型失败: {ret}")
|
||||
return False
|
||||
|
||||
print("文本识别模型转换完成!")
|
||||
return True
|
||||
|
||||
def generate_quantization_dataset(self, input_shape, num_samples=100, output_path="./dataset.txt"):
|
||||
"""
|
||||
生成量化数据集
|
||||
Args:
|
||||
input_shape: 输入形状 (N, C, H, W)
|
||||
num_samples: 样本数量
|
||||
output_path: 输出路径
|
||||
"""
|
||||
print(f"生成量化数据集: {output_path}")
|
||||
|
||||
# 创建数据集目录
|
||||
dataset_dir = os.path.dirname(output_path)
|
||||
if dataset_dir and not os.path.exists(dataset_dir):
|
||||
os.makedirs(dataset_dir)
|
||||
|
||||
# 生成随机数据并保存
|
||||
dataset_files = []
|
||||
for i in range(num_samples):
|
||||
# 生成随机数据 (0-255)
|
||||
data = np.random.randint(0, 256, input_shape, dtype=np.uint8)
|
||||
|
||||
# 保存为npy文件
|
||||
npy_path = f"{dataset_dir}/sample_{i}.npy"
|
||||
np.save(npy_path, data)
|
||||
dataset_files.append(npy_path)
|
||||
|
||||
# 创建数据集列表文件
|
||||
with open(output_path, 'w') as f:
|
||||
for file_path in dataset_files:
|
||||
f.write(f"{file_path}\n")
|
||||
|
||||
print(f"生成了 {num_samples} 个样本的量化数据集")
|
||||
return output_path
|
||||
|
||||
def test_model(self, rknn_path, input_shape):
|
||||
"""
|
||||
测试转换后的RKNN模型
|
||||
Args:
|
||||
rknn_path: RKNN模型路径
|
||||
input_shape: 输入形状
|
||||
"""
|
||||
print(f"测试RKNN模型: {rknn_path}")
|
||||
|
||||
# 重新初始化RKNN实例用于推理
|
||||
test_rknn = RKNN(verbose=False)
|
||||
|
||||
# 加载RKNN模型
|
||||
ret = test_rknn.load_rknn(rknn_path)
|
||||
if ret != 0:
|
||||
print(f"加载RKNN模型失败: {ret}")
|
||||
return False
|
||||
|
||||
# 初始化运行环境
|
||||
ret = test_rknn.init_runtime()
|
||||
if ret != 0:
|
||||
print(f"初始化运行环境失败: {ret}")
|
||||
return False
|
||||
|
||||
# 生成测试数据
|
||||
test_data = np.random.randint(0, 256, input_shape, dtype=np.uint8)
|
||||
|
||||
# 进行推理
|
||||
outputs = test_rknn.inference(inputs=[test_data])
|
||||
|
||||
print(f"推理成功! 输出形状: {[output.shape for output in outputs]}")
|
||||
|
||||
# 释放资源
|
||||
test_rknn.release()
|
||||
return True
|
||||
|
||||
def release(self):
|
||||
"""释放RKNN资源"""
|
||||
self.rknn.release()
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='将PaddleOCR ONNX模型转换为RKNN模型')
|
||||
parser.add_argument('--det_onnx', type=str, required=True, help='文本检测ONNX模型路径')
|
||||
parser.add_argument('--rec_onnx', type=str, required=True, help='文本识别ONNX模型路径')
|
||||
parser.add_argument('--output_dir', type=str, default='./rknn_models', help='RKNN模型输出目录')
|
||||
parser.add_argument('--platform', type=str, default='rk3588', help='目标平台')
|
||||
parser.add_argument('--quantize', action='store_true', help='是否进行量化')
|
||||
parser.add_argument('--test', action='store_true', help='是否测试转换后的模型')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# 创建输出目录
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
|
||||
# 初始化转换器
|
||||
converter = PaddleOCRConverter()
|
||||
|
||||
try:
|
||||
# 定义输入形状
|
||||
det_input_shape = (1, 3, 640, 640) # 文本检测模型输入形状
|
||||
rec_input_shape = (1, 3, 48, 320) # 文本识别模型输入形状
|
||||
|
||||
# 生成量化数据集（如果需要量化）
|
||||
det_dataset_path = None
|
||||
rec_dataset_path = None
|
||||
|
||||
if args.quantize:
|
||||
print("生成量化数据集...")
|
||||
det_dataset_path = converter.generate_quantization_dataset(
|
||||
det_input_shape, num_samples=100,
|
||||
output_path=os.path.join(args.output_dir, "det_dataset.txt")
|
||||
)
|
||||
rec_dataset_path = converter.generate_quantization_dataset(
|
||||
rec_input_shape, num_samples=100,
|
||||
output_path=os.path.join(args.output_dir, "rec_dataset.txt")
|
||||
)
|
||||
|
||||
# 转换文本检测模型
|
||||
det_rknn_path = os.path.join(args.output_dir, "text_detection.rknn")
|
||||
success = converter.convert_detection_model(
|
||||
args.det_onnx, det_rknn_path, args.platform, args.quantize, det_dataset_path
|
||||
)
|
||||
|
||||
if not success:
|
||||
print("文本检测模型转换失败!")
|
||||
return
|
||||
|
||||
# 转换文本识别模型
|
||||
rec_rknn_path = os.path.join(args.output_dir, "text_recognition.rknn")
|
||||
success = converter.convert_recognition_model(
|
||||
args.rec_onnx, rec_rknn_path, args.platform, args.quantize, rec_dataset_path
|
||||
)
|
||||
|
||||
if not success:
|
||||
print("文本识别模型转换失败!")
|
||||
return
|
||||
|
||||
# 测试模型（如果指定）
|
||||
if args.test:
|
||||
print("\n开始测试转换后的模型...")
|
||||
converter.test_model(det_rknn_path, det_input_shape)
|
||||
converter.test_model(rec_rknn_path, rec_input_shape)
|
||||
|
||||
print(f"\n所有模型转换完成!")
|
||||
print(f"文本检测模型: {det_rknn_path}")
|
||||
print(f"文本识别模型: {rec_rknn_path}")
|
||||
|
||||
finally:
|
||||
converter.release()
|
||||
|
||||
if __name__ == "__main__":
|
||||
# 如果直接运行脚本，提供示例用法
|
||||
'''
|
||||
python 016新的导出rknn模型脚本.py --det_onnx /home/admin-root/haotian/康达瑞贝斯机器狗/det_mobile_14_shape.onnx --rec_onnx /home/admin-root/haotian/康达瑞贝斯机 器狗/rec_mobile_14_shape.onnx --platform rk3588 --output_dir ./models
|
||||
'''
|
||||
if len(os.sys.argv) == 1:
|
||||
print("PaddleOCR ONNX to RKNN Model Converter")
|
||||
print("\n使用示例:")
|
||||
print("python paddleocr_converter.py \\")
|
||||
print(" --det_onnx ./text_detection.onnx \\")
|
||||
print(" --rec_onnx ./text_recognition.onnx \\")
|
||||
print(" --output_dir ./rknn_models \\")
|
||||
print(" --platform rk3588 \\")
|
||||
print(" --quantize \\")
|
||||
print(" --test")
|
||||
print("\n参数说明:")
|
||||
print(" --det_onnx: 文本检测ONNX模型路径")
|
||||
print(" --rec_onnx: 文本识别ONNX模型路径")
|
||||
print(" --output_dir: RKNN模型输出目录")
|
||||
print(" --platform: 目标平台 (rk3588, rk3566, rk3568等)")
|
||||
print(" --quantize: 是否进行量化")
|
||||
print(" --test: 是否测试转换后的模型")
|
||||
else:
|
||||
main()
|
||||
559
017能用的PaddleOCR_rknn脚本.py
Normal file
559
017能用的PaddleOCR_rknn脚本.py
Normal file
@ -0,0 +1,559 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
|
||||
"""
|
||||
RK3588 PaddleOCR RKNN推理程序
|
||||
使用转换后的RKNN模型在RK3588上进行OCR文本检测和识别
|
||||
"""
|
||||
|
||||
import cv2
|
||||
import yaml
|
||||
import numpy as np
|
||||
import math
|
||||
from rknn.api import RKNN
|
||||
import argparse
|
||||
import os
|
||||
import time
|
||||
|
||||
class RK3588OCR:
|
||||
def __init__(self, det_model_path, rec_model_path):
|
||||
"""
|
||||
初始化OCR推理器
|
||||
Args:
|
||||
det_model_path: 文本检测RKNN模型路径
|
||||
rec_model_path: 文本识别RKNN模型路径
|
||||
"""
|
||||
self.det_model_path = det_model_path
|
||||
self.rec_model_path = rec_model_path
|
||||
|
||||
# 初始化RKNN实例
|
||||
self.det_rknn = RKNN(verbose=False)
|
||||
self.rec_rknn = RKNN(verbose=False)
|
||||
|
||||
# 模型输入尺寸
|
||||
self.det_input_size = (640, 640)
|
||||
self.rec_input_size = (320, 48)
|
||||
|
||||
# 文本检测相关参数
|
||||
self.det_threshold = 0.3
|
||||
self.det_box_threshold = 0.6
|
||||
self.det_unclip_ratio = 1.5
|
||||
|
||||
# 加载模型
|
||||
self._load_models()
|
||||
|
||||
self.character = self.get_dict()
|
||||
|
||||
def get_dict(self, dict_path='/home/orangepi/Desktop/kangda_robotic_dog/机器狗后台服务/dict.yaml'):
|
||||
"""
|
||||
加载字典
|
||||
"""
|
||||
with open(dict_path, 'r', encoding='utf-8') as f:
|
||||
dict_rec = yaml.safe_load(f)
|
||||
return dict_rec.get('character_dict', [])
|
||||
|
||||
def _load_models(self):
|
||||
"""加载RKNN模型"""
|
||||
print("加载文本检测模型...")
|
||||
ret = self.det_rknn.load_rknn(self.det_model_path)
|
||||
if ret != 0:
|
||||
raise Exception(f"加载检测模型失败: {ret}")
|
||||
|
||||
ret = self.det_rknn.init_runtime(target='rk3588')
|
||||
if ret != 0:
|
||||
raise Exception(f"初始化检测模型运行环境失败: {ret}")
|
||||
|
||||
print("加载文本识别模型...")
|
||||
ret = self.rec_rknn.load_rknn(self.rec_model_path)
|
||||
if ret != 0:
|
||||
raise Exception(f"加载识别模型失败: {ret}")
|
||||
|
||||
ret = self.rec_rknn.init_runtime(target='rk3588')
|
||||
if ret != 0:
|
||||
raise Exception(f"初始化识别模型运行环境失败: {ret}")
|
||||
|
||||
print("模型加载完成!")
|
||||
|
||||
def resize_norm_img_det(self, img, input_shape=(640, 640)):
|
||||
"""
|
||||
检测模型的图像预处理 - 固定输入形状 [1, 3, 640, 640]
|
||||
"""
|
||||
h, w, _ = img.shape
|
||||
target_h, target_w = input_shape
|
||||
|
||||
# 计算缩放比例 - 保持宽高比
|
||||
ratio_h = target_h / h
|
||||
ratio_w = target_w / w
|
||||
ratio = min(ratio_h, ratio_w)
|
||||
|
||||
# 计算缩放后的尺寸
|
||||
new_h = int(h * ratio)
|
||||
new_w = int(w * ratio)
|
||||
|
||||
# 调整图像大小
|
||||
resized_img = cv2.resize(img, (new_w, new_h))
|
||||
|
||||
# 创建目标尺寸的图像，用灰色填充
|
||||
padded_img = np.ones((target_h, target_w, 3), dtype=np.float32) * 114.0 # 直接用float32
|
||||
|
||||
# 计算居中位置
|
||||
top = (target_h - new_h) // 2
|
||||
left = (target_w - new_w) // 2
|
||||
|
||||
# 将缩放后的图像放到居中位置
|
||||
padded_img[top:top+new_h, left:left+new_w] = resized_img.astype(np.float32)
|
||||
|
||||
# 归一化
|
||||
img = (padded_img / 255.0 - np.array([0.485, 0.456, 0.406], dtype=np.float32)) / np.array([0.229, 0.224, 0.225], dtype=np.float32)
|
||||
img = img.transpose(2, 0, 1).astype(np.float32)
|
||||
img = np.expand_dims(img, axis=0).astype(np.float32)
|
||||
|
||||
return img, ratio, (top, left)
|
||||
|
||||
def post_process_det(self, dt_boxes, ratio, padding_info, ori_shape):
|
||||
"""
|
||||
检测结果后处理 - 适配固定输入形状
|
||||
"""
|
||||
if dt_boxes is None:
|
||||
return None
|
||||
|
||||
ori_h, ori_w = ori_shape
|
||||
top, left = padding_info
|
||||
|
||||
# 将坐标从模型输出空间转换回原图空间
|
||||
dt_boxes[:, :, 0] = (dt_boxes[:, :, 0] - left) / ratio
|
||||
dt_boxes[:, :, 1] = (dt_boxes[:, :, 1] - top) / ratio
|
||||
|
||||
# 裁剪到原图范围内
|
||||
dt_boxes[:, :, 0] = np.clip(dt_boxes[:, :, 0], 0, ori_w)
|
||||
dt_boxes[:, :, 1] = np.clip(dt_boxes[:, :, 1], 0, ori_h)
|
||||
|
||||
return dt_boxes
|
||||
|
||||
def boxes_from_bitmap(self, pred, bitmap, dest_width, dest_height, max_candidates=1000, box_thresh=0.6):
|
||||
"""
|
||||
从位图中提取文本框
|
||||
"""
|
||||
bitmap = bitmap.astype(np.uint8)
|
||||
height, width = bitmap.shape
|
||||
|
||||
# 查找轮廓
|
||||
contours, _ = cv2.findContours(bitmap, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
|
||||
|
||||
num_contours = min(len(contours), max_candidates)
|
||||
boxes = []
|
||||
scores = []
|
||||
|
||||
for i in range(num_contours):
|
||||
contour = contours[i]
|
||||
points, sside = self.get_mini_boxes(contour)
|
||||
if sside < 5:
|
||||
continue
|
||||
|
||||
points = np.array(points)
|
||||
score = self.box_score_fast(pred, points.reshape(-1, 2))
|
||||
if box_thresh > score:
|
||||
continue
|
||||
|
||||
# 扩展box
|
||||
box = self.unclip(points, 1.5).reshape(-1, 1, 2)
|
||||
box, sside = self.get_mini_boxes(box)
|
||||
if sside < 5 + 2:
|
||||
continue
|
||||
|
||||
box = np.array(box)
|
||||
box[:, 0] = np.clip(box[:, 0] / width * dest_width, 0, dest_width)
|
||||
box[:, 1] = np.clip(box[:, 1] / height * dest_height, 0, dest_height)
|
||||
|
||||
boxes.append(box.astype(np.int16))
|
||||
scores.append(score)
|
||||
|
||||
return np.array(boxes), scores
|
||||
|
||||
def get_mini_boxes(self, contour):
|
||||
"""获取最小外接矩形"""
|
||||
bounding_box = cv2.minAreaRect(contour)
|
||||
points = sorted(list(cv2.boxPoints(bounding_box)), key=lambda x: x[0])
|
||||
|
||||
index_1, index_2, index_3, index_4 = 0, 1, 2, 3
|
||||
if points[1][1] > points[0][1]:
|
||||
index_1 = 0
|
||||
index_4 = 1
|
||||
else:
|
||||
index_1 = 1
|
||||
index_4 = 0
|
||||
|
||||
if points[3][1] > points[2][1]:
|
||||
index_2 = 2
|
||||
index_3 = 3
|
||||
else:
|
||||
index_2 = 3
|
||||
index_3 = 2
|
||||
|
||||
box = [points[index_1], points[index_2], points[index_3], points[index_4]]
|
||||
return box, min(bounding_box[1])
|
||||
|
||||
def box_score_fast(self, bitmap, _box):
|
||||
"""快速计算box得分"""
|
||||
h, w = bitmap.shape[:2]
|
||||
box = _box.copy()
|
||||
xmin = np.clip(np.floor(box[:, 0].min()).astype(int), 0, w - 1)
|
||||
xmax = np.clip(np.ceil(box[:, 0].max()).astype(int), 0, w - 1)
|
||||
ymin = np.clip(np.floor(box[:, 1].min()).astype(int), 0, h - 1)
|
||||
ymax = np.clip(np.ceil(box[:, 1].max()).astype(int), 0, h - 1)
|
||||
|
||||
mask = np.zeros((ymax - ymin + 1, xmax - xmin + 1), dtype=np.uint8)
|
||||
box[:, 0] = box[:, 0] - xmin
|
||||
box[:, 1] = box[:, 1] - ymin
|
||||
cv2.fillPoly(mask, box.reshape(1, -1, 2).astype(np.int32), 1)
|
||||
|
||||
return cv2.mean(bitmap[ymin:ymax + 1, xmin:xmax + 1], mask)[0]
|
||||
|
||||
def unclip(self, box, unclip_ratio):
|
||||
"""扩展文本框"""
|
||||
from shapely.geometry import Polygon
|
||||
import pyclipper
|
||||
|
||||
poly = Polygon(box)
|
||||
distance = poly.area * unclip_ratio / poly.length
|
||||
|
||||
offset = pyclipper.PyclipperOffset()
|
||||
offset.AddPath(box, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
|
||||
expanded = offset.Execute(distance)
|
||||
|
||||
if len(expanded) == 0:
|
||||
return box
|
||||
else:
|
||||
return np.array(expanded[0])
|
||||
|
||||
def resize_norm_img_rec(self, img, input_shape=(320, 48)):
|
||||
"""
|
||||
识别模型的图像预处理 - 固定输入形状 [1, 3, 48, 320]
|
||||
"""
|
||||
target_w, target_h = input_shape # 注意：宽度在前
|
||||
|
||||
h, w = img.shape[:2]
|
||||
|
||||
# 计算缩放比例，保持宽高比
|
||||
ratio_h = target_h / h
|
||||
ratio_w = target_w / w
|
||||
ratio = min(ratio_h, ratio_w)
|
||||
|
||||
# 计算缩放后的尺寸
|
||||
new_h = int(h * ratio)
|
||||
new_w = int(w * ratio)
|
||||
|
||||
# 调整图像大小
|
||||
resized_image = cv2.resize(img, (new_w, new_h))
|
||||
|
||||
# 创建目标尺寸的图像，用黑色填充
|
||||
padded_image = np.zeros((target_h, target_w, 3), dtype=np.float32) # 直接用float32
|
||||
|
||||
# 将缩放后的图像放到左上角（识别模型通常左对齐）
|
||||
padded_image[:new_h, :new_w] = resized_image.astype(np.float32)
|
||||
|
||||
# 归一化
|
||||
# padded_image = (padded_image / 255.0 - np.array([0.485, 0.456, 0.406], dtype=np.float32)) / np.array([0.229, 0.224, 0.225], dtype=np.float32)
|
||||
|
||||
# 不缩放反而会将识别结果再移后一个??
|
||||
padded_image = (padded_image / 255.0).astype(np.float32)
|
||||
padded_image = padded_image.transpose((2, 0, 1)).astype(np.float32)
|
||||
|
||||
return np.expand_dims(padded_image, axis=0).astype(np.float32)
|
||||
|
||||
def decode_rec_result(self, preds_prob):
|
||||
"""
|
||||
解码识别结果
|
||||
"""
|
||||
|
||||
# preds_idx = preds_idx[0]
|
||||
preds_prob = preds_prob[0]
|
||||
|
||||
preds_idx = np.argmax(preds_prob, axis=1)
|
||||
preds_prob = np.max(preds_prob, axis=1)
|
||||
|
||||
# CTC解码
|
||||
last_idx = 0
|
||||
preds_text = []
|
||||
preds_conf = []
|
||||
|
||||
|
||||
|
||||
# print("preds_id", len(preds_idx[0]))
|
||||
|
||||
for i, idx in enumerate(preds_idx):
|
||||
if idx != last_idx and idx != 0: # 0是blank
|
||||
if idx < len(self.character):
|
||||
# print("self.character[idx]", self.character[idx])
|
||||
# print("preds_prob[i]", preds_prob[i])
|
||||
preds_text.append(self.character[idx])
|
||||
preds_conf.append(preds_prob[i])
|
||||
last_idx = idx
|
||||
|
||||
text = ''.join(preds_text)
|
||||
conf = np.mean(preds_conf) if preds_conf else 0.0
|
||||
|
||||
return text, conf
|
||||
|
||||
def detect_text(self, image):
|
||||
"""
|
||||
文本检测 - 适配固定输入形状 [1, 3, 640, 640]
|
||||
"""
|
||||
ori_h, ori_w = image.shape[:2]
|
||||
|
||||
# 预处理
|
||||
det_img, ratio, padding_info = self.resize_norm_img_det(image)
|
||||
|
||||
# 推理
|
||||
det_output = self.det_rknn.inference(inputs=[det_img], data_format="nchw")[0]
|
||||
|
||||
# 后处理
|
||||
mask = det_output[0, 0, :, :]
|
||||
threshold = 0.3
|
||||
bitmap = (mask > threshold).astype(np.uint8) * 255
|
||||
|
||||
# 从位图中提取文本框（坐标是在640x640空间中的）
|
||||
boxes, scores = self.boxes_from_bitmap(mask, bitmap, 640, 640)
|
||||
|
||||
# 将坐标转换回原图空间
|
||||
if len(boxes) > 0:
|
||||
boxes = self.post_process_det(boxes, ratio, padding_info, (ori_h, ori_w))
|
||||
|
||||
|
||||
|
||||
print("*"*100, len(boxes))
|
||||
|
||||
return boxes, scores
|
||||
|
||||
def visualize_det_results(self, image_path, boxes):
|
||||
image = cv2.imread(image_path)
|
||||
|
||||
for box in boxes:
|
||||
box = np.array(box, dtype=np.int32)
|
||||
cv2.polylines(image, [box], True, (0, 255, 0), 2)
|
||||
cv2.imwrite('./visual_det.jpg', image)
|
||||
|
||||
def recognize_text(self, image):
|
||||
"""
|
||||
文本识别
|
||||
"""
|
||||
# 预处理
|
||||
rec_img = self.resize_norm_img_rec(image)
|
||||
|
||||
# 推理
|
||||
rec_output = self.rec_rknn.inference(inputs=[rec_img], data_format="nchw")
|
||||
|
||||
# 解码
|
||||
text, conf = self.decode_rec_result(rec_output[0])
|
||||
|
||||
# print("")
|
||||
|
||||
return text, conf
|
||||
|
||||
def get_rotate_crop_image(self, img, points):
|
||||
"""
|
||||
根据四个点坐标裁剪并矫正图像
|
||||
"""
|
||||
img_crop_width = int(
|
||||
max(
|
||||
np.linalg.norm(points[0] - points[1]),
|
||||
np.linalg.norm(points[2] - points[3])))
|
||||
img_crop_height = int(
|
||||
max(
|
||||
np.linalg.norm(points[0] - points[3]),
|
||||
np.linalg.norm(points[1] - points[2])))
|
||||
pts_std = np.float32([[0, 0], [img_crop_width, 0],
|
||||
[img_crop_width, img_crop_height],
|
||||
[0, img_crop_height]])
|
||||
M = cv2.getPerspectiveTransform(points, pts_std)
|
||||
dst_img = cv2.warpPerspective(
|
||||
img,
|
||||
M, (img_crop_width, img_crop_height),
|
||||
borderMode=cv2.BORDER_REPLICATE,
|
||||
flags=cv2.INTER_CUBIC)
|
||||
dst_img_height, dst_img_width = dst_img.shape[0:2]
|
||||
if dst_img_height * 1.0 / dst_img_width >= 1.5:
|
||||
dst_img = np.rot90(dst_img)
|
||||
return dst_img
|
||||
|
||||
def ocr(self, image_path):
|
||||
"""
|
||||
完整的OCR流程
|
||||
"""
|
||||
# 读取图像
|
||||
image = cv2.imread(image_path)
|
||||
if image is None:
|
||||
return []
|
||||
|
||||
# 1. 文本检测
|
||||
dt_boxes, scores = self.detect_text(image)
|
||||
|
||||
# 可视化检测框
|
||||
self.visualize_det_results(image_path, dt_boxes)
|
||||
|
||||
if dt_boxes is None or len(dt_boxes) == 0:
|
||||
return []
|
||||
|
||||
# 2. 文本识别
|
||||
ocr_results = []
|
||||
|
||||
text_list = []
|
||||
confidence_list = []
|
||||
for i, box in enumerate(dt_boxes):
|
||||
# 裁剪文本区域
|
||||
box_points = box.astype(np.float32)
|
||||
crop_img = self.get_rotate_crop_image(image, box_points)
|
||||
|
||||
# 识别文本
|
||||
text, conf = self.recognize_text(crop_img)
|
||||
|
||||
if conf > 0.4: # 置信度过滤
|
||||
ocr_results.append({
|
||||
'text': text,
|
||||
'confidence': conf,
|
||||
'box': box.tolist(),
|
||||
'score': scores[i] if i < len(scores) else 0.0
|
||||
})
|
||||
|
||||
text_list.append(text)
|
||||
confidence_list.append(round(conf.item(), 2))
|
||||
|
||||
# return ocr_results
|
||||
return [text_list, confidence_list]
|
||||
|
||||
def release(self):
|
||||
"""释放资源"""
|
||||
self.det_rknn.release()
|
||||
self.rec_rknn.release()
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='RK3588 PaddleOCR RKNN推理')
|
||||
parser.add_argument('--det_model', type=str, required=True, help='文本检测RKNN模型路径')
|
||||
parser.add_argument('--rec_model', type=str, required=True, help='文本识别RKNN模型路径')
|
||||
parser.add_argument('--image', type=str, help='输入图像路径')
|
||||
parser.add_argument('--video', type=str, help='输入视频路径')
|
||||
parser.add_argument('--camera', type=int, help='摄像头设备ID')
|
||||
parser.add_argument('--output', type=str, help='输出路径')
|
||||
parser.add_argument('--show', action='store_true', help='显示结果')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# 检查模型文件
|
||||
if not os.path.exists(args.det_model):
|
||||
print(f"检测模型文件不存在: {args.det_model}")
|
||||
return
|
||||
|
||||
if not os.path.exists(args.rec_model):
|
||||
print(f"识别模型文件不存在: {args.rec_model}")
|
||||
return
|
||||
|
||||
# 初始化OCR
|
||||
print("初始化RK3588 OCR...")
|
||||
ocr = RK3588OCR(args.det_model, args.rec_model)
|
||||
|
||||
try:
|
||||
if args.image:
|
||||
# 图像模式
|
||||
print(f"处理图像: {args.image}")
|
||||
|
||||
|
||||
# 进行OCR
|
||||
start_time = time.time()
|
||||
text, confidence = ocr.ocr(args.image)
|
||||
|
||||
print("text", text)
|
||||
print("confidence", confidence)
|
||||
|
||||
total_time = time.time() - start_time
|
||||
|
||||
# 打印结果
|
||||
print(f"\n总耗时: {total_time:.3f}s")
|
||||
print(f"识别结果:")
|
||||
for i in range(len(text)):
|
||||
print(f"{i+1}. 文本: '{text[i]}', 置信度: {confidence[i]:.3f}")
|
||||
|
||||
# 绘制结果
|
||||
# annotated_image = ocr.draw_results(image, results)
|
||||
|
||||
# 保存或显示结果
|
||||
# if args.output:
|
||||
# cv2.imwrite(args.output, annotated_image)
|
||||
# print(f"结果已保存到: {args.output}")
|
||||
|
||||
# if args.show:
|
||||
# cv2.imshow('OCR结果', annotated_image)
|
||||
# cv2.waitKey(0)
|
||||
# cv2.destroyAllWindows()
|
||||
|
||||
elif args.video or args.camera is not None:
|
||||
# 视频或摄像头模式
|
||||
if args.video:
|
||||
cap = cv2.VideoCapture(args.video)
|
||||
print(f"处理视频: {args.video}")
|
||||
else:
|
||||
cap = cv2.VideoCapture(args.camera)
|
||||
print(f"使用摄像头: {args.camera}")
|
||||
|
||||
if not cap.isOpened():
|
||||
print("无法打开视频源")
|
||||
return
|
||||
|
||||
while True:
|
||||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
break
|
||||
|
||||
# 进行OCR
|
||||
results = ocr.ocr(frame)
|
||||
|
||||
# 绘制结果
|
||||
annotated_frame = ocr.draw_results(frame, results)
|
||||
|
||||
# 显示结果
|
||||
cv2.imshow('Real-time OCR', annotated_frame)
|
||||
|
||||
# 按'q'退出
|
||||
if cv2.waitKey(1) & 0xFF == ord('q'):
|
||||
break
|
||||
|
||||
cap.release()
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
else:
|
||||
print("请指定输入源: --image, --video 或 --camera")
|
||||
|
||||
finally:
|
||||
ocr.release()
|
||||
|
||||
if __name__ == "__main__":
|
||||
# 如果直接运行脚本，提供示例用法
|
||||
'''
|
||||
启动命令示例
|
||||
python 010使用PaddleOCR_rknn.py --det_model ./text_detection.rknn --rec_model ./text_recognition.rknn --image ./image_test/632e474452d560edd7004f745319ff00_frame_000730.jpg --output ./result.jpg
|
||||
注:
|
||||
导出的额rknn模型没有进行归一化, 归一化参数mean=0,std=1
|
||||
'''
|
||||
if len(os.sys.argv) == 1:
|
||||
print("RK3588 PaddleOCR RKNN推理程序")
|
||||
print("\n使用示例:")
|
||||
print("# 处理单张图像")
|
||||
print("python rk3588_ocr.py \\")
|
||||
print(" --det_model ./rknn_models/text_detection.rknn \\")
|
||||
print(" --rec_model ./rknn_models/text_recognition.rknn \\")
|
||||
print(" --image ./test.jpg \\")
|
||||
print(" --output ./result.jpg \\")
|
||||
print(" --show")
|
||||
print()
|
||||
print("# 实时摄像头OCR")
|
||||
print("python rk3588_ocr.py \\")
|
||||
print(" --det_model ./rknn_models/text_detection.rknn \\")
|
||||
print(" --rec_model ./rknn_models/text_recognition.rknn \\")
|
||||
print(" --camera 0")
|
||||
print()
|
||||
print("# 处理视频文件")
|
||||
print("python rk3588_ocr.py \\")
|
||||
print(" --det_model ./rknn_models/text_detection.rknn \\")
|
||||
print(" --rec_model ./rknn_models/text_recognition.rknn \\")
|
||||
print(" --video ./input_video.mp4")
|
||||
else:
|
||||
main()
|
||||
Loading…
Reference in New Issue
Block a user