rknn部署rk3588

rknnlite_inference.py:

import os, cv2, time, numpy as np
from utils import *
from rknnlite.api import RKNNLite

conf_thres = 0.65
iou_thres = 0.95
input_width = 640
input_height = 640
model_name = 'green'
model_path = "./model"
config_path = "./config"
result_path = "./result"
image_path = "./dataset/0000.png"
video_path = "3.mp4"
video_inference = False
RKNN_MODEL = f'./model/green-640-640.rknn'
CLASSES = ['road', 'lane_line']


if __name__ == '__main__':
    isExist = os.path.exists(result_path)
    if not isExist:
        os.makedirs(result_path)
    rknn_lite = RKNNLite(verbose=False)
    ret = rknn_lite.load_rknn(RKNN_MODEL)
    ret = rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_AUTO)
    if video_inference == True:
        cap = cv2.VideoCapture(video_path)
        while(cap.isOpened()):
            ret, image_3c = cap.read()
            if not ret:
                break
            print('--> Running model for video inference')
            
            image_4c, image_3c = preprocess(image_3c, input_height, input_width)
            #ret = rknn_lite.init_runtime()
            start = time.time()
            image_3C = image_3c[np.newaxis,:]
            # print('111111',image_3C.shape)
            outputs = rknn_lite.inference(inputs=[image_3C])
            stop = time.time()
            fps = round(1/(stop-start), 2)
            
            outputs[0]=np.squeeze(outputs[0])
            
            outputs[0] = np.expand_dims(outputs[0], axis=0)
            
            colorlist = gen_color(len(CLASSES))
            results = postprocess(outputs, image_4c, image_3c, conf_thres, iou_thres, classes=len(CLASSES)) ##[box,mask,shape]
            
            results = results[0]              ## batch=1
            boxes, masks, shape = results
            if type(masks) != list and masks.ndim == 3:
                mask_img, vis_img = vis_result(image_3c, results, colorlist, CLASSES, result_path)
                #cv2.imshow("mask_img", mask_img)
                cv2.putText(vis_img, str(fps), (1,571),cv2.FONT_HERSHEY_SIMPLEX,1.0,(255,255,255))
                cv2.imshow("vis_img", vis_img)
            else:
                print("-------------No segmentation result-------------")
                #img5 = image_3c[185:455,:]
                #img2 = np.zeros_like(img5)
                #cv2.imshow("1", image_3c)
            cv2.waitKey(1)
    else:
        image_3c = cv2.imread(image_path) # (640,640,3)
        image_4c, image_3c = preprocess(image_3c, input_height, input_width)
        print('--> Running model for image inference')
        #ret = rknn_lite.init_runtime()
        start = time.time()
        image_3C2 = image_3c[np.newaxis,:] # (1, 640, 640, 3)
        outputs = rknn_lite.inference(inputs=[image_3C2]) # len(outputs)->2
        stop = time.time()
        fps = round(1/(stop-start), 2)  

        outputs[0]=np.squeeze(outputs[0])

        outputs[0] = np.expand_dims(outputs[0], axis=0)

        colorlist = [ (255,255,255), (0,0,0) ] #colorlist = gen_color(len(CLASSES))

        results = postprocess(outputs, image_4c, image_3c, conf_thres, iou_thres, classes=len(CLASSES)) ##[box,mask,shape]

        results = results[0]              ## batch=1

        boxes, masks, shape = results
        if masks.ndim == 2:
            masks = np.expand_dims(masks, axis=0).astype(np.float32)

        if type(masks) != list and masks.ndim == 3:
            mask_img, vis_img = vis_result(image_3c,  results, colorlist, CLASSES, result_path)
            print('--> Save inference result')
        else:
            print("-------------No segmentation result-------------")
    print("rknn_liteLite inference finish")
    rknn_lite.release()
    cv2.destroyAllWindows()

utils.py:

import cv2
import time
import numpy as np

def xywh2xyxy(x):
    y = np.copy(x)
    y[..., 0] = x[..., 0] - x[..., 2] / 2  # top left x
    y[..., 1] = x[..., 1] - x[..., 3] / 2  # top left y
    y[..., 2] = x[..., 0] + x[..., 2] / 2  # bottom right x
    y[..., 3] = x[..., 1] + x[..., 3] / 2  # bottom right y
    return y

