当前位置：首页 > news >正文

rknn yolo系列之量化前预处理，解决量化精度低以及出现类似未作nms的很多框子的问题

news 来源：原创 2024/9/19 9:37:51

rknn官方代码在处理onnx量化到rknn模型时并未提供脚本进行图像缩放，如果直接简单的将图片路径写入txt文件，然后进行量化输出的话就会导致输出的rknn模型文件推理精度低，或者没有精度的情况。
因此，参考

https://github.com/airockchip/rknn-toolkit2/blob/77b71094e08391c543d9c65fea5f7cf98cc16eee/rknpu2/examples/rknn_yolov5_demo/src/preprocess.cc#L17

里面的：

void letterbox(const cv::Mat &image, cv::Mat &padded_image, BOX_RECT &pads, const float scale, const cv::Size &target_size, const cv::Scalar &pad_color)
{// 调整图像大小cv::Mat resized_image;cv::resize(image, resized_image, cv::Size(), scale, scale);// 计算填充大小int pad_width = target_size.width - resized_image.cols;int pad_height = target_size.height - resized_image.rows;pads.left = pad_width / 2;pads.right = pad_width - pads.left;pads.top = pad_height / 2;pads.bottom = pad_height - pads.top;// 在图像周围添加填充cv::copyMakeBorder(resized_image, padded_image, pads.top, pads.bottom, pads.left, pads.right, cv::BORDER_CONSTANT, pad_color);
}

我改写了一版图像量化预处理的脚本。使用这个脚本执行完毕，会生成一个txt文件，然后把它复制到test.py里面进行替换就可以了。test.py见：

https://github.com/airockchip/rknn-toolkit2/blob/master/rknn-toolkit2/examples/onnx/yolov5/test.py

在这篇博客中，将展示如何使用Python和OpenCV实现图像调整和填充，并将处理后的图像保存到指定目录中。同时，我们将介绍如何使用shutil模块来清理目录内容。

环境准备

在开始之前，确保你已经安装了所需的Python库。你可以使用以下命令进行安装：

pip install opencv-python tqdm numpy

以下是实现图像调整、填充并保存结果的完整代码：

import os
import cv2
import numpy as np
from tqdm import tqdm
import shutilclass BOX_RECT:def __init__(self):self.left = 0self.right = 0self.top = 0self.bottom = 0def letterbox(image, target_size, pad_color=(114, 114, 114)):# 调整图像大小以适应目标尺寸的比例h, w = image.shape[:2]scale = min(target_size[0] / h, target_size[1] / w)new_w, new_h = int(scale * w), int(scale * h)resized_image = cv2.resize(image, (new_w, new_h))# 计算填充大小pad_width = target_size[1] - new_wpad_height = target_size[0] - new_hpads = BOX_RECT()pads.left = pad_width // 2pads.right = pad_width - pads.leftpads.top = pad_height // 2pads.bottom = pad_height - pads.top# 在图像周围添加填充padded_image = cv2.copyMakeBorder(resized_image, pads.top, pads.bottom, pads.left, pads.right, cv2.BORDER_CONSTANT, value=pad_color)return padded_image, pads, scale# 图像路径
images_path = "./train/images"
images_list = os.listdir(images_path)
cali_num = len(images_list)
jump_num = int(len(images_list) / cali_num)
cali_list = images_list[0::jump_num][0:cali_num]input_w = 320
input_h = input_w
target_size = (input_w, input_h)
pad_color = (0, 0, 0)save_txt = "dataset.txt"
save_path = "./calib"
absolute_save_path = os.path.abspath(save_path)# 如果保存路径存在则清空，不存在则创建
if os.path.exists(absolute_save_path):# 使用shutil来删除目录内容for filename in os.listdir(absolute_save_path):file_path = os.path.join(absolute_save_path, filename)try:if os.path.isfile(file_path) or os.path.islink(file_path):os.unlink(file_path)elif os.path.isdir(file_path):shutil.rmtree(file_path)except Exception as e:print(f'Failed to delete {file_path}. Reason: {e}')
else:os.makedirs(absolute_save_path)count = 0
with open(save_txt, "w") as f:for img in tqdm(cali_list):count += 1img_abs_path = os.path.join(images_path, img)img_suffix = os.path.splitext(img_abs_path)[-1]img_ndarry = cv2.imread(img_abs_path)processed_img, pads, scale = letterbox(img_ndarry, target_size, pad_color)cv2.imwrite(os.path.join(absolute_save_path, str(count) + img_suffix), processed_img)img_save_abs_path = os.path.join(absolute_save_path, str(count) + img_suffix)f.write(img_save_abs_path + "\n")

