openvino_yoloV3

import sys, os, cv2, time
import numpy as np, math
from argparse import ArgumentParser
from openvino.inference_engine import IECore

m_input_size = 416                  #模型输入416*416

yolo_scale_13 = 13                  #cell 13*13
yolo_scale_26 = 26                  #cell 26*26
yolo_scale_52 = 52                  #cell 52*52

classes = 2                         #类别个数
coords = 4                          #坐标占位 (中心点x,中心点y，高度，宽度)
num = 3                             #每个cell三个检测框
anchors = [10,13,16,30,33,23,30,61,62,45,59,119,116,90,156,198,373,326]     #先验框

LABELS = ("operator1", "operator2") #label列表

label_text_color = (255, 255, 255)
label_background_color = (125, 175, 75)
box_color = (255, 128, 0)
box_thickness = 1

def build_argparser():
    parser = ArgumentParser()
    parser.add_argument("-d", "--device", help="Specify the target device to infer on; CPU, GPU, FPGA or MYRIAD is acceptable. \
                                                Sample will look for a suitable plugin for device specified (CPU by default)", default="CPU", type=str)
    return parser

#获取展平后的数据的索引
def EntryIndex(side, lcoords, lclasses, location, entry):       
    n = int(location / (side * side))
    loc = location % (side * side)
    return int(n * side * side * (lcoords + lclasses + 1) + entry * side * side + loc)


class DetectionObject():
    xmin = 0
    ymin = 0
    xmax = 0
    ymax = 0
    class_id = 0
    confidence = 0.0

    def __init__(self, x, y, h, w, class_id, confidence, h_scale, w_scale):
        self.xmin = int((x - w / 2) * w_scale)
        self.ymin = int((y - h / 2) * h_scale)
        self.xmax = int(self.xmin + w * w_scale)
        self.ymax = int(self.ymin + h * h_scale)
        self.class_id = class_id
        self.confidence = confidence

# IOU交并比
def IntersectionOverUnion(box_1, box_2):
    width_of_overlap_area = min(box_1.xmax, box_2.xmax) - max(box_1.xmin, box_2.xmin)
    height_of_overlap_area = min(box_1.ymax, box_2.ymax) - max(box_1.ymin, box_2.ymin)
    area_of_overlap = 0.0
    if (width_of_overlap_area < 0.0 or height_of_overlap_area < 0.0):
        area_of_overlap = 0.0
    else:
        area_of_overlap = width_of_overlap_area * height_of_overlap_area
    box_1_area = (box_1.ymax - box_1.ymin)  * (box_1.xmax - box_1.xmin)
    box_2_area = (box_2.ymax - box_2.ymin)  * (box_2.xmax - box_2.xmin)
    area_of_union = box_1_area + box_2_area - area_of_overlap
    retval = 0.0
    if area_of_union <= 0.0:
        retval = 0.0
    else:
        retval = (area_of_overlap / area_of_union)
    return retval


def ParseYOLOV3Output(blob, resized_im_h, resized_im_w, original_im_h, original_im_w, threshold, objects):
    # 输出结果 例子V3
    # conv2d_58/BiasAdd （1, 13, 13, 255）  anchor(116,90, 156,198, 373,326)
    # conv2d_66/BiasAdd （1, 26, 26, 255）  anchor(30,61, 62,45, 59,119)
    # conv2d_74/BiasAdd （1, 52, 52, 255）  anchor(10,13, 16,30, 33,23)
    # 255(B, Cx, Cy, N*85)  N:num of cell   85:(x, y, h, w, box_score, class_no_1, …, class_no_80)
    # print(blob.shape)                                     #(B,N*(5+class),Cx,Cy)
    out_blob_h = blob.shape[2]                              #Cx
    out_blob_w = blob.shape[3]                              #Cy
    side = out_blob_h
    anchor_offset = 0

    if len(anchors) == 18:   ## YoloV3                      #根据anchor判断yolo类型
        if side == yolo_scale_13:
            anchor_offset = 2 * 6
        elif side == yolo_scale_26:
            anchor_offset = 2 * 3
        elif side == yolo_scale_52:
            anchor_offset = 2 * 0

    elif len(anchors) == 12: ## tiny-YoloV3
        if side == yolo_scale_13:
            anchor_offset = 2 * 3
        elif side == yolo_scale_26:
            anchor_offset = 2 * 0

    else:                    ## ???
        if side == yolo_scale_13:
            anchor_offset = 2 * 6
        elif side == yolo_scale_26:
            anchor_offset = 2 * 3
        elif side == yolo_scale_52:
            anchor_offset = 2 * 0

    side_square = side * side
    output_blob = blob.flatten()                    #展平（例：1*(5+2)*3*13*13）
    for i in range(side_square):
        row = int(i / side)                         #cell的x
        col = int(i % side)                         #cell的y
        for n in range(num):                        #3个预测框
            obj_index = EntryIndex(side, coords, classes, n * side * side + i, coords)      #置信度索引  box_score
            box_index = EntryIndex(side, coords, classes, n * side * side + i, 0)           #表示框x索引 x
            scale = output_blob[obj_index]                                                  #取得置信度
            if (scale < threshold):
                continue
            # 中心点坐标(x,y)
            x = (col + output_blob[box_index + 0 * side_square]) / side * resized_im_w      #根据原始坐标，得到绝对坐标
            y = (row + output_blob[box_index + 1 * side_square]) / side * resized_im_h      
            # 高度和宽度
            height = math.exp(output_blob[box_index + 3 * side_square]) * anchors[anchor_offset + 2 * n + 1]    #得到绝对坐标
            width = math.exp(output_blob[box_index + 2 * side_square]) * anchors[anchor_offset + 2 * n]
            for j in range(classes):
                class_index = EntryIndex(side, coords, classes, n * side_square + i, coords + 1 + j)            #类索引
                prob = scale * output_blob[class_index]                                                         #类置信度
                if prob < threshold:
                    continue
                # 将矩形框及置信度、类别进行存储
                obj = DetectionObject(x, y, height, width, j, prob, (original_im_h / resized_im_h), (original_im_w / resized_im_w))
                objects.append(obj)
    return objects


