# -*- 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

VideoFormat = ['.mp4', '.avi', '.ts']
def video2imgs(videof, imgdir):
    cap = cv2.VideoCapture(videof)
    i = 0
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        imgp = os.path.join(imgdir, f"{i}.png")
        i += 1
        cv2.imwrite(imgp, frame)
        
        if i == 400:
            break
    cap.release()
    
    print(os.path.basename(videof) + f" haved resolved")

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()
    
def main():
    videopath = r'\\192.168.1.28\share\测试_202406\0822\A_1724314806144'
    videopath = r'D:\videos'
    savepath = r'D:\videos'
    # video2imgs(videopath, savepath)
    k = 0
    for filename in os.listdir(videopath):
        # filename = "20240822-163506_88e6409d-f19b-4e97-9f01-b3fde259cbff.ts"
        
        file, ext = os.path.splitext(filename)
        if ext not in VideoFormat:
            continue
        
        basename = os.path.basename(videopath)
        imgbase = basename + '-&-' + file
        imgdir = os.path.join(savepath, imgbase)
        if not os.path.exists(imgdir):
            os.mkdir(imgdir)

        videof = os.path.join(videopath, filename)
        video2imgs(videof, imgdir)
        
        k += 1
        if k == 1:
            break

    
    
if __name__ == '__main__':
    main()