opencv_webcam/opencv_webcam.py

# OpenCV Webcam Script v0.2
# 创建人：曾逸夫
# 创建时间：2022-01-02

import cv2
from pathlib import Path
import argparse
import time
import os
import glob
import re
import os


def parse_args(known=False):
    parser = argparse.ArgumentParser(description='OpenCV Webcam Script')
    parser.add_argument('--device', '-dev', default=0,
                        type=int, help='device index for webcam')
    parser.add_argument('--quit', '-q', default="q",
                        type=str, help='quit key for webcam')
    parser.add_argument('--is_autoSaveFrame', '-isasf',
                        action='store_true', help='is auto save frame')
    parser.add_argument('--is_headSaveFrame', '-ishsf',
                        action='store_true', help='is head save frame')
    parser.add_argument('--is_resizeFrame', '-isrf',
                        action='store_true', help='is resize frame')
    parser.add_argument('--frame_saveDir', '-fsd',
                        default="./WebcamFrame", type=str, help='save frame dir')
    parser.add_argument('--frame_nSave', '-fns', default=1,
                        type=int, help='n frames save a frame (auto save frame)')
    parser.add_argument('--frame_capKey', '-fck', default="a",
                        type=str, help='frame capture key (head save frame)')
    parser.add_argument('--resize_frame', '-rf',
                        default=[640, 480], type=int, nargs='+', help='resize frame save')
    parser.add_argument('--resizeRatio_frame', '-rrf',
                        default=1.0, type=float, help='resize ratio frame save')
    parser.add_argument('--frame_namePrefix', '-fnp',
                        default="frame", type=str, help='frame name prefix')
    parser.add_argument('--frame_saveStyle', '-fss',
                        default="jpg", type=str, help='frame save style')
    parser.add_argument('--jpg_quality', '-jq',
                        default=95, type=int, help='frame save jpg quality (0-100) default 95')
    parser.add_argument('--png_quality', '-pq',
                        default=3, type=int, help='frame save jpg quality (0-9) default 3')

    args = parser.parse_known_args()[0] if known else parser.parse_args()
    return args


# 保存路径管理
def increment_path(path, exist_ok=False, sep='', mkdir=False):
    # 引用：https://github.com/ultralytics/yolov5/blob/master/utils/general.py
    path = Path(path)  # os-agnostic
    if path.exists() and not exist_ok:
        suffix = path.suffix
        path = path.with_suffix('')
        dirs = glob.glob(f"{path}{sep}*")  # similar paths
        matches = [re.search(rf"%s{sep}(\d+)" % path.stem, d) for d in dirs]
        i = [int(m.groups()[0]) for m in matches if m]  # indices
        n = max(i) + 1 if i else 2  # increment number
        path = Path(f"{path}{sep}{n}{suffix}")  # update path
    dir = path if path.suffix == '' else path.parent  # directory
    if not dir.exists() and mkdir:
        dir.mkdir(parents=True, exist_ok=True)  # make directory
    return path


# 帧保存
def frame_opt(frame, frame_savePath, frame_num, is_resizeFrame, resize_frame, resizeRatio_frame,
              frame_namePrefix, frame_saveStyle, jpg_quality, png_quality):
    if (is_resizeFrame):  # resize frame
        w_resize = int(
            resize_frame[0] * resizeRatio_frame)  # 重塑宽度
        h_resize = int(
            resize_frame[1] * resizeRatio_frame)  # 重塑高度
        frame_new = cv2.resize(
            frame, (w_resize, h_resize), interpolation=cv2.INTER_AREA)  # 重塑

        if (frame_saveStyle == 'jpg'):
            cv2.imwrite(
                f'./{frame_savePath}/{frame_namePrefix}-{frame_num}.{frame_saveStyle}', frame_new,
                [int(cv2.IMWRITE_JPEG_QUALITY), jpg_quality])
        elif (frame_saveStyle == 'png'):
            cv2.imwrite(
                f'./{frame_savePath}/{frame_namePrefix}-{frame_num}.{frame_saveStyle}', frame_new,
                [int(cv2.IMWRITE_PNG_COMPRESSION), png_quality])
        else:
            print(f'帧格式有问题！')
            return
    else:
        if (frame_saveStyle == 'jpg'):
            cv2.imwrite(
                f'./{frame_savePath}/{frame_namePrefix}-{frame_num}.{frame_saveStyle}', frame,
                [int(cv2.IMWRITE_JPEG_QUALITY), jpg_quality])
        elif (frame_saveStyle == 'png'):
            cv2.imwrite(
                f'./{frame_savePath}/{frame_namePrefix}-{frame_num}.{frame_saveStyle}', frame,
                [int(cv2.IMWRITE_PNG_COMPRESSION), png_quality])
        else:
            print(f'帧格式有问题！')
            return