def clip_boxes(boxes, shape):
    boxes[..., [0, 2]] = boxes[..., [0, 2]].clip(0, shape[1])  # x1, x2
    boxes[..., [1, 3]] = boxes[..., [1, 3]].clip(0, shape[0])  # y1, y2


def scale_boxes(img1_shape, boxes, img0_shape, ratio_pad=None):
    if ratio_pad is None:  # calculate from img0_shape
        gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1])  # gain  = old / new
        pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2  # wh padding
    else:
        gain = ratio_pad[0][0]
        pad = ratio_pad[1]

    boxes[..., [0, 2]] -= pad[0]  # x padding
    boxes[..., [1, 3]] -= pad[1]  # y padding
    boxes[..., :4] /= gain
    clip_boxes(boxes, img0_shape)
    return boxes


def crop_mask(masks, boxes):
    n, h, w = masks.shape
    x1, y1, x2, y2 = np.split(boxes[:, :, None], 4, axis=1)
    r = np.arange(w, dtype=np.float32)[None, None, :]  # rows shape(1,w,1)
    c = np.arange(h, dtype=np.float32)[None, :, None]  # cols shape(h,1,1)

    return masks * ((r >= x1) * (r < x2) * (c >= y1) * (c < y2))


def sigmoid(x):
    return 1.0 / (1 + np.exp(-x))


def process_mask(protos, masks_in, bboxes, shape):
    c, mh, mw = protos.shape  # CHW
   
    ih, iw = shape
    masks = sigmoid(masks_in @ protos.reshape(c, -1)).reshape(-1, mh, mw)  # CHW 【lulu】
    
    downsampled_bboxes = bboxes.copy()
    downsampled_bboxes[:, 0] *= mw / iw
    downsampled_bboxes[:, 2] *= mw / iw
    downsampled_bboxes[:, 3] *= mh / ih
    downsampled_bboxes[:, 1] *= mh / ih

    masks = crop_mask(masks, downsampled_bboxes)  # CHW
    masks = np.transpose(masks, [1, 2, 0])
    # masks = cv2.resize(masks, (shape[1], shape[0]), interpolation=cv2.INTER_NEAREST)
    masks = cv2.resize(masks, (shape[1], shape[0]), interpolation=cv2.INTER_LINEAR)
    
    if masks.ndim == 3:
        masks = np.transpose(masks, [2, 0, 1])
    return np.where(masks > 0.5, masks, 0)


def nms(bboxes, scores, threshold=0.5):
    x1 = bboxes[:, 0]
    y1 = bboxes[:, 1]
    x2 = bboxes[:, 2]
    y2 = bboxes[:, 3]
    areas = (x2 - x1) * (y2 - y1)

    order = scores.argsort()[::-1]
    keep = []
    
    while order.size > 0:
        i = order[0]
        #print('i:',i)
        keep.append(i)

        if order.size == 1: break
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])
        w = np.maximum(0.0, (xx2 - xx1))
        h = np.maximum(0.0, (yy2 - yy1))
        inter = w * h

        iou = inter / (areas[i] + areas[order[1:]] - inter)
        ids = np.where(iou <= threshold)[0]
        order = order[ids + 1]

    return keep


