自制数据集处理

1. 视频抽帧

每90帧抽一帧

import cv2
import os
# 处理视频为图片

def extract_frames(video_path, output_folder, frame_interval=30):
    # 打开视频文件
    video_capture = cv2.VideoCapture(video_path)
    # 确保视频文件已成功打开
    if not video_capture.isOpened():
        print("Error: Unable to open video file.")
        return

    # 确保输出文件夹存在
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    frame_count = 0
    while True:
        # 读取一帧
        ret, frame = video_capture.read()
        if not ret:
            break

        # 在每个指定的帧间隔保存帧
        if frame_count % frame_interval == 0:
            # 保存当前帧为图像文件
            frame_filename = os.path.join(output_folder, f"frame_{frame_count}.jpg")
            cv2.imwrite(frame_filename, frame)

        frame_count += 1

    # 释放视频捕获对象
    video_capture.release()


# 视频文件路径
video_path = "maize_tassel.mp4"
# 输出文件夹路径
output_folder = ("maize_tassel")
# 每30帧抽取一张图片
extract_frames(video_path, output_folder, frame_interval=90)

2. 图像裁剪

裁剪5120*2700图像为8份1280*1280图像

import os
from PIL import Image

def crop_images(input_folder, output_folder, crop_width, crop_height):
    # 确保输出文件夹存在
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # 获取输入文件夹中的所有图像文件
    images = [f for f in os.listdir(input_folder) if f.endswith(('jpg', 'jpeg', 'png', 'bmp', 'tiff', 'png'))]

    for image_file in images:
        # 打开图像
        with Image.open(os.path.join(input_folder, image_file)) as img:
            img_width, img_height = img.size

            # 确保图像尺寸为5120x2700
            if img_width != 5120 or img_height != 2700:
                print(f"图像 {image_file} 尺寸不符合要求，跳过裁剪。")
                continue

            # 裁剪并保存图像
            for i in range(2):  # 高度方向裁剪2个块
                for j in range(4):  # 宽度方向裁剪4个块
                    left = j * crop_width
                    top = i * crop_height
                    right = left + crop_width
                    bottom = top + crop_height

                    cropped_img = img.crop((left, top, right, bottom))
                    cropped_img_name = f"{os.path.splitext(image_file)[0]}_crop_{i}_{j}.png"
                    cropped_img.save(os.path.join(output_folder, cropped_img_name))
                    print(f"保存裁剪图像: {cropped_img_name}")

# 示例用法
input_folder = './maize1'  # 输入文件夹路径
output_folder = './maize_split_8'  # 输出文件夹路径
crop_width = 1280  # 裁剪宽度
crop_height = 1280  # 裁剪高度

crop_images(input_folder, output_folder, crop_width, crop_height)

3. 图像画框

将一个voc格式的xml文件中的框画在一个图像中。

import os
import cv2
import xml.etree.ElementTree as ET


def draw_bounding_boxes(image_path, xml_path, output_path):
    # 读取图像
    image = cv2.imread(image_path)
    if image is None:
        print(f"Error: Unable to open image file {image_path}")
        return

    # 解析XML文件
    tree = ET.parse(xml_path)
    root = tree.getroot()

    # 获取所有object标签
    for obj in root.findall('object'):
        # 获取bndbox标签
        bndbox = obj.find('bndbox')
        if bndbox is None:
            continue

        # 提取边框坐标
        xmin = int(bndbox.find('xmin').text)
        ymin = int(bndbox.find('ymin').text)
        xmax = int(bndbox.find('xmax').text)
        ymax = int(bndbox.find('ymax').text)

        # 在图像上绘制红色边框，线条宽度为2
        cv2.rectangle(image, (xmin, ymin), (xmax, ymax), (0, 0, 255), 10)

    # 保存结果图像
    cv2.imwrite(output_path, image)
    print(f"Output saved to {output_path}")


def process_directories(xml_dir, img_dir, output_dir):
    # 确保输出目录存在
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 遍历XML目录中的所有文件
    for xml_file in os.listdir(xml_dir):
        if xml_file.endswith('.xml'):
            xml_path = os.path.join(xml_dir, xml_file)
            img_filename = xml_file.replace('.xml', '.jpg')
            image_path = os.path.join(img_dir, img_filename)

            if not os.path.exists(image_path):
                img_filename = xml_file.replace('.xml', '.png')
                image_path = os.path.join(img_dir, img_filename)

            if os.path.exists(image_path):
                output_path = os.path.join(output_dir, img_filename)
                draw_bounding_boxes(image_path, xml_path, output_path)
            else:
                print(f"Warning: No image file found for {xml_file}")


# 示例用法
xml_dir = './outputs'  # 替换为包含XML标签文件的目录路径
img_dir = './images'  # 替换为包含图像文件的目录路径
output_dir = './end'  # 替换为输出图像的保存目录路径

process_directories(xml_dir, img_dir, output_dir)