# webcam opencv
def webcam_opencv(device_index, quit_key, is_autoSaveFrame, frame_saveDir, frame_nSave, is_headSaveFrame,
                  frame_capKey, is_resizeFrame, resize_frame, resizeRatio_frame, frame_namePrefix, frame_saveStyle,
                  jpg_quality, png_quality):
    if (quit_key == frame_capKey):  # 判断快捷键冲突
        print(f'快捷键冲突! 程序结束！')
        return

    s_time = time.time()  # 起始时间

    cap = cv2.VideoCapture(device_index)  # 连接设备
    frame_width = cap.get(3)  # 宽度
    frame_height = cap.get(4)  # 高度
    fps = cap.get(5)  # 帧率
    print(f'宽度：{frame_width}, 高度：{frame_height}， FPS：{fps}')
    success, _ = cap.read()  # 读取设备

    if success:
        print(f'摄像头连接成功！')

        # 帧保存路径管理
        if is_autoSaveFrame or is_headSaveFrame:
            frame_savePath = increment_path(
                Path(f"{frame_saveDir}") / "frames", exist_ok=False)  # 增量运行
            frame_savePath.mkdir(parents=True, exist_ok=True)  # 创建目录

        frame_num = 0  # 帧数
        while(1):
            ret, frame = cap.read()  # 捕获画面
            frame_num += 1
            print(f'第{frame_num}帧')
            cv2.imshow("capture", frame)  # 显示画面
            # print(frame.shape) # h,w,c
            if (is_autoSaveFrame):  # 自动保存
                if (frame_num % frame_nSave == 0):  # 每隔n帧保存一次
                    frame_opt(frame, frame_savePath, frame_num, is_resizeFrame, resize_frame, resizeRatio_frame,
                              frame_namePrefix, frame_saveStyle, jpg_quality, png_quality)
            if (is_headSaveFrame):  # 手动保存
                if cv2.waitKey(20) & 0xFF == ord(frame_capKey):  # 保存键
                    frame_opt(frame, frame_savePath, frame_num, is_resizeFrame, resize_frame, resizeRatio_frame,
                              frame_namePrefix, frame_saveStyle, jpg_quality, png_quality)
            if cv2.waitKey(20) & 0xFF == ord(quit_key):  # 退出 ord：字符转ASCII码
                break

        print(f'一共处理了{frame_num}帧')
        # 释放缓存
        cap.release()
        cv2.destroyAllWindows()
    else:
        print(f'摄像头连接异常！')

    print(f'程序结束！')
    e_time = time.time()  # 终止时间
    print(f'用时：{round((e_time - s_time), 3)}秒, {round((e_time - s_time)/60, 3)}分, {round((e_time - s_time)/3600, 3)}小时')


def main(args):
    device_index = args.device
    quit_key = args.quit
    is_autoSaveFrame = args.is_autoSaveFrame
    is_headSaveFrame = args.is_headSaveFrame
    frame_saveDir = args.frame_saveDir
    frame_nSave = args.frame_nSave
    frame_capKey = args.frame_capKey
    resize_frame = args.resize_frame
    is_resizeFrame = args.is_resizeFrame
    resizeRatio_frame = args.resizeRatio_frame
    frame_namePrefix = args.frame_namePrefix
    frame_saveStyle = args.frame_saveStyle
    jpg_quality = args.jpg_quality
    png_quality = args.png_quality

    # 调用webcam opencv
    webcam_opencv(device_index, quit_key, is_autoSaveFrame, frame_saveDir, frame_nSave,
                  is_headSaveFrame, frame_capKey, is_resizeFrame, resize_frame, resizeRatio_frame,
                  frame_namePrefix, frame_saveStyle, jpg_quality, png_quality)


if __name__ == '__main__':
    args = parse_args()
    main(args)