import cv2
import json
import codecs
import time, datetime
import numpy as np
import ipywidgets
from IPython.display import display
from PIL import Image, ImageDraw, ImageFont
from timeit import default_timer as timer

# 定义视频读取函数
def read_video(input_video_path, video_start_time, video_end_time):
    cap = cv2.VideoCapture(input_video_path)  # 打开视频
    total_frame_num = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))  # 获取视频总帧数
    fps = cap.get(cv2.CAP_PROP_FPS)  # 视频帧率
    s_time_split = video_start_time.split(':')
    start_time = int(s_time_split[0]) * 3600 + int(s_time_split[1]) * 60 + int(s_time_split[2])
    e_time_split = video_end_time.split(':')
    end_time = int(e_time_split[0]) * 3600 + int(e_time_split[1]) * 60 + int(e_time_split[2])
    start_frame_id = int(start_time * fps)  # 设置需要处理的开始帧
    end_frame_id = int(end_time * fps)  # 设置需要处理的结束帧
    if end_frame_id > total_frame_num:
        end_frame_id = total_frame_num - 1
    cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame_id)  # 设置开始帧
    return cap, start_frame_id, end_frame_id


video_path = '/home/aistudio/work/faster_rcnn/homework/example2.mp4'
output_path = '/home/aistudio/work/faster_rcnn/homework/example2_1.mp4'

# start time
video_start_time = '00:00:00'
# end time
video_end_time = '00:00:20'

vid, start_frame_id, end_frame_id = read_video(video_path, video_start_time, video_end_time)
if not vid.isOpened():
    raise IOError("Couldn't open webcam or video")
video_FourCC    = int(vid.get(cv2.CAP_PROP_FOURCC))
video_fps       = vid.get(cv2.CAP_PROP_FPS)
video_size      = (int(vid.get(cv2.CAP_PROP_FRAME_WIDTH)), int(vid.get(cv2.CAP_PROP_FRAME_HEIGHT)))
out = cv2.VideoWriter(output_path, video_FourCC, video_fps, video_size)
accum_time = 0
curr_fps = 0
fps = "FPS: ??"
prev_time = timer()
index = 0
for frame_id in range(start_frame_id, end_frame_id):
    return_value, frame = vid.read()
    image = Image.fromarray(frame)
    result = np.asarray(image)
    curr_time = timer()
    exec_time = curr_time - prev_time
    prev_time = curr_time
    accum_time = accum_time + exec_time
    curr_fps = curr_fps + 1
    if accum_time > 1:
        accum_time = accum_time - 1
        fps = "FPS: " + str(curr_fps)
        curr_fps = 0
    cv2.putText(result, text=fps, org=(3, 15), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.50, color=(255, 0, 0), thickness=2)
    out.write(result)
    index += 1
print(index)
out.release()