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

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 shutil

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

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_w
    pad_height = target_size[0] - new_h

    pads = BOX_RECT()
    pads.left = pad_width // 2
    pads.right = pad_width - pads.left
    pads.top = pad_height // 2
    pads.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 += 1
        img_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 = 0
	        self.right = 0
	        self.top = 0
	        self.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_w
	    pad_height = target_size[0] - new_h
	    pads = BOX_RECT()
	    pads.left = pad_width // 2
	    pads.right = pad_width - pads.left
	    pads.top = pad_height // 2
	    pads.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

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 = 320
	input_h = input_w
	target_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 = 0
	with open(save_txt, "w") as f:
	    for img in tqdm(cali_list):
	        count += 1
	        img_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")

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