import os
import cv2
import numpy as np
from caffe_object_detection import caffe_get_detector, caffe_detect_obj


def detect_and_highlight_object_in_img(img, language='en'):
    for obj_name, confidence, col_min, col_max, row_min, row_max in caffe_detect_obj(detector=detector, image=img, language=language):
        text_to_show = f"{obj_name} {confidence}"
        img = cv2.rectangle(img, (col_min, row_min), (col_max, row_max), (255, 0, 0), 2)
        # 获取文字区域框大小
        t_size = cv2.getTextSize(text_to_show, cv2.FONT_HERSHEY_TRIPLEX, 0.5, 1)[0]
        # 获取 文字区域右下角坐标
        textlbottom = (col_min, row_min+10) + np.array(list(t_size))
        # 绘制文字区域矩形框
        cv2.rectangle(img, (col_min, row_min), tuple(textlbottom), (255, 0, 0), -1)
        # 计算文字起始位置偏移
        row_min = int(row_min + (t_size[1]/2 + 12))
        # 绘字
        cv2.putText(img, text_to_show, (col_min, row_min), cv2.FONT_HERSHEY_TRIPLEX, 0.5, (255, 255, 255), 1)
    return img


def capture_and_detect(language='en'):
    capture = cv2.VideoCapture(0)
    while True:
        ret, frame = capture.read()
        img_to_show = detect_and_highlight_object_in_img(frame, language)
        cv2.imshow('streaming', img_to_show)
        #cv2.waitKey(1)
        if cv2.waitKey(1) & 0xff == ord('q'):
            break
    capture.release()
    cv2.destroyAllWindows()


def detect_and_show(img_filepath: str, language='en'):
    img = cv2.imread(img_filepath)
    img_to_show = detect_and_highlight_object_in_img(img, language)
    cv2.imshow('Pic', img_to_show)
    cv2.waitKey(0)
    cv2.destroyAllWindows()


if __name__ == "__main__":
    detector = caffe_get_detector(
        os.path.join(os.getcwd(), 'static', 'models', 'MobileNetSSD', 'MobileNetSSD_deploy.prototxt'),
        os.path.join(os.getcwd(), 'static', 'models', 'MobileNetSSD', 'MobileNetSSD_deploy.caffemodel')
    )

    detect_and_show('riding_horse.jpg')
    # capture_and_detect()