Создание собственного датасета для OpenCV на Python: методы и примеры

Датасет — это основа для обучения моделей компьютерного зрения. В этой статье мы разберем, как собрать и подготовить собственный датасет для работы с OpenCV в Python. Рассмотрим ручной сбор данных, парсинг из интернета, аугментацию и разметку.


1. Ручной сбор данных с веб-камеры

Пример: Создание датасета для распознавания жестов.

import cv2
import os

# Создаем папку для сохранения изображений
class_name = "gesture_ok"
os.makedirs(f"dataset/{class_name}", exist_ok=True)

cap = cv2.VideoCapture(0)
counter = 0

while True:
    ret, frame = cap.read()
    if not ret:
        break

    cv2.imshow('Frame', frame)

    # Сохраняем кадр при нажатии пробела
    key = cv2.waitKey(1)
    if key == 32:  # Код пробела
        cv2.imwrite(f"dataset/{class_name}/{counter}.jpg", frame)
        print(f"Saved image {counter}")
        counter += 1
    elif key == 27:  # ESC для выхода
        break

cap.release()
cv2.destroyAllWindows()

Совет:

  • Собирайте данные при разном освещении и ракурсах.
  • Храните изображения в папках с названиями классов (например, dataset/ok, dataset/not_ok).

2. Использование готовых датасетов

Популярные источники:

Пример загрузки изображений с Kaggle:

  1. Установите библиотеку: pip install kaggle.
  2. Скачайте датасет через API:
kaggle datasets download -d username/dataset-name

3. Парсинг изображений из интернета

Пример с использованием библиотеки requests и BeautifulSoup:

import requests
from bs4 import BeautifulSoup
import os

query = "cats"
url = f"https://www.pexels.com/search/{query}/"

response = requests.get(url)
soup = BeautifulSoup(response.text, 'html.parser')
images = soup.find_all('img', {'class': 'photo-item__img'})

os.makedirs("dataset/cats", exist_ok=True)

for i, img in enumerate(images[:50]):  # Скачиваем первые 50 изображений
    img_url = img['src']
    img_data = requests.get(img_url).content
    with open(f"dataset/cats/{i}.jpg", 'wb') as f:
        f.write(img_data)

Важно: Убедитесь, что использование изображений соответствует лицензии источника.


4. Аугментация данных

Увеличьте размер датасета, применяя преобразования:

  • Повороты
  • Изменение яркости
  • Зеркальное отражение

Пример с библиотекой imgaug:

import imgaug.augmenters as iaa
import cv2
import os

# Загрузка изображения
image = cv2.imread("dataset/cats/0.jpg")

# Определение аугментаций
augmenter = iaa.Sequential([
    iaa.Fliplr(0.5),  # Зеркальное отражение с вероятностью 50%
    iaa.GaussianBlur(sigma=(0, 1.0)),
    iaa.AdditiveGaussianNoise(scale=0.1*255)
])

# Применение аугментаций
augmented_images = augmenter(images=[image] * 5)  # Создаем 5 вариантов

# Сохранение результатов
for i, img in enumerate(augmented_images):
    cv2.imwrite(f"dataset/cats/augmented_{i}.jpg", img)

5. Разметка данных

Для задач детекции объектов используйте инструменты вроде LabelImg:

  1. Установите LabelImg:
pip install labelImg
labelImg  # Запуск
  1. Рисуйте bounding boxes вокруг объектов и сохраняйте разметку в формате XML (PascalVOC) или JSON (COCO).

Пример чтения разметки в Python:

import xml.etree.ElementTree as ET

tree = ET.parse('dataset/cats/0.xml')
root = tree.getroot()

for obj in root.findall('object'):
    name = obj.find('name').text
    bbox = obj.find('bndbox')
    xmin = int(bbox.find('xmin').text)
    ymin = int(bbox.find('ymin').text)
    xmax = int(bbox.find('xmax').text)
    ymax = int(bbox.find('ymax').text)
    print(f"Class: {name}, BBox: [{xmin}, {ymin}, {xmax}, {ymax}]")

6. Предобработка данных

Пример нормализации и изменения размера:

import cv2
import numpy as np

def preprocess_image(image_path, target_size=(224, 224)):
    image = cv2.imread(image_path)
    image = cv2.resize(image, target_size)
    image = image / 255.0  # Нормализация [0, 1]
    return image

# Применение ко всему датасету
preprocessed_images = []
for img_path in os.listdir("dataset/cats"):
    preprocessed = preprocess_image(f"dataset/cats/{img_path}")
    preprocessed_images.append(preprocessed)

7. Структура датасета

Рекомендуемая организация папок:

dataset/
├── train/
│   ├── class1/
│   └── class2/
└── test/
    ├── class1/
    └── class2/

Используйте sklearn.model_selection.train_test_split для разделения:

from sklearn.model_selection import train_test_split

X = [...]  # Список путей к изображениям
y = [...]  # Метки классов

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)

Заключение

Сбор датасета — важный этап для проектов компьютерного зрения. Используйте комбинацию ручного сбора, парсинга и аугментации. Не забывайте о качественной разметке и предобработке. С готовым датасетом вы сможете обучать модели в OpenCV, TensorFlow или PyTorch для решения задач классификации, детекции или сегментации.

Комментарии

Добавить комментарий

Ваш адрес email не будет опубликован. Обязательные поля помечены *