def main_IE_infer():
    camera_width = 320
    camera_height = 240
    fps = ""
    framepos = 0
    frame_count = 0
    vidfps = 0
    skip_frame = 0
    elapsedTime = 0
    new_w = int(camera_width * m_input_size/camera_width)
    new_h = int(camera_height * m_input_size/camera_height)
    args = build_argparser().parse_args()
    model_xml = "/home/bhc/darknet-master/backup1/16/16.xml" #<--- CPU
    #model_xml = "lrmodels/YoloV3/FP16/frozen_yolo_v3.xml" #<--- MYRIAD
    model_bin = os.path.splitext(model_xml)[0] + ".bin"

    # cap = cv2.VideoCapture(0)
    # cap.set(cv2.CAP_PROP_FPS, 30)
    # cap.set(cv2.CAP_PROP_FRAME_WIDTH, camera_width)
    # cap.set(cv2.CAP_PROP_FRAME_HEIGHT, camera_height)

    cap = cv2.VideoCapture("/home/bhc/BHC/Q3/1_Astemo_DL/video/D16_operator/D16_20211025135235.mp4")
    camera_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))                   #视频的宽度
    camera_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))                 #视频的高度
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))                    #视频总帧数
    vidfps = int(cap.get(cv2.CAP_PROP_FPS))                                 #视频fps
    print("videosFrameCount =", str(frame_count))
    print("videosFPS =", str(vidfps))

    time.sleep(1)
    ie = IECore()
    net = ie.read_network(model=model_xml)                                  #读取模型
    input_blob = next(iter(net.input_info))                                 #模型输入
    exec_net = ie.load_network(network=net, device_name="CPU")              #加载模型

    while cap.isOpened():                                                   #视频是否打开
        t1 = time.time()

        ## Uncomment only when playing video files
        #cap.set(cv2.CAP_PROP_POS_FRAMES, framepos)

        ret, image = cap.read()                                             #读取帧
        if not ret:
            break
        resized_image = cv2.resize(image, (new_w, new_h), interpolation = cv2.INTER_CUBIC)  #resize操作到输入要求
        canvas = np.full((m_input_size, m_input_size, 3), 128)
        canvas[(m_input_size-new_h)//2:(m_input_size-new_h)//2 + new_h,(m_input_size-new_w)//2:(m_input_size-new_w)//2 + new_w,  :] = resized_image
        prepimg = canvas
        prepimg = prepimg[np.newaxis, :, :, :]                              # Batch size axis add
        prepimg = prepimg.transpose((0, 3, 1, 2))                           # NHWC to NCHW
        outputs = exec_net.infer(inputs={input_blob: prepimg})              # 推理BHWC（1,416,416,3） 3:RGB
        objects = []
        for output in outputs.values():
            objects = ParseYOLOV3Output(output, new_h, new_w, camera_height, camera_width, 0.7, objects) #解析推理结果

        # Filtering overlapping boxes
        objlen = len(objects)
        for i in range(objlen):
            if (objects[i].confidence == 0.0):
                continue
            for j in range(i + 1, objlen):
                if (IntersectionOverUnion(objects[i], objects[j]) >= 0.4):      #如果IOU大于0.4的则将置信度设置为0
                    objects[j].confidence = 0

        
        # Drawing boxes
        for obj in objects:
            if obj.confidence < 0.2:
                continue
            label = obj.class_id
            confidence = obj.confidence
            if confidence > 0.2:
                label_text = LABELS[label] + " (" + "{:.1f}".format(confidence * 100) + "%)"
                cv2.rectangle(image, (obj.xmin, obj.ymin), (obj.xmax, obj.ymax), box_color, box_thickness)
                cv2.putText(image, label_text, (obj.xmin, obj.ymin - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, label_text_color, 1)

        cv2.putText(image, fps, (camera_width - 170, 15), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (38, 0, 255), 1, cv2.LINE_AA)
        cv2.imshow("Result", image)

        if cv2.waitKey(1)&0xFF == ord('q'):
            break
        elapsedTime = time.time() - t1
        fps = "(Playback) {:.1f} FPS".format(1/elapsedTime)

        ## frame skip, video file only
        #skip_frame = int((vidfps - int(1/elapsedTime)) / int(1/elapsedTime))
        #framepos += skip_frame

    cv2.destroyAllWindows()
    del net
    del exec_net

if __name__ == '__main__':
    sys.exit(main_IE_infer() or 0)