# import StreamManagerApi.py
from StreamManagerApi import *
import cv2
import json

if __name__ == '__main__':
    # init stream manager
    streamManagerApi = StreamManagerApi()
    ret = streamManagerApi.InitManager()
    if ret != 0:
        print("Failed to init Stream manager, ret=%s" % str(ret))
        exit()

    # create streams by pipeline config file
    with open("./Sample.pipeline", 'rb') as f:
        pipelineStr = f.read()
    ret = streamManagerApi.CreateMultipleStreams(pipelineStr)
    if ret != 0:
        print("Failed to create Stream, ret=%s" % str(ret))
        exit()

    # Construct the input of the stream
    dataInput = MxDataInput()
    with open("test.jpg", 'rb') as f:
        dataInput.data = f.read()

    # The following is how to set the dataInput.roiBoxs
    """
    roiVector = RoiBoxVector()
    roi = RoiBox()
    roi.x0 = 100
    roi.y0 = 100
    roi.x1 = 200
    roi.y1 = 200
    roiVector.push_back(roi)
    dataInput.roiBoxs = roiVector
    """

    # Inputs data to a specified stream based on streamName.
    streamName = b'classification'
    inPluginId = 0
    uniqueId = streamManagerApi.SendDataWithUniqueId(streamName, inPluginId, dataInput)
    if uniqueId < 0:
        print("Failed to send data to stream.")
        exit()

    # Obtain the inference result by specifying streamName and uniqueId.
    inferResult = streamManagerApi.GetResultWithUniqueId(streamName, uniqueId, 3000)
    if inferResult.errorCode != 0:
        print("GetResultWithUniqueId error. errorCode=%d, errorMsg=%s" % (
            inferResult.errorCode, inferResult.data.decode()))
        exit()

    # print the infer result
    #infer_result = inferResult.data.decode()
    infer_result = json.loads(inferResult.data.decode())
    print(infer_result)
    print(type(infer_result))

    img2 = cv2.imread('test.jpg')
    for i, _ in enumerate(infer_result['MxpiObject']):
        y0 = infer_result['MxpiObject'][i]['y0']
        x0 = infer_result['MxpiObject'][i]['x0']
        y1 = infer_result['MxpiObject'][i]['y1']
        x1 = infer_result['MxpiObject'][i]['x1']

        # Draw detection bounding box
        bboxes = []
        bboxes = {'x0': int(x0),
                'x1': int(x1),
                'y0': int(y0),
                'y1': int(y1),
                'confidence': round(infer_result['MxpiObject'][i]['classVec'][0]['confidence'], 2),
                'text': infer_result['MxpiObject'][i]['classVec'][0]['className']}

        text = "{}{}".format(str(bboxes['confidence']), " ")
        for item in bboxes['text']:
            text += item
        cv2.rectangle(img2, (bboxes['x0'], bboxes['y0']), (bboxes['x1'], bboxes['y1']), (255, 0, 0), 2)
        cv2.putText(img2, text + str(i + 1), (bboxes['x0'], bboxes['y0'] + 15), cv2.FONT_HERSHEY_SIMPLEX, 0.5,
                    (0, 0, 255), 1)

    cv2.imwrite("resultout.jpg", img2)

    # destroy streams
    streamManagerApi.DestroyAllStreams()