onnx推理maskrcnn

￡神魔★判官ぃ 2022-08-28 00:37 184阅读 0赞

import onnxruntime as rt
    import cv2
    from torchvision.ops.boxes import nms
    def preprocess(img):
        norm_mean = (0.5, 0.5, 0.5)
        norm_std = (0.5, 0.5, 0.5)
        data = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        data = data.astype(np.float32) / 255.0
        data = (data - norm_mean) / norm_std
        data = np.transpose(data, (2, 0, 1)).astype(np.float32)
        data = np.expand_dims(data, axis=0)
        return data
    
    img = cv2.imread('/home/lixuan/workspace/project/segmentation/mmdetection/demo/demo.jpg')
    data = preprocess(img)
    sess = rt.InferenceSession('/home/lixuan/snap/mmdetection/testmaskrcnn.onnx/mask_rcnn_r50_fpn_1x_coco.onnx')
    input_name = sess.get_inputs()[0].name
    boxes = sess.run(['boxes'], {input_name:data.astype(np.float32)})[0]
    labels = sess.run(['labels'], {input_name:data.astype(np.float32)})[0]
    masks = sess.run(['masks'], {input_name:data.astype(np.float32)})[0]
    needindexs = np.where(boxes[:,-1] > 0.2)
    boxes = boxes[needindexs]
    labels = labels[needindexs]
    masks = masks[needindexs]
    boxs = []
    score = []
    for index,box in enumerate(boxes):
        x1 = int(box[0])
        y1 = int(box[1])
        x2 = int(box[2])
        y2 = int(box[3])
        score.append(float(box[4]))
        boxs.append([x1,y1,x2,y2])
    boxs = torch.Tensor(boxs)
    score = torch.Tensor(score)
    needindexs = nms(boxs,score,0.3)
    boxes = boxes[needindexs]
    labels = labels[needindexs]
    masks = masks[needindexs]
    for index,box in enumerate(boxes):
        x1 = int(box[0])
        y1 = int(box[1])
        x2 = int(box[2])
        y2 = int(box[3])
        w = x2 - x1
        h = y2 - y1
        mask = masks[index]
        mask = cv2.resize(mask,(w,h)) * 255
        ret, mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)
    
        black = np.zeros_like(img)
        black = cv2.cvtColor(black, cv2.COLOR_BGR2GRAY)
        black[y1:y2,x1:x2] = mask
    
        contours, hierarchy = cv2.findContours(black, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        if len(contours) == 1:
            contours = np.array(contours).reshape((-1, 1, 2))
            cv2.polylines(img, [contours], isClosed=True, color=(0, 0, 255), thickness=1)
        else:
            for contour in contours:
                contour = np.array(contour).reshape((-1, 1, 2))
                cv2.polylines(img, [contour], isClosed=True, color=(0, 0, 255), thickness=1)
        # cv2.imshow('',black)
        cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
    cv2.imshow('',img)
    cv2.waitKey(0)