added train

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

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()

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()

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 <onnx_model> <TARGET_PLATFORM> <dtype(optional)> <output_rknn_path(optional)>
+
+## 比如
+python convert.py yolov5.onnx rk1808 u8 yolov5.rknn
+```
+
+## 3. 测试
+将模型部署到1808上之后，使用下面的命令测试一张图片：
+
+```shell
+python test.py <rknn_model_name> <test_image>
+```