数据处理记录

数据处理

# 这是一个示例 Python 脚本。

# 按 Shift+F10 执行或将其替换为您的代码。
# 按 双击 Shift 在所有地方搜索类、文件、工具窗口、操作和设置。
import os
import cv2
import random
import shutil

def get_all_video_path(path, write_txt=False, txt_save_path=""):
    if not os.path.exists(path):
        print(f' \"{path}\" is not existed')
        exit()

    if write_txt:
        if not os.path.exists(txt_save_path):
            os.makedirs(txt_save_path)
        txt_path = os.path.join(txt_save_path, "test_path.txt")
        with open(txt_path, 'w', encoding='utf-8') as f:
            for root, dirs, files in os.walk(path):
                for file in files:
                    file_path = os.path.join(root, file)
                    f.write(file_path + '\n')
        return txt_path  #
    else:
        for root, dirs, files in os.walk(path):
            for file in files:
                file_path = os.path.join(root, file)
                yield str(file_path)


def cap_img_from_video(input_path, output_path, time_interval=60):
    # if not os.path.exists(output_path):
    #     os.makedirs(output_path)
    # 获得视频地址
    video_path_list = get_all_video_path(input_path)
    for video_path in video_path_list:
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            print(f'{input_path} Video not opened')
            exit()
        else:
            print(f'{input_path} Video opened')
        fps = cap.get(cv2.CAP_PROP_FPS)
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        total_frames = cap.get(cv2.CAP_PROP_FRAME_COUNT)
        print(f"FPS: {fps:.2f}, Resolution: {width} x {height}, Total Frames: {total_frames}")
        video_tail = os.path.split(video_path)[-1]
        video_output_path = os.path.join(output_path, video_tail.split('.')[0])
        # print(video_output_path)

        frame_count = 0
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            if frame_count % int(fps * time_interval) == 0:

                crop_a_img(frame, video_output_path, frame_count)
            frame_count += 1


def crop_a_img(img, output_path,frame_count,debug= True):
    if not os.path.exists(output_path):
        os.makedirs(output_path)
    elif debug: # 删除已经产生的文件夹
        shutil.rmtree(output_path)
        os.makedirs(output_path)

    height, width = img.shape[:2]
    # 选取需要截取图像的中心点
    points = [[int(width / 2), int(height / 2)],
              [int(width / 4), int(height / 4)],
              [int(width / 4), int(3 * height / 4)],
              [int(3 * width / 4), int(height / 4)],
              [int(3 * width / 4), int(3 * height / 4)]]
    shapes = [50, 100, 150, 200, 250, 300]
    shrinks = [10, 20, 30, 40, 40, 40]
    for point in points:
        point_key = "".join([str(point[0]), "_", str(point[1])])
        for i, shape in enumerate(shapes):
            shr = shrinks[i]
            # 生成五个候选框
            h_w = [[shape, shape],
                   [shape + random.randint(0, shr), shape + random.randint(0, shr)],
                   [shape - random.randint(0, shr), shape - random.randint(0, shr)],
                   [shape + random.randint(0, shr), shape - random.randint(0, shr)]]
            for h, w in h_w:
                begin_x, begin_y = max(0, int(point[0] - w / 2)), max(0, int(point[1] - h / 2))
                end_x, end_y = min(int(point[0] + w / 2), width), min(int(point[1] + h / 2), height)

                # 将截取的图片保存
                new_img = img[begin_y:end_y, begin_x:end_x, :]
                img_path = os.path.join(output_path, str(frame_count)+ '_' +str(shape)+'_'+ point_key + '_' + str(begin_x) + '_' + str(begin_y) + '_' + str(
                    end_x) + '_' + str(end_y) + '.jpg')
                # print(img_path)
                cv2.imwrite(img_path, new_img)


# 按装订区域中的绿色按钮以运行脚本。
if __name__ == '__main__':
    # 获取地址
    in_path = 'C:\\Users\\Triumphal\\Desktop\\video\\test_1'
    out_path = 'C:\\Users\\Triumphal\\Desktop\\video\\test_crop'

    cap_img_from_video(in_path, out_path)

# 访问 https://www.jetbrains.com/help/pycharm/ 获取 PyCharm 帮助