6.2 deeplabv3+的rknn模型检测

       该部分的检测分为检测前的图片处理和检测两部分。

       首先，检测前的图片处理基本上套用了deeplabv3+源代码的图片处理方式，只是在检测前做了略微的改动.其次，rknn模型得到的结果需要稍微的处理，与pth模型得到的结果格式略有不同。

       具体的检测代码如下：

import os
import urllib
import traceback
import time
import sys
import numpy as np
import cv2
from rknn.api import RKNN
from PIL import Image
import copy
import torch
import torch.nn.functional as F

ONNX_MODEL = 'best.onnx'
RKNN_MODEL = 'best.rknn'
DATASET = './dataset.txt'

colors = [ (0, 0, 0), (128, 0, 0), (0, 128, 0), (128, 128, 0), (0, 0, 128), (128, 0, 128), (0, 128, 128), 
                            (128, 128, 128), (64, 0, 0), (192, 0, 0), (64, 128, 0), (192, 128, 0), (64, 0, 128), (192, 0, 128), 
                            (64, 128, 128), (192, 128, 128), (0, 64, 0), (128, 64, 0), (0, 192, 0), (128, 192, 0), (0, 64, 128), 
                            (128, 64, 12)]




#将图像转换成RGB图像，防止灰度图在预测时报错。
def cvtColor(image):
    if len(np.shape(image)) == 3 and np.shape(image)[2] == 3:
        return image 
    else:
        image = image.convert('RGB')
        return image 
        
#对输入图像进行resize
def resize_image(image, size):
    iw, ih  = image.size
    w, h    = size

    scale   = min(w/iw, h/ih)
    nw      = int(iw*scale)
    nh      = int(ih*scale)

    image   = image.resize((nw,nh), Image.BICUBIC)
    new_image = Image.new('RGB', size, (128,128,128))
    new_image.paste(image, ((w-nw)//2, (h-nh)//2))

    return new_image, nw, nh

def preprocess_input(image):
    image /= 255.0
    return image


def processes(outputs, nw, nh,orininal_w, orininal_h,old_img):
    input_shape=[512,512]
    for item in outputs:
    	
    	item=item.reshape((2,512,512))
    	item=torch.tensor(item)
    	pr = F.softmax(item.permute(1,2,0),dim = -1).cpu().numpy()
    	#--------------------------------------#
        #   将灰条部分截取掉
        #--------------------------------------#
    	pr=pr[int((input_shape[0] - nh) // 2) : int((input_shape[0] - nh) // 2 + nh),int((input_shape[1] - nw) // 2) : int((input_shape[1] - nw) // 2 + nw)]
    	#---------------------------------------------------#
        #   进行图片的resize
        #---------------------------------------------------#
    	pr = cv2.resize(pr, (orininal_w, orininal_h), interpolation = cv2.INTER_LINEAR)
    	#---------------------------------------------------#
        #   取出每一个像素点的种类
        #---------------------------------------------------#
    	pr = pr.argmax(axis=-1)
    	
    	
    	seg_img = np.reshape(np.array(colors, np.uint8)[np.reshape(pr, [-1])], [orininal_h, orininal_w, -1])
    	image   = Image.fromarray(np.uint8(seg_img))
    	#image   = Image.blend(old_img, image, 0.5)
    	return image
    	


if __name__ == '__main__':

    # Create RKNN object
    rknn = RKNN(verbose=True)

    # pre-process config
    print('--> config model')
    rknn.config(mean_values=[0, 0, 0], std_values=[255,255, 255], target_platform='rk3588')
    # Load model
    print('--> Loading model')
    ret = rknn.load_onnx(model=ONNX_MODEL)
    if ret != 0:
        print('Load model failed!')
        exit(ret)
    print('done')

    # Build model
    print('--> Building model')
    ret = rknn.build(do_quantization=True, dataset=DATASET)
    if ret != 0:
        print('Build model failed!')
        exit(ret)
    print('done')
    
    # Export RKNN model
    print('--> Export rknn model')
    ret = rknn.export_rknn(RKNN_MODEL)
    if ret != 0:
        print('Export rknn model failed!')
        exit(ret)
    print('done')

    # Init runtime environment
    print('--> Init runtime environment')
    ret = rknn.init_runtime()
    if ret != 0:
        print('Init runtime environment failed!')
        exit(ret)
    print('done')
    
    # Set inputs
    img = Image.open("ss.jpg")
    #代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
    image = cvtColor(img)
    # 将plt图片转换为array格式的cv2的图片，后续作为输出的处理
    new_imgs=np.array(image)
    #对输入图像进行一个备份，后面用于绘图
    old_img     = copy.deepcopy(image)
    orininal_h  = np.array(image).shape[0]
    orininal_w  = np.array(image).shape[1]
    # 给图像增加灰条，实现不失真的resize
    image_data, nw, nh  = resize_image(image, (512,512))
    #添加上batch_size维度
    image_data  = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, np.float32)), (0, 1, 2)), 0)#(2, 0, 1)
    
    outputs = rknn.inference(inputs=[image_data])
    r_image=processes(outputs, nw, nh,orininal_w, orininal_h,old_img)
    r_image.save("deep.jpg")
    
   

        注意，以上代码运行不会报错，但得到的结果不准确，后续需要对数据做进一步处理，处理后再更新。