使用OpenCV GUI清理数据集 | 为目标检测模型创建更好的数据集

点击下方卡片，关注“小白玩转Python”公众号

在深度学习中有几件重要的事情，我认为数据是最关键的。如果没有合适的数据，要取得好的结果是非常困难的。即使你用强大的预训练模型和GPU训练模型，你的模型也可能表现不佳。在本文中，我将讨论如何获取数据集、分析数据集并使用简单的OpenCV GUI应用程序对其进行清理。我将使用Python作为编程语言。

标记的图像

获取数据的最佳网站

有两种主要方法可以获取数据。第一种方法是生成自己的数据。你可以使用手机相机拍照，或者从互联网上抓取。获取图像后，你需要为它们标记（用于目标检测），这个过程可能需要很长时间。第二种方法更容易，你可以从网站上获取数据，互联网上有许多免费的数据集。

Kaggle：https://www.kaggle.com/datasets
Roboflow：https://universe.roboflow.com/

在选择数据集之前，我总是检查这两个网站。如果我找不到合适的数据集，我就会使用第一种方法。

如何判断一个数据集是否可以直接使用？

如果你创建了自己的数据集，你可能不需要清理它，因为你自己拍摄了照片并对它们进行了标记。你可以在创建过程中清理它，无需额外的清理步骤。
如果你直接从互联网上下载数据集，你将需要检查它。可能会有无关的图像、错误的标签或低质量的图像。

我通常从随机检查数据集开始。我只是检查图像的质量，并观察对象的分布，我不太关注数据集。如果一切都看起来可以接受，我就训练一个小模型。在评估结果之后，我通常可以理解数据是否需要清理或是否可以用于训练。

如果我决定数据还不准备好使用，我会使用我创建的简单OpenCV GUI对其进行清理。现在我将向你展示如何创建此GUI并用于清理数据集。

清理数据的GUI

这个GUI非常简单易用，而且非常有效。用户可以在一个屏幕上看到标签、边界框和图像本身，决定是否可以将其用于数据集。

如果用户决定图像是合适的，图像和标签将被复制到另一个文件夹中，稍后将用于训练。（按“s”键）
如果用户认为图像不合适，图像和标签将被删除。（按空格键）

实际上，这非常有效，因为你可以看到图像，同时也可以看到标签。使用此GUI，用户可以将不同的数据集合并到一个文件夹中，并为其选择最佳图像。

代码

导入必要的库

import cv2
import os
import shutil
import random

创建图像和标签的文件夹

# Define the path to your folder containing images# folders for all images and labels
image_folder = 'original-images' 
label_folder = 'original-labels'# folders for selected images and labels
saved_images_folder= 'saved-images'
saved_labels_folder = 'saved-labels'

获取图像文件路径并对它们进行随机排序

注意：在目标检测数据集中，每个图像文件都有一个相应的带有相同名称的标签文件。这种命名约定是通用的。例如，如果有一个名为“image23.jpg”的图像文件，则存在一个名为“image23.txt”的文件，其中包含关于标签和边界框的信息。

# Get a list of all image files in the folder
image_files = [f for f in os.listdir(image_folder) if f.endswith(('.jpg', '.png', '.jpeg'))]# Randomly shuffle the list of image files
random.shuffle(image_files)

显示图像、标签和边界框的函数

# Function to display image and handle keypress
def display_image(image_path, label_path):img = cv2.imread(image_path)with open(label_path, 'r') as f:lines = f.readlines()for line in lines:class_id, x_center, y_center, width, height = map(float, line.strip().split())# Convert YOLO format to pixel coordinatesimg_h, img_w, _ = img.shapex_center *= img_wy_center *= img_hwidth *= img_wheight *= img_hx_min = int(x_center - (width / 2))y_min = int(y_center - (height / 2))x_max = int(x_center + (width / 2))y_max = int(y_center + (height / 2))# Draw bounding boxcv2.rectangle(img, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)# Get class name based on class_idclass_names =   ['black-bishop', 'black-king', 'black-knight', 'black-pawn', 'black-queen', 'black-rook', 'white-bishop', 'white-king', 'white-knight', 'white-pawn', 'white-queen', 'white-rook']class_name = class_names[int(class_id)]# Write class namecv2.putText(img, class_name, (x_min, y_min - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)cv2.imshow('Image', img)key = cv2.waitKey(0)return key

主循环

# Main loop to display images and handle selection
for image_file in image_files:image_path = os.path.join(image_folder, image_file)label_path = os.path.join(label_folder, image_file.replace('.jpg', '.txt'))  # Assuming label files have the same name as image files with a .txt extension# Check if the label file existsif os.path.exists(label_path):key = display_image(image_path,label_path)# Press 's' to save the image and its corresponding label file to the selected folderif key == ord('s'):selected_image_path = os.path.join(saved_images_folder, image_file)selected_label_path = os.path.join(saved_labels_folder, image_file.replace('.jpg', '.txt'))shutil.move(image_path, selected_image_path)  # Move the image file to the selected foldershutil.move(label_path, selected_label_path)  # Move the label file to the saved labels folderprint(f"Image '{image_file}' and its corresponding label file saved.")elif key == ord(' '):os.remove(image_path)os.remove(label_path)continue  # Move to the next imageelif key == ord('q'):breakelse:print(f"No label file found for image '{image_file}'.")continue  # Skip to the next imagecv2.destroyAllWindows()