python图像/视频处理函数

python图像/视频处理函数

import cv2
import os
import argparse

from tqdm import tqdm


def show_video(path, fps=3):
    """ opencv显示视频,按q中断
    """
    if not os.path.exists(path):
        raise Exception('video dont exists!')

    cap = cv2.VideoCapture(path)

    while(cap.isOpened()):
        ret, frame = cap.read()
        cv2.imshow('frame', frame)
        if cv2.waitKey(fps) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()


def video2img(input_path, out_path, save=True, show=True, time=3):
    # 参考博客:https://blog.csdn.net/qq_25436597/article/details/79621833
    ''' 将读取的视频转化为图像并保存在out_path文件夹下
    params:
        input_path: 原始视频文件文件夹
        out_path:   保存的文件夹
        save:   是否需要保存
        show:   是否需要显示视频
        time:   视频显示帧率
    '''
    if not os.path.exists(input_path):
        raise Exception('video dont exists!')

    if not os.path.exists(out_path):
        os.makedirs(out_path)
        print('{} is not exists,will makedir {}'.format(out_path, out_path))

    i = 0
    cap = cv2.VideoCapture(input_path)
    ret, frame = cap.read()
    while(ret):
        ret, frame = cap.read()
        if save:
            # cv2.imwrite(frame,)
            img_name = str(i)+'.jpg'
            save_path = os.path.join(out_path, img_name)
            cv2.imwrite(save_path, frame)
            i += 1
        # gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

        if show:
            cv2.imshow('frame', frame)
            if cv2.waitKey(time) & 0xFF == ord('q'):
                break

    if save:
        print('save video to img in :{}!'.format(out_path))
    cap.release()
    cv2.destroyAllWindows()


def img2video(input_path, out_path, fps=30, img_size=None, num_img=None):
    # 参考博客:https://theailearner.com/2018/10/15/creating-video-from-images-using-opencv-python/
    """ 将input_path文件夹下面的img变成视频.avi，保存在out_path文件下
    param:  
        input_path: 图像文件夹路径
        out_path:   输出视频文件
        fps:    输出视频的帧率
        img_size:   图像的大小，默认为None
        num_img:    文件夹下面图像的数量，默认为None

    """
    if not os.path.exists(input_path):
        raise Exception('img dir dont exists!')

    # if not os.path.exists(out_path):
    #     os.makedirs(out_path)
    #     print('{} is not exists,will makedir {}'.format(out_path,out_path))

    img_name_list = os.listdir(input_path)
    img_name_list.sort(key=lambda x: int(x.split('.')[0]))

    if not num_img:
        num_img = len(img_name_list)
    print('图像总数:{}'.format(num_img))

    if not img_size:
        # 读取第一张图片
        first_img_name = os.path.join(input_path, img_name_list[0])
        first_img = cv2.imread(first_img_name)
        img_h,img_w = first_img.shape[:2]  # [H,W,3]
    print('img_size:', img_size)

    # 转为视频

    # img_array = []
    # for img_name in img_name_list:
    #     img_path = os.path.join(input_path, img_name)
    #     img = cv2.imread(img_path)
    #     img_array.append(img)
    #     print(img_name)

    out = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(
        *'DIVX'), fps, (img_w, img_h))   # 注意这里的size跟img_size是反过来的
    # 生成mp4格式
    # out = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc('m', 'p', '4', 'v'), fps, (img_w, img_h))   # 注意这里的size跟img_size是反过来的
    
    bar = tqdm(img_name_list)
    for img_name in bar:
        img_path = os.path.join(input_path, img_name)
        img = cv2.imread(img_path)
        out.write(img)

    out.release()



if __name__ == "__main__":
    input_path = r'result/result_8'
    out_path = r'video/LW_Bbox_001.avi'
    img2video(input_path,out_path)
    # video2img
    # parser = argparse.ArgumentParser(description='video2img')
    # parser.add_argument('--input_path', type=str, default='video/test001.mp4',
    #                     help='path of video')
    # parser.add_argument('--out_path', type=str, default='img',
    #                     help='path of output img')
    # opt = parser.parse_args()
    # video2img(opt.input_path, opt.out_path)