#!/usr/bin/env python3

import cv2
import numpy as np
import openvino as ov

model_path = '../Models/yolov8n_openvino_model/yolov8n.xml'
image_path = '../../../assets/bus.jpg'
device_name = 'CPU'


def main():
    # Загрузка OpenVINO
    core = ov.Core()
    model = core.read_model(model_path)

    # Загрузка изображения
    image = cv2.imread(image_path)

    # Добавить N измерений
    input_tensor = np.expand_dims(image, 0)

    # Изменение формы модели в соответствии с ВxШ изображения
    n, h, w, c = input_tensor.shape
    model.reshape({model.input().get_any_name(): ov.PartialShape((n, c, h, w))})

    # Предварительная обработка
    ppp = ov.preprocess.PrePostProcessor(model)
    ppp.input().tensor().set_element_type(ov.Type.u8).set_layout(ov.Layout('NHWC'))
    ppp.input().model().set_layout(ov.Layout('NCHW'))
    ppp.output().tensor().set_element_type(ov.Type.f32)
    model = ppp.build()

    compiled_model = core.compile_model(model, device_name)
    results = compiled_model.infer_new_request({0: input_tensor})

    # Output
    predictions = next(iter(results.values()))
    detections = predictions.reshape(-1, 7)

    for detection in detections:
        confidence = detection[2]

        if confidence > 0.25:
            class_id = int(detection[1])

            xmin = int(detection[3] * w)
            ymin = int(detection[4] * h)
            xmax = int(detection[5] * w)
            ymax = int(detection[6] * h)

            cv2.rectangle(image, (xmin, ymin), (xmax, ymax), (0, 255, 0), 2)

    cv2.imwrite('/tmp/py_openvino_result.bmp', image)


if __name__ == '__main__':
    main()