# -*- coding: utf-8 -*-
"""
Created on Wed Sep 20 14:10:09 2023

@author: ym
"""

import numpy as np
import os
import cv2

# from pathlib import Path
# import math
# import sys
# from scipy.spatial.distance import cdist


def video2imgs(videopath):
# =============================================================================
#     videopath：视频文件地址，在该地址的 "/file_imgs/" 文件加下存储视频帧图像
# =============================================================================

    path, filename = os.path.split(videopath)
    file, ext = os.path.splitext(filename)
    
    savepath = os.path.join(path, "{}_imgs".format(file))
    if not os.path.exists(savepath):
        os.makedirs(savepath)

    cap = cv2.VideoCapture(videopath)
    k = 0
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        
        k += 1
        cv2.imwrite(os.path.join(savepath, "{}.png".format(k)), frame)
        


def videosave(bboxes, videopath="100_1688009697927.mp4"):

    cap = cv2.VideoCapture(videopath)
    
    fps = int(cap.get(cv2.CAP_PROP_FPS))  # integer required, floats produce error in MP4 codec
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    
    ## =========================================== 在当前模块地址下存储图像和视频
    path = os.path.split(os.path.realpath(__file__))[0]
    _, filename = os.path.split(videopath)
    file, ext = os.path.splitext(filename)
    
    ## ======================================================== 视频保存设置
    fourcc = cv2.VideoWriter_fourcc(*'MP4V')
    save_video_path = os.path.join(path, "{}_show_1.mp4".format(file))    
    vid_writer = cv2.VideoWriter(save_video_path, fourcc, fps, (width, height))
    
    ## ======================================================== 图像保存路径设置
    save_img_path = os.path.join(path, "{}_show".format(file))
    if not os.path.exists(save_img_path):
        os.makedirs(save_img_path)
    
    cout = 0
    while cap.isOpened():
        ret, frame = cap.read()
        if ret:
            
            idx = np.where(bboxes[:, 7] == cout)[0]
            box = bboxes[idx, 0:4].astype(int)
            for i in range(box.shape[0]):
                x1, y1 = box[i, :2]
                x2, y2 = box[i, 2:4]
                cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 125, 255), 2)
                
                
            cv2.imwrite(os.path.join(save_img_path, "{}.png".format(cout)), frame)
            vid_writer.write(frame)
            
            cout += 1
        else:
            print("end!!!!!!!!!!!!!!!!!!!")
            break
    
    vid_writer.release()
    cap.release()