import cv2
import os
import os.path as osp

#location = [左上角(x,y)，右下角(x,y);截图的左上角(x,y),右下角(x,y)]
#information = [start_frame_o ,end_frame_o,idx_o,location_origin, start_frame_s,end_frame_s,idx_s,location_single]
def video_label(video_path,information,output_path):
    """
        给指定帧的图像，添加标注信息
        video_path:视频输入文件
        information:视频的信息包括原视频起始帧、索引以及对象坐标，截取后视频起始帧、索引以及对象坐标
        ouput_path：标注后视频保存地址

    """
    start_frame_o ,end_frame_o,idx_o,location_origin, start_frame_s,end_frame_s,idx_s,location_single = information

    cap = cv2.VideoCapture(video_path)
    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    codec = cv2.VideoWriter_fourcc(*'XVID')
    video_name = os.path.basename(video_path)
    """labels取值取决于information中的信息"""
    labels = idx_s
    out = cv2.VideoWriter(osp.join(output_path,str(labels),video_name), codec, fps, (frame_width, frame_height))

    if labels == 0:
        labels = "Nodding!"
    elif labels == 1:
        labels = "no playing phone!"
    elif labels == 2:
        labels = "not sleep!"
    elif labels == 3:
        labels = "playing phone!"
    elif labels == 4:
        labels = "sleep!"
    else:
        pass

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        frame_idx = cap.get(cv2.CAP_PROP_POS_FRAMES)
        if start_frame_o + start_frame_s <= frame_idx and frame_idx <= end_frame_o + end_frame_s:
            x0_o, y0_o, x1_o, y1_o = int(location_origin[0]),int(location_origin[1]),int(location_origin[2]),int(location_origin[3])
            x0_s, y0_s, x1_s, y1_s = int(location_single[0]),int(location_single[1]),int(location_single[2]),int(location_single[3])
            cv2.rectangle(frame, (x0_o + x0_s, y0_o + y0_s), (x1_s + x0_o, y1_s + y0_o), (0, 255, 0), 2)
            cv2.putText(frame, labels, (0, 40), cv2.FONT_HERSHEY_SIMPLEX, 3, (0, 255, 0), 2)
        out.write(frame)
        cv2.namedWindow("Detection_results", cv2.WINDOW_AUTOSIZE)

        cv2.imshow("Detection_results", frame)

        if cv2.waitKey(10) & 0xFF == ord('q'):
            break

    cap.release()
    out.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    video_path = "C:/Users/Administrator/PycharmProjects/pythonProject1/data/0711-7_5.avi"
    information = [0,201,0,[1140,722,1393,1197],0,10,0,[82,6,242,159]]
    output_path = "C:/Users/Administrator/PycharmProjects/pythonProject1/output"

    video_label(video_path,information,output_path)