from AIDetector_pytorch import Detector
import imutils
import cv2
import os
os.environ['KMP_DUPLICATE_LIB_OK']='TRUE'
file = open('track_result.txt', 'w')

def main():

    name = 'demo'

    det = Detector()
    # cap = cv2.VideoCapture('D:/TH/5_smoke/Yolov5-Deepsort-main/smogfire_video14.mp4')
    cap = cv2.VideoCapture('/home/thsw/WJ/nyh/CODE/Yolov5-Deepsort-main/video/roadsign_orin.mp4')
    fps = int(cap.get(5))
    print('fps:', fps)
    t = int(1000/fps)

    videoWriter = None
    frame_count = 0

    while True:

        # try:
        _, im = cap.read()
        if im is None:
            break
        frame_count += 1
        file.write(str('第'+str(frame_count)+'帧\n*********************************************'))
        result = det.feedCap(im)
        result = result['frame']
        result = imutils.resize(result, height=500)
        if videoWriter is None:
            fourcc = cv2.VideoWriter_fourcc(
                'm', 'p', '4', 'v')  # opencv3.0
            videoWriter = cv2.VideoWriter(
                './detect_result/result_roadsign.mp4', fourcc, fps, (result.shape[1], result.shape[0]))

        videoWriter.write(result)
        # cv2.imshow(name, result)
        # cv2.waitKey(t)

        # if cv2.getWindowProperty(name, cv2.WND_PROP_AUTOSIZE) < 1:
        #     # 点x退出
        #     break
        # except Exception as e:
        #     print(e)
        #     break

    cap.release()
    videoWriter.release()
    # cv2.destroyAllWindows()

if __name__ == '__main__':
    
    main()