Датасет — это основа для обучения моделей компьютерного зрения. В этой статье мы разберем, как собрать и подготовить собственный датасет для работы с 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:
- Установите библиотеку:
pip install kaggle
. - Скачайте датасет через 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:
- Установите LabelImg:
pip install labelImg
labelImg # Запуск
- Рисуйте 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 для решения задач классификации, детекции или сегментации.
Добавить комментарий