Как я могу загрузить пользовательские данные в загрузчик данных (pytorch)

Вопрос или проблема

[[img1_filename,img2_filename...,img16_filename],label1], [[img1_filename,img2_filename...,img16_filename],label2],...]

[[[16 изображений (0-16)], label1]], [[16 изображений (4,20)], label2], [[16 изображений (9,24)], label3], [[16 изображений (13,28)], label3]]

X_train = []
        y_train = []
        for data in tqdm(data):  # Цикл по каждой партии
            # Загрузка изображения (X)
            x = data[0]
            y = data[1]
            temp_data_list = []
            for img in x:
                try:
                    img = cv2.imread(img)
                    # примените любую предобработку здесь
                    # img = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
                    # img = self.preprocess_image(img)
                    temp_data_list.append(img)

                except Exception as e:
                    print(e)
                    print('ошибка чтения файла: ', img)
                    # Читать метку (y)
            # label = label_names[y]
            # Добавить пример в массивы
            X_train.append(temp_data_list)
            y_train.append(y)

from pathlib import Path
from torch.utils.data import Dataset

class CustomDataset(Dataset):
    def __init__(self, directory):
        self.directory = directory
        self.images = Path(self.directory).glob("**/*.jpg")

    def __len__(self):
        return len(self.images)

    def __getitem__(self, index):
        # реализуйте вашу пользовательскую логику, которая загружает и
        # преобразует изображения по мере необходимости
        pass

Ответ или решение

Чтобы загрузить кастомные данные в PyTorch DataLoader, вам необходимо создать собственный класс Dataset, который будет наследовать от PyTorch Dataset и реализует методы __len__ и __getitem__. Это позволит вам настроить загрузку данных в соответствии с вашими требованиями и структурой папок.

Шаг 1: Создание класса CustomDataset

В вашем случае, структура данных представляет собой изображения, извлеченные из видео, где каждая выборка содержит 16 изображений и соответствующую метку. Ниже представлен пример реализации этого класса:

import os
import cv2
import numpy as np
from pathlib import Path
import torch
from torch.utils.data import Dataset, DataLoader

class CustomDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = Path(root_dir)
        self.transform = transform
        self.data = []
        self.labels = []

        # Сбор данных и меток
        for class_folder in os.listdir(self.root_dir):
            class_path = self.root_dir / class_folder
            if os.path.isdir(class_path):  # Проверка, что это папка
                for video in os.listdir(class_path):
                    video_path = class_path / video
                    if os.path.isdir(video_path):  # Проверка, что это папка
                        frames = sorted(video_path.glob("*.jpg"))  # Сортируем файлы для последовательности
                        # Генерация выборок
                        num_samples = max(0, len(frames) - 16) // 4 + 1
                        for i in range(num_samples):
                            start_index = i * 4
                            frame_indices = list(range(start_index, start_index + 16))
                            if frame_indices[-1] < len(frames):  # Проверка выхода за границы
                                self.data.append([str(frames[idx]) for idx in frame_indices])
                                self.labels.append(class_folder)  # Используем имя папки как метку

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        # Загрузка изображений
        frame_paths = self.data[index]
        images = []
        for img_path in frame_paths:
            img = cv2.imread(img_path)
            if img is not None:
                img = cv2.resize(img, (112, 112))  # Изменение размера
                images.append(img)
            else:
                print(f'Ошибка загрузки изображения: {img_path}')
        images = np.array(images)  # Преобразуем в NumPy массив

        label = self.labels[index]
        if self.transform:
            images = self.transform(images)  # Применение трансформаций, если есть

        return images, label

# Использование CustomDataset с DataLoader
root_dir = './train_data'  # Укажите путь к вашему набору данных
dataset = CustomDataset(root_dir)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

# Пример итерации по DataLoader
for batch_images, batch_labels in dataloader:
    print(batch_images.shape, batch_labels)

Шаг 2: Объяснение кода

1. Инициализация: В методе __init__ происходит инициализация класса, где мы проходимся по папкам и формируем список изображений и меток.

2. Длина набора данных: Метод __len__ возвращает общее количество образцов данных, что необходимо для DataLoader.

3. Получение данных: Метод __getitem__ загружает изображения по заданному индексу и применяет, при необходимости, предобработку. Мы используем OpenCV для загрузки изображений и изменения их размера до 112×112.

4. Создание DataLoader: Мы создаем экземпляр DataLoader, используя наш класс CustomDataset, указывая размер пакета и возможность перемешивания данных для эффективного обучения.

Шаг 3: Обработка ошибок

Убедитесь, что при загрузке изображений вы обрабатываете возможные ошибки, например, если файл отсутствует или файл оказал не изображением. В приведенном выше примере добавляется вывод сообщений об ошибках.

Заключение

Теперь у вас есть полноценное решение для загрузки вашего кастомного набора данных в PyTorch. Это позволит вам эффективно осуществлять обучение модели, используя сложные датасеты. Не забывайте экспериментировать с предварительной обработкой изображений и параметрами загрузки, чтобы повысить качество обучения вашей нейронной сети.