retinanet_resnet模型

from __future__ import print_function
import cv2
import time
import logging as log
from openvino.inference_engine import IECore
labels_to_names = {0: 'person', 1: 'bicycle', 2: 'car', 3: 'motorcycle', 4: 'airplane', 5: 'bus', 6: 'train', 7: 'truck', 8: 'boat', 9: 'traffic light', 10: 'fire hydrant', 11: 'stop sign', 12: 'parking meter', 13: 'bench', 14: 'bird', 15: 'cat', 16: 'dog', 17: 'horse', 18: 'sheep', 19: 'cow', 20: 'elephant', 21: 'bear', 22: 'zebra', 23: 'giraffe', 24: 'backpack', 25: 'umbrella', 26: 'handbag', 27: 'tie', 28: 'suitcase', 29: 'frisbee', 30: 'skis', 31: 'snowboard', 32: 'sports ball', 33: 'kite', 34: 'baseball bat', 35: 'baseball glove', 36: 'skateboard', 37: 'surfboard', 38: 'tennis racket', 39: 'bottle', 40: 'wine glass', 41: 'cup', 42: 'fork', 43: 'knife', 44: 'spoon', 45: 'bowl', 46: 'banana', 47: 'apple', 48: 'sandwich', 49: 'orange', 50: 'broccoli', 51: 'carrot', 52: 'hot dog', 53: 'pizza', 54: 'donut', 55: 'cake', 56: 'chair', 57: 'couch', 58: 'potted plant', 59: 'bed', 60: 'dining table', 61: 'toilet', 62: 'tv', 63: 'laptop', 64: 'mouse', 65: 'remote', 66: 'keyboard', 67: 'cell phone', 68: 'microwave', 69: 'oven', 70: 'toaster', 71: 'sink', 72: 'refrigerator', 73: 'book', 74: 'clock', 75: 'vase', 76: 'scissors', 77: 'teddy bear', 78: 'hair drier', 79: 'toothbrush'}
def retinanet_detection():
    model_xml = "D:/projects/models/retinanet/retinanet_resnet50_coco_best_v2.1.0.xml"
    model_bin = "D:/projects/models/retinanet/retinanet_resnet50_coco_best_v2.1.0.bin"

    log.info("Creating Inference Engine")
    ie = IECore()
    # Read IR
    net = ie.read_network(model=model_xml, weights=model_bin)

    log.info("Preparing input blobs")
    input_it = iter(net.input_info)
    input_blob = next(input_it)
    print(input_blob)
    out_blob = next(iter(net.outputs))

    # Read and pre-process input images
    print(net.input_info[input_blob].input_data.shape)

    image = cv2.imread("D:/images/city-walk.png")
    ih, iw, ic = image.shape                                                                        #模型输入( 1, 1333, 1333, 3)
    image_blob = cv2.dnn.blobFromImage(image, 1.0, (1333, 1333), (103.939, 116.779, 123.68), False) #图像预处理

    # Loading model to the plugin
    exec_net = ie.load_network(network=net, device_name="CPU")

    # Start sync inference
    log.info("Starting inference in synchronous mode")
    inf_start1 = time.time()
    res = exec_net.infer(inputs={input_blob: [image_blob]})
    inf_end1 = time.time() - inf_start1
    print("inference time(ms) : %.3f" % (inf_end1 * 1000))

    # Processing output blob
    log.info("Processing output blob")
    res = res[out_blob]
    for obj in res[0][0]:                                                           #模型输出(1, 1, N, 7)
        if obj[2] > 0.5:
            index = int(obj[1])
            xmin = int(obj[3] * iw)
            ymin = int(obj[4] * ih)
            xmax = int(obj[5] * iw)
            ymax = int(obj[6] * ih)
            cv2.putText(image, labels_to_names[index], (xmin, ymin-10), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (255, 0, 0), 2)
            cv2.rectangle(image, (xmin, ymin), (xmax, ymax), (0, 0, 255), 2, 8, 0)
    cv2.imshow("retinanet detection", image)
    cv2.waitKey(0)
    cv2.destroyAllWindows()


if __name__ == '__main__':
    retinanet_detection()

 

posted @ 2022-02-25 16:56  wuyuan2011woaini  阅读(54)  评论(0编辑  收藏  举报