def non_max_suppression(
        prediction,
        conf_thres=0.25,
        iou_thres=0.45,
        classes=None,
        agnostic=False,
        multi_label=False,
        labels=(),
        max_det=300,
        nc=0,  # number of classes (optional)
):
    # Checks
    assert 0 <= conf_thres <= 1, f'Invalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0'
    assert 0 <= iou_thres <= 1, f'Invalid IoU {iou_thres}, valid values are between 0.0 and 1.0'

    # 【lulu】prediction.shape[1]：box + cls + num_masks
    bs = prediction.shape[0]  # batch size
    nc = nc or (prediction.shape[1] - 4)  # number of classes
    nm = prediction.shape[1] - nc - 4  # num_masks
    mi = 4 + nc  # mask start index
    xc = np.max(prediction[:, 4:mi], axis=1) > conf_thres  ## 【lulu】    

    # Settings
    # min_wh = 2  # (pixels) minimum box width and height
    max_wh = 7680  # (pixels) maximum box width and height
    max_nms = 30000  # maximum number of boxes into torchvision.ops.nms()
    time_limit = 0.5 + 0.05 * bs  # seconds to quit after
    redundant = True  # require redundant detections
    multi_label &= nc > 1  # multiple labels per box (adds 0.5ms/img)
    merge = False  # use merge-NMS

    t = time.time()
    output = [np.zeros((0, 6 + nm))] * bs  ## 【lulu】
    
    for xi, x in enumerate(prediction):  # image_3c index, image_3c inference
        # Apply constraints
        # x[((x[:, 2:4] < min_wh) | (x[:, 2:4] > max_wh)).any(1), 4] = 0  # width-height

        x = np.transpose(x, [1, 0])[xc[xi]]  ## 【lulu】#x.shape[0]----->10
        # If none remain process next image_3c
 
        if not x.shape[0]: continue

        # Detections matrix nx6 (xyxy, conf, cls)
        box, cls, mask = np.split(x, [4, 4 + nc], axis=1)  ## 【lulu】

        box = xywh2xyxy(box)  # center_x, center_y, width, height) to (x1, y1, x2, y2)
	
        j = np.argmax(cls, axis=1)  ## 【lulu】
        conf = cls[np.array(range(j.shape[0])), j].reshape(-1, 1)
        x = np.concatenate([box, conf, j.reshape(-1, 1), mask], axis=1)[conf.reshape(-1, ) > conf_thres]

        # Check shape
        n = x.shape[0]  # number of boxes

        if not n: continue
        x = x[np.argsort(x[:, 4])[::-1][:max_nms]]  # sort by confidence and remove excess boxes 【lulu】

        # Batched NMS
        c = x[:, 5:6] * max_wh  # classes
        boxes, scores = x[:, :4] + c, x[:, 4]  # boxes (offset by class), scores
        i = nms(boxes, scores, iou_thres)  ## 【lulu】
        i = i[:max_det]  # limit detections

        output[xi] = x[i]

        if (time.time() - t) > time_limit:
            # LOGGER.warning(f'WARNING ️ NMS time limit {time_limit:.3f}s exceeded')
            break  # time limit exceeded
    return output


def make_anchors(feats_shape, strides, grid_cell_offset=0.5):
    """Generate anchors from features."""
    anchor_points, stride_tensor = [], []
    assert feats_shape is not None
    dtype_ = np.float
    for i, stride in enumerate(strides):
        _, _, h, w = feats_shape[i]
        sx = np.arange(w, dtype=dtype_) + grid_cell_offset  # shift x
        sy = np.arange(h, dtype=dtype_) + grid_cell_offset  # shift y

        sy, sx = np.meshgrid(sy, sx, indexing='ij')
        anchor_points.append(np.stack((sx, sy), -1).reshape(-1, 2))
        stride_tensor.append(np.full((h * w, 1), stride, dtype=dtype_))
    return np.concatenate(anchor_points), np.concatenate(stride_tensor)


def dist2bbox(distance, anchor_points, xywh=True, dim=-1):
    """Transform distance(ltrb) to box(xywh or xyxy)."""
    lt, rb = np.split(distance, 2, dim)
    x1y1 = anchor_points - lt
    x2y2 = anchor_points + rb
    if xywh:
        c_xy = (x1y1 + x2y2) / 2
        wh = x2y2 - x1y1
        return np.concatenate((c_xy, wh), dim)  # xywh bbox
    return np.concatenate((x1y1, x2y2), dim)  # xyxy bbox


def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True, stride=32):
    # Resize and pad image while meeting stride-multiple constraints
    shape = im.shape[:2]  # current shape [height, width]
    if isinstance(new_shape, int):
        new_shape = (new_shape, new_shape)

    # Scale ratio (new / old)
    r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
    if not scaleup:  # only scale down, do not scale up (for better val mAP)
        r = min(r, 1.0)

    # Compute padding
    ratio = r, r  # width, height ratios
    new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
    dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding
    if auto:  # minimum rectangle
        dw, dh = np.mod(dw, stride), np.mod(dh, stride)  # wh padding
    elif scaleFill:  # stretch
        dw, dh = 0.0, 0.0
        new_unpad = (new_shape[1], new_shape[0])
        ratio = new_shape[1] / shape[1], new_shape[0] / shape[0]  # width, height ratios

    dw /= 2  # divide padding into 2 sides
    dh /= 2

    if shape[::-1] != new_unpad:  # resize
        im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
    im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add border
    return im, ratio, (dw, dh)