代码解释

1. 引入所需库

import os
import cv2
import numpy as np
from tqdm import tqdm
import shutil

我们使用了OpenCV进行图像处理，tqdm用于显示进度条，numpy用于数值计算，shutil用于文件操作。

2. 定义BOX_RECT类

	class BOX_RECT:def __init__(self):self.left = 0self.right = 0self.top = 0self.bottom = 0

BOX_RECT类用于存储图像填充的尺寸。

3. 定义letterbox函数

	def letterbox(image, target_size, pad_color=(114, 114, 114)):h, w = image.shape[:2]scale = min(target_size[0] / h, target_size[1] / w)new_w, new_h = int(scale * w), int(scale * h)resized_image = cv2.resize(image, (new_w, new_h))pad_width = target_size[1] - new_wpad_height = target_size[0] - new_hpads = BOX_RECT()pads.left = pad_width // 2pads.right = pad_width - pads.leftpads.top = pad_height // 2pads.bottom = pad_height - pads.toppadded_image = cv2.copyMakeBorder(resized_image, pads.top, pads.bottom, pads.left, pads.right, cv2.BORDER_CONSTANT, value=pad_color)return padded_image, pads, scale

letterbox函数用于调整图像大小，并在图像周围添加填充，使其符合目标尺寸。

4. 设置路径和参数

	images_path = "/media/lindsay/data/rockchip_algorithem/algo/zaozi/train/images"images_list = os.listdir(images_path)cali_num = len(images_list)jump_num = int(len(images_list) / cali_num)cali_list = images_list[0::jump_num][0:cali_num]input_w = 320input_h = input_wtarget_size = (input_w, input_h)pad_color = (0, 0, 0)save_txt = "datasettxt"save_path = "./calib"absolute_save_path = os.path.abspath(save_path)

我们定义了图像路径、目标尺寸和填充颜色，并确定了保存路径。

5. 清理或创建保存目录

	if os.path.exists(absolute_save_path):for filename in os.listdir(absolute_save_path):file_path = os.path.join(absolute_save_path, filename)try:if os.path.isfile(file_path) or os.path.islink(file_path):os.unlink(file_path)elif os.path.isdir(file_path):shutil.rmtree(file_path)except Exception as e:print(f'Failed to delete {file_path}. Reason: {e}')else:os.makedirs(absolute_save_path)

如果保存目录存在，则清空目录内容，否则创建目录。

6. 处理图像并保存

	count = 0with open(save_txt, "w") as f:for img in tqdm(cali_list):count += 1img_abs_path = os.path.join(images_path, img)img_suffix = os.path.splitext(img_abs_path)[-1]img_ndarry = cv2.imread(img_abs_path)processed_img, pads, scale = letterbox(img_ndarry, target_size, pad_color)cv2.imwrite(os.path.join(absolute_save_path, str(count) + img_suffix), processed_img)img_save_abs_path = os.path.join(absolute_save_path, str(count) + img_suffix)f.write(img_save_abs_path + "\n")

我们遍历图像列表，调整图像大小并填充，最后将处理后的图像保存到指定目录，并将路径写入文本文件。