From 4f26f54017f57ac60237c2de1b34328f1e638029 Mon Sep 17 00:00:00 2001 From: asxalex Date: Fri, 22 Aug 2025 14:32:50 +0800 Subject: [PATCH] added train --- train/anchors_yolov5.txt | 18 +++ train/convert.py | 75 ++++++++++++ train/test.py | 242 +++++++++++++++++++++++++++++++++++++++ train/train.md | 32 ++++++ 4 files changed, 367 insertions(+) create mode 100644 train/anchors_yolov5.txt create mode 100644 train/convert.py create mode 100644 train/test.py create mode 100644 train/train.md diff --git a/train/anchors_yolov5.txt b/train/anchors_yolov5.txt new file mode 100644 index 0000000..caba463 --- /dev/null +++ b/train/anchors_yolov5.txt @@ -0,0 +1,18 @@ +10.0 +13.0 +16.0 +30.0 +33.0 +23.0 +30.0 +61.0 +62.0 +45.0 +59.0 +119.0 +116.0 +90.0 +156.0 +198.0 +373.0 +326.0 \ No newline at end of file diff --git a/train/convert.py b/train/convert.py new file mode 100644 index 0000000..2d1dcf2 --- /dev/null +++ b/train/convert.py @@ -0,0 +1,75 @@ +import sys + +from rknn.api import RKNN + +DATASET_PATH = '../../../datasets/COCO/coco_subset_20.txt' +DEFAULT_RKNN_PATH = '../model/yolov5.rknn' +DEFAULT_QUANT = True + +def parse_arg(): + if len(sys.argv) < 3: + print("Usage: python3 {} onnx_model_path [platform] [dtype(optional)] [output_rknn_path(optional)]".format(sys.argv[0])) + print(" platform choose from [rk3562, rk3566, rk3568, rk3576, rk3588, rv1103, rv1106, rv1126b, rv1109, rv1126, rk1808]") + print(" dtype choose from [i8, fp] for [rk3562, rk3566, rk3568, rk3576, rk3588, rv1103, rv1106, rv1126b]") + print(" dtype choose from [u8, fp] for [rv1109, rv1126, rk1808]") + exit(1) + + model_path = sys.argv[1] + platform = sys.argv[2] + + do_quant = DEFAULT_QUANT + if len(sys.argv) > 3: + model_type = sys.argv[3] + if model_type not in ['i8', 'u8', 'fp']: + print("ERROR: Invalid model type: {}".format(model_type)) + exit(1) + elif model_type in ['i8', 'u8']: + do_quant = True + else: + do_quant = False + + if len(sys.argv) > 4: + output_path = sys.argv[4] + else: + output_path = DEFAULT_RKNN_PATH + + return model_path, platform, do_quant, output_path + +if __name__ == '__main__': + model_path, platform, do_quant, output_path = parse_arg() + + # Create RKNN object + rknn = RKNN(verbose=False) + + # Pre-process config + print('--> Config model') + rknn.config(mean_values=[[0, 0, 0]], std_values=[ + [255, 255, 255]], target_platform=platform) + print('done') + + # Load model + print('--> Loading model') + ret = rknn.load_onnx(model=model_path) + if ret != 0: + print('Load model failed!') + exit(ret) + print('done') + + # Build model + print('--> Building model') + ret = rknn.build(do_quantization=do_quant, dataset=DATASET_PATH) + if ret != 0: + print('Build model failed!') + exit(ret) + print('done') + + # Export rknn model + print('--> Export rknn model') + ret = rknn.export_rknn(output_path) + if ret != 0: + print('Export rknn model failed!') + exit(ret) + print('done') + + # Release + rknn.release() \ No newline at end of file diff --git a/train/test.py b/train/test.py new file mode 100644 index 0000000..6e8388d --- /dev/null +++ b/train/test.py @@ -0,0 +1,242 @@ +import cv2 +import os +import numpy as np +from rknnlite.api import RKNNLite +import sys +from copy import copy + +OBJ_THRESH = 0.25 +NMS_THRESH = 0.45 + +# The follew two param is for map test +# OBJ_THRESH = 0.001 +# NMS_THRESH = 0.65 + +IMG_SIZE = (640, 640) # (width, height), such as (1280, 736) + +CLASSES = ("person", "bicycle", "car","motorbike ","aeroplane ","bus ","train","truck ","boat","traffic light", + "fire hydrant","stop sign ","parking meter","bench","bird","cat","dog ","horse ","sheep","cow","elephant", + "bear","zebra ","giraffe","backpack","umbrella","handbag","tie","suitcase","frisbee","skis","snowboard","sports ball","kite", + "baseball bat","baseball glove","skateboard","surfboard","tennis racket","bottle","wine glass","cup","fork","knife ", + "spoon","bowl","banana","apple","sandwich","orange","broccoli","carrot","hot dog","pizza ","donut","cake","chair","sofa", + "pottedplant","bed","diningtable","toilet ","tvmonitor","laptop ","mouse ","remote ","keyboard ","cell phone","microwave ", + "oven ","toaster","sink","refrigerator ","book","clock","vase","scissors ","teddy bear ","hair drier", "toothbrush ") + +coco_id_list = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 27, 28, 31, 32, 33, 34, + 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, + 64, 65, 67, 70, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90] + + +def get_real_box(self, box, in_format='xyxy'): + bbox = copy(box) + if self.enable_ltter_box == True: + # unletter_box result + if in_format=='xyxy': + bbox[:,0] -= self.letter_box_info_list[-1].dw + bbox[:,0] /= self.letter_box_info_list[-1].w_ratio + bbox[:,0] = np.clip(bbox[:,0], 0, self.letter_box_info_list[-1].origin_shape[1]) + + bbox[:,1] -= self.letter_box_info_list[-1].dh + bbox[:,1] /= self.letter_box_info_list[-1].h_ratio + bbox[:,1] = np.clip(bbox[:,1], 0, self.letter_box_info_list[-1].origin_shape[0]) + + bbox[:,2] -= self.letter_box_info_list[-1].dw + bbox[:,2] /= self.letter_box_info_list[-1].w_ratio + bbox[:,2] = np.clip(bbox[:,2], 0, self.letter_box_info_list[-1].origin_shape[1]) + + bbox[:,3] -= self.letter_box_info_list[-1].dh + bbox[:,3] /= self.letter_box_info_list[-1].h_ratio + bbox[:,3] = np.clip(bbox[:,3], 0, self.letter_box_info_list[-1].origin_shape[0]) + return bbox + + +def filter_boxes(boxes, box_confidences, box_class_probs): + """Filter boxes with object threshold. + """ + box_confidences = box_confidences.reshape(-1) + class_max_score = np.max(box_class_probs, axis=-1) + classes = np.argmax(box_class_probs, axis=-1) + + _class_pos = np.where(class_max_score* box_confidences >= OBJ_THRESH) + scores = (class_max_score* box_confidences)[_class_pos] + + boxes = boxes[_class_pos] + classes = classes[_class_pos] + + return boxes, classes, scores + +def nms_boxes(boxes, scores): + """Suppress non-maximal boxes. + # Returns + keep: ndarray, index of effective boxes. + """ + x = boxes[:, 0] + y = boxes[:, 1] + w = boxes[:, 2] - boxes[:, 0] + h = boxes[:, 3] - boxes[:, 1] + + areas = w * h + order = scores.argsort()[::-1] + + keep = [] + while order.size > 0: + i = order[0] + keep.append(i) + + xx1 = np.maximum(x[i], x[order[1:]]) + yy1 = np.maximum(y[i], y[order[1:]]) + xx2 = np.minimum(x[i] + w[i], x[order[1:]] + w[order[1:]]) + yy2 = np.minimum(y[i] + h[i], y[order[1:]] + h[order[1:]]) + + w1 = np.maximum(0.0, xx2 - xx1 + 0.00001) + h1 = np.maximum(0.0, yy2 - yy1 + 0.00001) + inter = w1 * h1 + + ovr = inter / (areas[i] + areas[order[1:]] - inter) + inds = np.where(ovr <= NMS_THRESH)[0] + order = order[inds + 1] + keep = np.array(keep) + return keep + + +def box_process(position, anchors): + grid_h, grid_w = position.shape[2:4] + col, row = np.meshgrid(np.arange(0, grid_w), np.arange(0, grid_h)) + col = col.reshape(1, 1, grid_h, grid_w) + row = row.reshape(1, 1, grid_h, grid_w) + grid = np.concatenate((col, row), axis=1) + stride = np.array([IMG_SIZE[1]//grid_h, IMG_SIZE[0]//grid_w]).reshape(1,2,1,1) + + col = col.repeat(len(anchors), axis=0) + row = row.repeat(len(anchors), axis=0) + anchors = np.array(anchors) + anchors = anchors.reshape(*anchors.shape, 1, 1) + + box_xy = position[:,:2,:,:]*2 - 0.5 + box_wh = pow(position[:,2:4,:,:]*2, 2) * anchors + + box_xy += grid + box_xy *= stride + box = np.concatenate((box_xy, box_wh), axis=1) + + # Convert [c_x, c_y, w, h] to [x1, y1, x2, y2] + xyxy = np.copy(box) + xyxy[:, 0, :, :] = box[:, 0, :, :] - box[:, 2, :, :]/ 2 # top left x + xyxy[:, 1, :, :] = box[:, 1, :, :] - box[:, 3, :, :]/ 2 # top left y + xyxy[:, 2, :, :] = box[:, 0, :, :] + box[:, 2, :, :]/ 2 # bottom right x + xyxy[:, 3, :, :] = box[:, 1, :, :] + box[:, 3, :, :]/ 2 # bottom right y + + return xyxy + +def post_process(input_data, anchors): + boxes, scores, classes_conf = [], [], [] + # 1*255*h*w -> 3*85*h*w + input_data = [_in.reshape([len(anchors[0]),-1]+list(_in.shape[-2:])) for _in in input_data] + for i in range(len(input_data)): + boxes.append(box_process(input_data[i][:,:4,:,:], anchors[i])) + scores.append(input_data[i][:,4:5,:,:]) + classes_conf.append(input_data[i][:,5:,:,:]) + + def sp_flatten(_in): + ch = _in.shape[1] + _in = _in.transpose(0,2,3,1) + return _in.reshape(-1, ch) + + boxes = [sp_flatten(_v) for _v in boxes] + classes_conf = [sp_flatten(_v) for _v in classes_conf] + scores = [sp_flatten(_v) for _v in scores] + + boxes = np.concatenate(boxes) + classes_conf = np.concatenate(classes_conf) + scores = np.concatenate(scores) + + # filter according to threshold + boxes, classes, scores = filter_boxes(boxes, scores, classes_conf) + + # nms + nboxes, nclasses, nscores = [], [], [] + + for c in set(classes): + inds = np.where(classes == c) + b = boxes[inds] + c = classes[inds] + s = scores[inds] + keep = nms_boxes(b, s) + + if len(keep) != 0: + nboxes.append(b[keep]) + nclasses.append(c[keep]) + nscores.append(s[keep]) + + if not nclasses and not nscores: + return None, None, None + + boxes = np.concatenate(nboxes) + classes = np.concatenate(nclasses) + scores = np.concatenate(nscores) + + return boxes, classes, scores + + +def draw(image, boxes, scores, classes): + for box, score, cl in zip(boxes, scores, classes): + top, left, right, bottom = [int(_b) for _b in box] + print("%s @ (%d %d %d %d) %.3f" % (CLASSES[cl], top, left, right, bottom, score)) + cv2.rectangle(image, (top, left), (right, bottom), (255, 0, 0), 2) + cv2.putText(image, '{0} {1:.2f}'.format(CLASSES[cl], score), + (top, left - 6), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2) + + + +INPUT_SIZE = 640 + +if __name__ == '__main__': + rknn_lite = RKNNLite() + model = sys.argv[1] + image_name = sys.argv[2] + + target = None + + # load RKNN model + print('--> Load RKNN model') + ret = rknn_lite.load_rknn(model) + if ret != 0: + print('Load RKNN model failed') + exit(ret) + print('done') + + ori_img = cv2.imread(image_name) + img = cv2.cvtColor(ori_img, cv2.COLOR_BGR2RGB) + img = cv2.resize(img, (IMG_SIZE,IMG_SIZE)) + + # init runtime environment + print('--> Init runtime environment') + ret = rknn_lite.init_runtime(target=target) + if ret != 0: + print('Init runtime environment failed') + exit(ret) + print('done') + + + with open("anchors_yolov5.txt", 'r') as f: + values = [float(_v) for _v in f.readlines()] + anchors = np.array(values).reshape(3,-1,2).tolist() + print("use anchors from '{}', which is {}".format("anchors_yolov5.txt", anchors)) + + # Inference + print('--> Running model') + outputs = rknn_lite.inference(inputs=[img]) + + boxes, classes, scores = post_process(outputs, anchors) + + img_p = ori_img.copy() + if boxes is not None: + draw(img_p, get_real_box(boxes), scores, classes) + + if not os.path.exists('./result'): + os.mkdir('./result') + result_path = os.path.join('./result', image_name) + cv2.imwrite(result_path, img_p) + print('Detection result save to {}'.format(result_path)) + + rknn_lite.release() diff --git a/train/train.md b/train/train.md new file mode 100644 index 0000000..d853d29 --- /dev/null +++ b/train/train.md @@ -0,0 +1,32 @@ +# yolov5 train +这里介绍rknn yolov5训练和转换。 + +## 1. 训练 +源码使用[https://github.com/airockchip/yolov5](https://github.com/airockchip/yolov5)这个仓库。 + +训练过程与官方的yolov5训练方法一致，训练完成之后，使用该仓库的`export.py`进行转换： + +```shell +# for detection model +python export.py --rknpu --weight yolov5s.pt + +# for segmentation model +python export.py --rknpu --weight yolov5s-seg.pt +``` + +## 2. 转换 +之后，使用这里的`convert.py`将onnx转换为`rknn`: + +```shell +python convert.py + +## 比如 +python convert.py yolov5.onnx rk1808 u8 yolov5.rknn +``` + +## 3. 测试 +将模型部署到1808上之后，使用下面的命令测试一张图片： + +```shell +python test.py +``` \ No newline at end of file