图片视频分类标记

import os
import cv2
import time
import argparse
import subprocess
import numpy as np


def to_to_mp4(ts_file: str) -> str:
    mp4_file = ts_file.replace(".ts", ".mp4")
    if os.path.exists(mp4_file):
        return mp4_file
    cmd = [
        "ffmpeg",
        "-y",
        "-i",
        ts_file,
        "-c:v",
        "copy",
        "-c:a",
        "copy",
        "-movflags",
        "+faststart",
        mp4_file,
    ]
    proc = subprocess.run(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=False,
    )
    if proc.returncode == 0:
        return mp4_file
    else:
        return None


def classify_video(video_path):
    print(f"开始处理视频: {video_path}")
    global index, jiedi_count, yandian_count, jueyuanzhebi_count, anzhuangyinxian_count, zuoye_count
    base_time = int(time.time())
    video = None
    if video_path.endswith(".mp4"):
        video = cv2.VideoCapture(video_path, cv2.CAP_FFMPEG)
    elif video_path.endswith((".png", ".jpg", ".jpeg")):
        frame = cv2.imdecode(np.fromfile(video_path, dtype=np.uint8), cv2.IMREAD_COLOR)
    while True:
        if video:
            ret, frame = video.read()
            if not ret:
                break
        h, w = frame.shape[:2]
        scale = 640 / max(w, h)
        frame = cv2.resize(
            frame, (int(w * scale), int(h * scale)), interpolation=cv2.INTER_AREA
        )
        # fmt: off
        show_frame = frame.copy()
        cv2.putText(show_frame, f"1 - jiedi",          (10,  30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1,)
        cv2.putText(show_frame, f"2 - yandian",        (10,  60), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1,)
        cv2.putText(show_frame, f"3 - jueyuanzhebi",   (10,  90), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1,)
        cv2.putText(show_frame, f"4 - anhuangyinxian", (10, 120), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1,)
        cv2.putText(show_frame, f"5 - zuoye",          (10, 150), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1,)
        # fmt: on
        cv2.imshow("frame", show_frame)
        key_code = cv2.waitKey(0) & 0xFF
        skip_frames = 8
        if key_code == ord("q"):
            break
        elif key_code == ord("1"):
            path = os.path.join("jiedi", f"{base_time}_{index}.jpg")
            cv2.imwrite(path, frame)
            index += 1
            skip_frames = 4
            jiedi_count += 1
            print(f"保存接地: {path}")
        elif key_code == ord("2"):
            path = os.path.join("yandian", f"{base_time}_{index}.jpg")
            cv2.imwrite(path, frame)
            index += 1
            skip_frames = 4
            yandian_count += 1
            print(f"保存验电: {path}")
        elif key_code == ord("3"):
            path = os.path.join("jueyuanzhebi", f"{base_time}_{index}.jpg")
            cv2.imwrite(path, frame)
            index += 1
            skip_frames = 4
            jueyuanzhebi_count += 1
            print(f"保存绝缘遮蔽: {path}")
        elif key_code == ord("4"):
            path = os.path.join("anzhuangyinxian", f"{base_time}_{index}.jpg")
            cv2.imwrite(path, frame)
            index += 1
            skip_frames = 4
            anzhuangyinxian_count += 1
            print(f"保存安装引线: {path}")
        elif key_code == ord("5"):
            path = os.path.join("zuoye", f"{base_time}_{index}.jpg")
            cv2.imwrite(path, frame)
            index += 1
            skip_frames = 16
            zuoye_count += 1
            print(f"保存作业: {path}")
        elif key_code == 8:
            skip_frames = -32
        elif key_code == ord("\r"):
            skip_frames = 32
        elif key_code == ord(" "):
            skip_frames = 64
        else:
            skip_frames = 16
        print(f"skip_frames = {skip_frames}")
        if video:
            video.set(
                cv2.CAP_PROP_POS_FRAMES,
                video.get(cv2.CAP_PROP_POS_FRAMES) + skip_frames,
            )
        else:
            break
    cv2.destroyAllWindows()


parse = argparse.ArgumentParser(description="视频分类")
parse.add_argument("--dir", type=str, required=True, help="视频路径")
args = parse.parse_args()
dir_path = args.dir
index = 0
image_size = 0
video_size = 0
jiedi_count = 0
yandian_count = 0
jueyuanzhebi_count = 0
anzhuangyinxian_count = 0
zuoye_count = 0
os.makedirs("jiedi", exist_ok=True)
os.makedirs("yandian", exist_ok=True)
os.makedirs("jueyuanzhebi", exist_ok=True)
os.makedirs("anzhuangyinxian", exist_ok=True)
os.makedirs("zuoye", exist_ok=True)
for root, dirs, files in os.walk(dir_path):
    for file in files:
        if file.endswith((".png", ".jpg", ".jpeg")):
            image_size += 1
            classify_video(os.path.join(root, file))
        elif file.endswith((".ts")):
            video_size += 1
            ts_file = os.path.join(root, file)
            mp4_file = to_to_mp4(ts_file)
            if mp4_file:
                classify_video(mp4_file)
            else:
                print(f"转换失败: {file}")
        elif file.endswith((".mp4")):
            mp4_file = ts_file = os.path.join(root, file)
            ts_file = mp4_file.replace(".mp4", ".ts")
            if os.path.exists(ts_file):
                continue
            video_size += 1
            classify_video(mp4_file)
        else:
            print(f"未知文件类型: {file}")
print("- 原始资料:")
print(f"  * 图片数量: {image_size}张")
print(f"  * 视频数量: {video_size}个")
print("- 数据集:")
print(f"  * 接地: {jiedi_count}张")
print(f"  * 验电: {yandian_count}张")
print(f"  * 绝缘: {jueyuanzhebi_count}张")
print(f"  * 引线: {anzhuangyinxian_count}张")
print(f"  * 作业: {zuoye_count}张")