import os
import numpy as np
import matplotlib.pyplot as plt
from IPython.display import display, clear_output
import tracker
from detector import Detector
import cv2

if __name__ == '__main__':
    px = 1/plt.rcParams['figure.dpi']

    font_draw_number = cv2.FONT_HERSHEY_SIMPLEX
    draw_text_postion = (int(660 * 1), int(440 * 1))

    # 初始化 yolov5
    detector = Detector()

    # 打开视频
#     capture = cv2.VideoCapture('./video/test.mp4')
#     capture = cv2.VideoCapture("./water_pics/ljk1/20211%09d.jpg")
    import glob
    capture = sorted(glob.glob('./water_pics/*/*.jpg'))
    
    # capture = cv2.VideoCapture('/mnt/datasets/datasets/towncentre/TownCentreXVID.avi')

#     while True:
    for im in capture:
        # 读取每帧图片
#         _, im = capture.read()
        im_path = im
        print(f"当前正在处理的图片是：{im_path}")
        im = cv2.imread(im)
        if im is None:
            break

        # 缩小尺寸，1920x1080->960x540
        im = cv2.resize(im, (960, 540))

        list_bboxs = []
        bboxes = detector.detect(im)
        result = {}
        text_draw = ''
        # 如果画面中 有bbox
        if len(bboxes) > 0:
            # re-id结果。
            list_bboxs = tracker.update(bboxes, im)

            # 画框
            # 撞线检测点，(x1，y1)，y方向偏移比例 0.0~1.0
            output_image_frame = tracker.draw_bboxes(im, list_bboxs, line_thickness=None)
        else:
            # 如果画面中 没有bbox
            output_image_frame = im

        # 输出图片
        # output_image_frame = cv2.add(output_image_frame, color_polygons_image)

        if len(list_bboxs) > 0:
            # ----------------------判断撞线----------------------
            for item_bbox in list_bboxs:
                x1, y1, x2, y2, label, track_id = item_bbox
                result[label] = result.get(label, 0) + 1
            for key, value in result.items():
               text_draw += f"{key}:{value}-"
            cv2.imwrite(f"./water_pic_result/result_{os.path.basename(im_path)}", output_image_frame)
            # pass
        else:
            # 如果图像中没有任何的bbox，则清空list
            # list_overlapping_blue_polygon.clear()
            # list_overlapping_yellow_polygon.clear()
            pass
        pass
        pass
    pass