def preprocess(image, input_height, input_width):
    image_3c = image

    # Convert the image_3c color space from BGR to RGB
    image_3c = cv2.cvtColor(image_3c, cv2.COLOR_BGR2RGB)

    # Resize the image_3c to match the input shape
    image_3c, ratio, dwdh = letterbox(image_3c, new_shape=[input_height, input_width], auto=False)

    # Normalize the image_3c data by dividing it by 255.0
    image_4c = np.array(image_3c) / 255.0

    # Transpose the image_3c to have the channel dimension as the first dimension
    image_4c = np.transpose(image_4c, (2, 0, 1))  # Channel first

    # Expand the dimensions of the image_3c data to match the expected input shape
    image_4c = np.expand_dims(image_4c, axis=0).astype(np.float32)

    image_4c = np.ascontiguousarray(image_4c)  # contiguous

    # Return the preprocessed image_3c data
    return image_4c, image_3c


def postprocess(preds, img, orig_img, OBJ_THRESH, NMS_THRESH, classes=None):
    p = non_max_suppression(preds[0],
                            OBJ_THRESH,
                            NMS_THRESH,
                            agnostic=False,
                            max_det=300,
                            nc=classes,
                            classes=None)

    results = []
    proto = preds[1]

    for i, pred in enumerate(p):
        shape = orig_img.shape

        if not len(pred):
            results.append([[], [], []])  # save empty boxes
            continue
        
        masks = process_mask(proto[i], pred[:, 6:], pred[:, :4], img.shape[2:])  # HWC
        
        pred[:, :4] = scale_boxes(img.shape[2:], pred[:, :4], shape).round()
        results.append([pred[:, :6], masks, shape[:2]])
    return results


def gen_color(class_num):
    color_list = []
    np.random.seed(1)
    while 1:
        a = list(map(int, np.random.choice(range(255), 3)))
        if (np.sum(a) == 0): continue
        color_list.append(a)
        if len(color_list) == class_num: break
    return color_list


def vis_result(image_3c, results, colorlist, CLASSES, result_path):
    boxes, masks, shape = results

    #if masks.ndim == 2:
    #    masks = np.expand_dims(masks, axis=0).astype(np.float32)
    # Convert the image_3c color space from BGR to RGB
    image_3c = cv2.cvtColor(image_3c, cv2.COLOR_RGB2BGR)
    vis_img = image_3c.copy()
    mask_img = np.zeros_like(image_3c)
    cls_list = []
     
    center_list = []
    #print('boxes.shape,masks.shape',len(boxes),masks.shape)
    for box, mask in zip(boxes, masks):
        cls = int(box[-1])
        cls_list.append(cls)
        #print('box.shape,mask.shape,cls.shape',box.shape,mask.shape,cls)
        dummy_img = np.zeros_like(image_3c)
        dummy_img[mask != 0] = colorlist[int(box[-1])]  #(640, 640, 3) 
        mask_img[mask != 0] = colorlist[int(box[-1])]  #(640, 640, 3) 

        centroid = np.mean(np.argwhere(dummy_img), axis=0) #(x,y)

        if np.isnan(centroid).all() == False:
            centroid_x, centroid_y = int(centroid[1]), int(centroid[0])
            center_list.append([centroid_x, centroid_y])
    #print('center_list',center_list)
    
    vis_img = cv2.addWeighted(vis_img, 0.5, mask_img, 0.5, 0)
    #print('vis_img',vis_img.shape)
    '''
    for i, box in enumerate(boxes):
        cls = int(box[-1])
        cv2.rectangle(vis_img, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), (0, 0, 255), 3, 4)
        
        cv2.putText(vis_img, f"{CLASSES[cls]}:{round(box[4], 2)}", (int(box[0]), int(box[1])), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
    '''
                    
    for j in range(len(center_list)):
        cv2.circle(vis_img, (center_list[j][0], center_list[j][1]), radius=5, color=(0, 0, 255), thickness=-1)
    vis_img = np.concatenate([image_3c, mask_img, vis_img], axis=1)
    for i in range(len(CLASSES)):
        num = cls_list.count(i)
        if num != 0:
            print(f">>>>>>>>>>>>>Found {num} {CLASSES[i]}<<<<<<<<<<<<<<<<")
    cv2.imwrite(f"./{result_path}/origin_image.jpg", image_3c)
    cv2.imwrite(f"./{result_path}/mask_image.jpg", mask_img)
    cv2.imwrite(f"./{result_path}/visual_image.jpg", vis_img)
    return mask_img, vis_img