Тенденция применения беспилотных летательных аппаратов (БПЛА) продолжает развиваться и процветать. Оснащение беспилотников камерами и навигационным оборудованием геодезического класса точности позволяет получать ортофотопланы с сантиметровой точностью. Расширить возможности БПЛА можно применив нейронные сети, способные распознавать объекты на фотографиях. В статье рассмотрен процесс подготовки фотографий с БПЛА, разметки объектов для обучения нейронной сети, ее обучения и получения результата в виде выявления объекта на новом фото на реальном участке железнодорожного перегона, определяемые объекты — пикетные столбики. Исходный код обработки данных и обучения модели выгружен на GitHub.

Постановка задачи

Железнодорожный путь требуют регулярного мониторинга различных элементов инфраструктуры дляобеспечения безопасности и бесперебойного движения поездов. Один из элементов инфраструктуры — пикетные столбики. Для автоматизации обработки и анализа собранных данных, в частности изображений, можно применить методы компьютерного зрения. Задача данной работы — автоматизировать процесс обнаружения пикетов на фотографиях, полученных с беспилотника.

В рамках данной задачи разберем систему, способную детектировать пикеты на изображениях, снятых с высоты порядка 200 метров над землей. Система должна быть способна:

• определять пикеты на изображениях с камеры БПЛА в разрешении 3264x4928 пикселей;

• справляться с возможной вариативностью условий съёмки, таких как освещение, угол обзора и фон (в данной работе не реализовано, но нужно учесть на будущее для повышения качества модели в разных условиях);

• работать на новых изображениях с аналогичной точностью после обучения.

Для этого применена нейросеть YOLOv5 — модель, широко используемая для детекции объектов, обладающая высокой точностью и скоростью.

Исходные данные

Для распознавания пикетных столбиков были собраны данные в виде 594 цветных фотографии железнодорожного полотна, сделанных с высоты 200 метров с использованием БПЛА. Разрешение каждого изображения составляет 3264×4928 пикселей в формате JPG, фото в вертикальной ориентации.

В основном на каждом фото имеется по 2 пикета, реже — один или три пикета, размер их на фото будет достаточно мал, примерно 50×50 пикселей. Расположение пикетов на фото чаще всего в средней полосе кадра, но бывают и в разных местах.

Настройка окружения для работы над задачей

Для разработки и обучения модели распознавания пикетных столбиков необходим фреймворк PyTorch и предобученная модель YOLOv5. Окружение можно развернуть в Anaconda с использованием изолированной среды, в которой необходимо установить Python версии 3.10, эта версия выбрана для совместимости с PyTorch и YOLOv5. Помимо этого, для обучения нейронной сети на GPU необходимы пакеты CUDA Toolkit и cuDNN Library. Как скачать, установить и проверить рассказано ЗДЕСЬ:

Основные этапы настройки окружения

1. Создание и настройка виртуальной среды

Изолированная виртуальная среда в Anaconda позволяет управлять версиями библиотек и исключить конфликт зависимости между ними. Среда создается командой:

conda create -n yolov5_env python=3.10

где yolov5_env – название создаваемой среды

2. Установка основных библиотек

В среде yolov5_env необходимо установить/импортировать следующие ключевые библиотеки:

PyTorch: для реализации и обучения модели YOLOv5;

Torchvision: для предварительной обработки изображений и работы с датасетами;

OpenCV: для операций с изображениями, включая разрезание, изменение размера и аугментацию данных;

Pillow: для работы с изображениями и преобразования их формата;

Matplotlib и Seaborn: для визуализации данных и результатов;

Numpy (np): библиотека для работы с многомерными массивами и высокоуровневыми математическими функциями;

lxml : для работы с аннотациями в формате XML;

os: стандартная библиотека Python для работы с файловой системой (создание папок, перемещение файлов и др.);

shutil: стандартная библиотека Python для операций с файлами и директориями (копирование, перемещение и удаление);

random: стандартная библиотека Python для генерации случайных чисел.

3. Установка и настройка YOLOv5

YOLOv5 установливается с использованием репозитория на GitHub, для этого репозиторий клонируется (командой git clone) в свою рабочую папку локально на ПК. Далее модель настраивается для работы в созданной виртуальной среде, что позволяет использовать её предобученные веса и адаптировать под задачу распознавания пикетных столбиков. Установка YOLOv5 в командной строке:

git clone https://github.com/ultralytics/yolov5
cd yolov5
pip install -r requirements.txt

Конфигурация модели и гиперпараметров. YOLOv5 поддерживает гибкую настройку конфигурации модели, включая выбор предобученных весов, параметры обучения (размер батча, скорость обучения), а также адаптацию аннотаций и формата изображений к входным требованиям модели. Дополнительно были заданы гиперпараметры, специфичные для условий распознавания мелких объектов (пикетных столбиков) на большом фоне.

Подготовка и разметка данных для обучения модели

Для обучения модели YOLOv5 необходимо провести предварительную обработку данных, включающая разметку объектов, адаптацию изображений их нарезку, конвертацию и сохранение патчей.

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

Разметка объектов проводится вручную, например, инструментом LabelImg.

# Запуск labelImg
!labelImg

Запускается отдельное окно для разметки пикетов, выбираем папку со всеми фотографиями (Open Dir), далее клавишу W, выделяем область объекта, сохраняем название объекта «piket», для следующих выделенных пикетов просто выбирается имя объекта из списка.

Результаты сохраняются в формате XML. В XML-файлах содержатся следующие данные:

• координаты ограничивающей рамки (bounding box) для каждого столбика.

• категория объекта — в данном случае, все объекты относятся к классу «пикетный столбик», обозначенные как piket.

Аннотации формата XML используются для обработки данных с помощью библиотек компьютерного зрения и машинного обучения, в том числе при обучении модели YOLOv5.

Пример содержимого XML — файла данной работы:

<?xml version='1.0' encoding='utf-8'?>
<annotation>
	<folder>piket</folder>
	<filename>R0025940.JPG</filename>
	<path>D:\MyProjects\Piket_detection\fotos\R0025940.JPG</path>
	<source>
		<database>Unknown</database>
	</source>
	<size>
		<width>3264</width>
		<height>4928</height>
		<depth>3</depth>
	</size>
	<segmented>0</segmented>
	<object>
		<name>piket</name>
		<pose>Unspecified</pose>
		<truncated>0</truncated>
		<difficult>0</difficult>
		<bndbox>
			<xmin>1500</xmin>
			<ymin>3187</ymin>
			<xmax>1527</xmax>
			<ymax>3213</ymax>
		</bndbox>
	</object>
	<object>
		<name>piket</name>
		<pose>Unspecified</pose>
		<truncated>0</truncated>
		<difficult>0</difficult>
		<bndbox>
			<xmin>1490</xmin>
			<ymin>3530</ymin>
			<xmax>1514</xmax>
			<ymax>3552</ymax>
		</bndbox>
	</object>
	<object>
		<name>piket</name>
		<pose>Unspecified</pose>
		<truncated>0</truncated>
		<difficult>0</difficult>
		<bndbox>
			<xmin>3223</xmin>
			<ymin>3627</ymin>
			<xmax>3253</xmax>
			<ymax>3652</ymax>
		</bndbox>
	</object>
</annotation>

Разделим по папкам аннотации и фотографии:

# Путь к папкам с фотографиями и разметками
image_folder = r"D:\MyProjects\Piket_detection\fotos"
annotation_folder = r"D:\MyProjects\Piket_detection\pikets"

После разметки пикетов для проверки правильности их расположения, были выведены несколько случайных изображений из подготовленной выборки:

# Функция для чтения координат пикетов из .xml файла
def parse_annotation(xml_file):
    tree = ET.parse(xml_file)
    root = tree.getroot()
    boxes = []
    
    for obj in root.findall('object'):
        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)
        boxes.append((xmin, ymin, xmax, ymax))
    
    return boxes

# Функция для отображения изображения с нанесенными на него пикетами
def show_image_with_boxes(image_path, boxes):
    image = cv2.imread(image_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # Для корректного отображения цветов
    
    for (xmin, ymin, xmax, ymax) in boxes:
        cv2.rectangle(image, (xmin, ymin), (xmax, ymax), (255, 0, 0), 3)  # Отрисовка красных рамок вокруг пикетов
    
    plt.figure(figsize=(10, 10))
    plt.imshow(image)
    plt.axis('off')
    plt.show()

# Получение списка всех изображений и соответствующих xml файлов
images = [f for f in os.listdir(image_folder) if f.endswith('.JPG')]
annotations = [f.replace('.JPG', '.xml') for f in images]

# Выбор 2 случайных изображения для визуализации
random_images = random.sample(images, 2)

# Визуализация двух случайных изображений с нанесенными пикетами
for image_file in random_images:
    image_path = os.path.join(image_folder, image_file)
    annotation_path = os.path.join(annotation_folder, image_file.replace('.JPG', '.xml'))
    
    if os.path.exists(annotation_path):
        boxes = parse_annotation(annotation_path)
        show_image_with_boxes(image_path, boxes)
    else:
        print(f"Разметка для изображения {image_file} не найдена.")

Предварительная обработка изображений

Изображения с разрешением 3264×4928 пикселей слишком большие для эффективного обучения модели, что могло привести к увеличению требований по памяти и замедлению процесса обучения. Поэтому были предприняты следующие шаги по предварительной обработке изображений:

1. Нарезка изображений: кадрирование изображений до размеров 1024x1024 пикселей с перекрытием (stride) 512 пикселей для выделения областей с пикетными столбиками, что позволяет уменьшить размер изображения без потери ключевых объектов;

2. Обновление координат bounding box: для каждого патча вычисляются новые координаты боксов (пикетов), если они попадают в область патча.

3. Сохранение патчей: сохраняются только те патчи, в которых есть пикеты, вместе с их разметкой в формате YOLOv5 (координаты центров и размеры боксов).

# Параметры нарезки
patch_size = 1024
stride = 512  # Шаг нарезки

# Функция для нарезки изображения на патчи
def slice_image(image, patch_size, stride):
    patches = []
    h, w, _ = image.shape
    for y in range(0, h - patch_size + 1, stride):
        for x in range(0, w - patch_size + 1, stride):
            patch = image[y:y + patch_size, x:x + patch_size]
            patches.append((patch, x, y))
    return patches

# Функция для обновления координат боксов для патча
def update_boxes_for_patch(boxes, x_offset, y_offset, patch_size):
    updated_boxes = []
    for (xmin, ymin, xmax, ymax) in boxes:
        # Переводим координаты бокса в систему координат патча
        xmin_patch = xmin - x_offset
        ymin_patch = ymin - y_offset
        xmax_patch = xmax - x_offset
        ymax_patch = ymax - y_offset
        
        # Проверяем, попадает ли бокс в границы патча
        if (0 <= xmin_patch <= patch_size and 0 <= ymin_patch <= patch_size and 
            0 <= xmax_patch <= patch_size and 0 <= ymax_patch <= patch_size):
            updated_boxes.append((xmin_patch, ymin_patch, xmax_patch, ymax_patch))
    
    return updated_boxes

# Функция для сохранения патчей и обновленной разметки
def save_patches(image_file, boxes, save_dir):
    image_path = os.path.join(image_folder, image_file)
    image = cv2.imread(image_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    
    patches = slice_image(image, patch_size, stride)
    
    for i, (patch, x_offset, y_offset) in enumerate(patches):
        patch_boxes = update_boxes_for_patch(boxes, x_offset, y_offset, patch_size)
        
        if len(patch_boxes) > 0:  # Сохраняем только патчи с пикетами
            patch_name = f"{image_file.replace('.JPG', '')}_patch_{i}.jpg"
            patch_path = os.path.join(save_dir, patch_name)
            cv2.imwrite(patch_path, cv2.cvtColor(patch, cv2.COLOR_RGB2BGR))
            
            # Также нужно сохранить соответствующую разметку для каждого патча (если пикеты в нём есть)
            annotation_file = patch_path.replace('.jpg', '.txt')  # Используем .txt для YOLOv5 формата
            with open(annotation_file, 'w') as f:
                for (xmin, ymin, xmax, ymax) in patch_boxes:
                    # Преобразование координат в формат YOLO (x_center, y_center, width, height)
                    x_center = (xmin + xmax) / 2 / patch_size
                    y_center = (ymin + ymax) / 2 / patch_size
                    width = (xmax - xmin) / patch_size
                    height = (ymax - ymin) / patch_size
                    f.write(f"0 {x_center} {y_center} {width} {height}\n")  # '0' - класс пикет

# Папка для сохранения нарезанных изображений и разметок
save_dir = r"D:\MyProjects\Piket_detection\sliced_images"
os.makedirs(save_dir, exist_ok=True)

# Запускаем процесс нарезки всех изображений и сохранения патчей
for image_file in images:
    annotation_file = os.path.join(annotation_folder, image_file.replace('.JPG', '.xml'))
    
    if os.path.exists(annotation_file):
        boxes = parse_annotation(annotation_file)
        save_patches(image_file, boxes, save_dir)
    else:
        print(f"Разметка для изображения {image_file} не найдена.")           

Каждое нарезанное изображение проверялось на наличие пикетов, и, если bounding box полностью находился в пределах фрагмента, его координаты обновлялись в соответствии с координатами патча. Эти обновленные данные сохранялись в формате, который YOLOv5 использует для обучения — координаты центра bounding box, его ширина и высота. Например, для координат пикетного столбика (xmin, ymin, xmax, ymax) в XML они были преобразованы в формат (x_center, y_center, width, height).

Проверка правильности нарезки

Визуализировать правильность нарезки фотографий, сопоставление с ними аннотаций можно так:

1. Чтение аннотаций: аннотации загружаются из файлов формата YOLOv5 и преобразуются в стандартные координаты (xmin, ymin, xmax, ymax).

2. Отображение изображений: для каждого изображения рисуются bounding box вокруг пикетов, если такие есть.

3. Визуализация: три случайных изображения из папки sliced_images визуализируются с нанесенными на них bounding box.

# Путь к папке с нарезанными изображениями и аннотациями
sliced_image_folder = r"D:\MyProjects\Piket_detection\sliced_images"
# Функция для чтения аннотаций в формате YOLOv5
def read_yolo_annotation(annotation_file, img_size):
    boxes = []
    with open(annotation_file, 'r') as f:
        lines = f.readlines()
        for line in lines:
            _, x_center, y_center, width, height = map(float, line.strip().split())
            # Преобразование координат из YOLOv5 в обычные координаты (xmin, ymin, xmax, ymax)
            x_center *= img_size
            y_center *= img_size
            width *= img_size
            height *= img_size
            xmin = int(x_center - width / 2)
            ymin = int(y_center - height / 2)
            xmax = int(x_center + width / 2)
            ymax = int(y_center + height / 2)
            boxes.append((xmin, ymin, xmax, ymax))
    return boxes

# Функция для отображения изображения с нанесенными на него пикетами
def show_sliced_image_with_boxes(image_path, annotation_path, img_size=1024):
    image = cv2.imread(image_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # Для корректного отображения цветов
    
    if os.path.exists(annotation_path):
        boxes = read_yolo_annotation(annotation_path, img_size)
        for (xmin, ymin, xmax, ymax) in boxes:
            cv2.rectangle(image, (xmin, ymin), (xmax, ymax), (255, 0, 0), 2)  # Отрисовка красных рамок вокруг пикетов
    
    plt.figure(figsize=(8, 8))
    plt.imshow(image)
    plt.axis('off')
    plt.show()

# Получаем список всех нарезанных изображений
sliced_images = [f for f in os.listdir(sliced_image_folder) if f.endswith('.jpg')]
random_images = random.sample(sliced_images, 3)

# Визуализация трех случайных изображений с нанесенными пикетами
for image_file in random_images:
    image_path = os.path.join(sliced_image_folder, image_file)
    annotation_path = image_path.replace('.jpg', '.txt')  # Путь к файлу аннотаций
    show_sliced_image_with_boxes(image_path, annotation_path)

Создание обучающей и тестовой выборок

После подготовки данных изображения были случайным образом разделены на две части: обучающую (80%) и тестовую (20%) выборки. Каждое изображение из выборок имело соответствующую аннотацию в формате YOLOv5. Папки с обучающей и тестовой выборками были структурированы в соответствии с требованиями фреймворка YOLOv5.

# Создание папок для обучающей и тестовой выборок
train_folder = r"D:\MyProjects\Piket_detection\train_data"
val_folder = r"D:\MyProjects\Piket_detection\val_data"
os.makedirs(train_folder, exist_ok=True)
os.makedirs(val_folder, exist_ok=True)

# Получение списка всех нарезанных изображений
sliced_images = [f for f in os.listdir(sliced_image_folder) if f.endswith('.jpg')]

# Процент для обучающей 80% и тестовой выборки 20%
train_ratio = 0.8

# Перемешивание данных случайным образом
random.shuffle(sliced_images)

# Разделение на обучающую и тестовую выборки
train_size = int(len(sliced_images) * train_ratio)
train_images = sliced_images[:train_size]
val_images = sliced_images[train_size:]

# Функция для копирования изображений и аннотаций
def copy_files(image_list, source_folder, dest_folder):
    for image_file in image_list:
        # Путь к изображению и аннотации
        image_path = os.path.join(source_folder, image_file)
        annotation_path = image_path.replace('.jpg', '.txt')

        # Убедимся, что копируем файлы в разные папки
        if not os.path.exists(os.path.join(dest_folder, image_file)):
            # Копируем изображение
            shutil.copy(image_path, os.path.join(dest_folder, image_file))

        # Копируем аннотацию, если она существует и не копируется в ту же папку
        if os.path.exists(annotation_path):
            dest_annotation_path = os.path.join(dest_folder, os.path.basename(annotation_path))
            if not os.path.exists(dest_annotation_path):
                shutil.copy(annotation_path, dest_annotation_path)

# Копируем обучающую выборку
copy_files(train_images, sliced_image_folder, train_folder)
print(f"Обучающая выборка: {len(train_images)} изображений")

# Копируем тестовую выборку
copy_files(val_images, sliced_image_folder, val_folder)
print(f"Тестовая выборка: {len(val_images)} изображений")

# Функция для чтения аннотации и вывода изображения с bounding box'ами
def plot_image_with_pikets(image_path, annotation_path):
    # Чтение изображения
    img = cv2.imread(image_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    
    # Чтение аннотации
    with open(annotation_path, 'r') as file:
        lines = file.readlines()
    
    # Пробегаем по всем аннотациям и рисуем bounding box'ы
    for line in lines:
        class_id, x_center, y_center, width, height = map(float, line.strip().split())
        img_h, img_w, _ = img.shape
        
        # Преобразование относительных координат в пиксели
        x_center, y_center = int(x_center * img_w), int(y_center * img_h)
        width, height = int(width * img_w), int(height * img_h)
        
        # Вычисляем координаты верхнего левого и нижнего правого углов
        xmin = int(x_center - width / 2)
        ymin = int(y_center - height / 2)
        xmax = int(x_center + width / 2)
        ymax = int(y_center + height / 2)
        
        # Рисуем bounding box
        cv2.rectangle(img, (xmin, ymin), (xmax, ymax), (255, 0, 0), 2)
    
    # Вывод изображения
    plt.figure(figsize=(10, 10))
    plt.imshow(img)
    plt.axis('off')
    plt.show()

# Функция для вывода случайного изображения из указанной папки с наложением bounding box
def plot_random_images_with_pikets(folder, num_images=2):
    image_files = [f for f in os.listdir(folder) if f.endswith('.jpg')]
    
    # Выбираем случайные изображения
    random_images = random.sample(image_files, num_images)
    
    for image_file in random_images:
        # Путь к изображению
        image_path = os.path.join(folder, image_file)
        
        # Путь к аннотации
        annotation_path = image_path.replace('.jpg', '.txt')
        
        # Используем ранее написанную функцию для отображения изображений
        if os.path.exists(annotation_path):
            plot_image_with_pikets(image_path, annotation_path)
        else:
            print(f"Аннотация для {image_file} не найдена.")

# Вывод одного случайного изображения из обучающей выборки
print("Случайные изображения из обучающей выборки:")
plot_random_images_with_pikets(train_folder, num_images=1)

# Вывод одного случайного изображения из тестовой выборки
print("Случайные изображения из тестовой выборки:")
plot_random_images_with_pikets(val_folder, num_images=1)

Обучение модели

Подготовка файла конфигурации

Перед тем, как запустить обучение модели YOLOv5 необходимо создать файл data.yaml, содержащий пути к обучающей и тестовой выборкам, а также информацию о количестве классов (в нашем случае — один класс «piket»). Пример содержимого файла:

train: D:/MyProjects/Piket_detection/train_data  # Путь к папке с обучающими изображениями
val: D:/MyProjects/Piket_detection/val_data      # Путь к папке с валидационными изображениями

nc: 1  # Количество классов (у нас один класс — пикет)
names: ['piket']  # Название класса

Активация окружения, установка зависимостей

После подготовки файла конфигурации, следующим шагом — обучение модели YOLOv5 для распознавания пикетных столбиков. Процедура запуска обучения модели на GPU (видеокарте) осуществляется через командную строку в созданном ранее окружении:  

conda activate yolov5_env

Переходим в папку со скаченным репозиторием:

cd D:\MyProjects\Piket_detection\yolov5

Устанавливаем зависимости:

pip install -r requirements.txt

Проверяем, доступна ли видеокарта для обучения (по результату True/False), в противном случае обучать на CPU, а это в десятки раз (а то и сотни) медленнее:

python -c "import torch; print(torch.cuda.is_available())"

Запуск обучения

Теперь можно запустить обучение в командной строке

python train.py --img 1024 --batch 8 --epochs 15 --data D:/MyProjects/Piket_detection/data.yaml --weights yolov5s.pt --name MY_MODEL --device 0

Параметры обучения:

--img 1024: задает размер входного изображения для модели. Здесь размер установлен на 1024 пикселя, что соответствует разрешению нарезанных изображений. Этот параметр важен, так как модели YOLO требуется фиксированный размер изображения для работы, а более крупное разрешение позволяет точнее распознавать мелкие детали на картинке, такие как пикетные столбики.

--batch 8: размер батча, то есть количество изображений, обрабатываемых моделью за один шаг обучения. В данном случае он установлен на 8. Размер батча влияет на стабильность и скорость обучения, а также на использование видеопамяти (GPU memory). Если ресурсов недостаточно, иногда его уменьшают, чтобы избежать переполнения памяти.

--epochs 15: количество эпох обучения. Количество часто подбирается эмпирически: малое количество может привести к недообучению, а слишком большое — к переобучению.

--data D:/MyProjects/Piket_detection/data.yaml: путь к файлу конфигурации данных (data.yaml). Этот YAML-файл содержит информацию о путях к тренировочной и валидационной выборкам, а также о классах, которые должна распознавать модель. Его наличие позволяет модели знать, с какими данными она работает и какие классы объектов искать.

--weights yolov5s.pt: путь к файлу с предобученными весами. Здесь используется yolov5s.pt, который является легковесной версией модели YOLOv5. Предобученные веса служат базой для дообучения, ускоряя обучение и улучшая точность, поскольку модель уже имеет общее представление о типах объектов, и ей нужно адаптироваться только к специфике пикетных столбиков.

--name MY_MODEL: задает имя для сессии обучения. Результаты обучения (включая веса и логи) будут сохранены в папке с этим именем, что облегчает организацию и последующее использование файлов, особенно если проводится несколько экспериментов.

--device 0: указывает, на каком устройстве будет выполняться обучение. Здесь 0 означает использование первого доступного GPU. Это помогает ускорить процесс, так как обучение модели на GPU значительно быстрее, чем на CPU. Параметр «--device cpu» запустит обучение на ЦП.

Результаты обучения модели

На предоставленных графиках можно проанализировать результаты обучения модели YOLOv5 для детекции пикетов.

После завершения обучения модели, результаты будут сохранены в папке yolov5\runs\train\ в нашем случае путь выглядит так: «D:\MyProjects\Piket_detection\yolov5\runs\train\MY_MODEL».

Файл results

В табличной форме в виде картинки

Пояснение результатов:

train/box_loss и val/box_loss: показатели ошибок (loss) на обучающем и валидационном наборах. Обе кривые показывают, что ошибки (loss) уменьшаются по мере обучения, что указывает на то, что модель становится лучше в нахождении bounding boxes. Валидационная ошибка (val/box_loss) немного колеблется в начале, но в конце она стабилизируется, что хорошо.

train/obj_loss и val/obj_loss: ошибки объектной функции, которые измеряют уверенность модели в наличии объектов в кадре. Видно, что эти ошибки также уменьшаются, что подтверждает улучшение в уверенности модели.

train/cls_loss и val/cls_loss: ошибки классификации (cls_loss) равны нулю. Это может быть связано с тем, что в данном случае только один класс для детекции — “piket”, что исключает необходимость в многоклассовой классификации.

metrics/precision и metrics/recall: метрики точности (precision) и полноты (recall) увеличиваются и достигают высоких значений ближе к завершению обучения, что говорит о том, что модель успешно находит и классифицирует пикеты.

metrics/mAP_0.5 и metrics/mAP_0.5:0.95: метрики средней точности (Mean Average Precision) на разных уровнях IoU (intersection over union). Видно, что mAP на уровне 0.5 достаточно высок (около 0.9), а при более строгом mAP 0.5:0.95 также увеличивается до ~0.6. Это хороший результат, показывающий, что модель точно детектирует пикеты.

Модель хорошо обучена и показывает стабильные результаты на обучающей и валидационной выборках.

Метрики mAP показывают, что модель имеет достаточно высокую точность в детекции пикетов, хотя есть возможности для дальнейшего улучшения на более строгих уровнях IoU.

Файл F1_curve

График F1-Confidence Curve показывает соотношение между метрикой F1 и уверенностью модели. Видно, что при уровне уверенности около 0.6 модель демонстрирует наилучшее значение F1 (приблизительно 0.99). Это говорит о том, что модель эффективно сбалансировала precision и recall, если порог доверия установлен на 0.6.

Для дальнейшего использования модели можно рассмотреть установление порога уверенности около 0.6, что максимизирует баланс между precision и recall.

Файл P_curve

График Precision-Confidence показывает, как изменяется точность при различных уровнях уверенности. График демонстрирует высокую точность, достигающую максимума около уровня уверенности 0.89.

Файл R_curve

График Recall-Confidence иллюстрирует, как изменяется полнота модели в зависимости от уровня уверенности. Полнота держится на уровне около 1.0 до уровня уверенности примерно 0.8, что указывает на уверенное обнаружение большинства объектов при высоких значениях уверенности.

Файл PR_curve

График Precision-Recall демонстрирует соотношение между точностью и полнотой для модели, показывая высокие значения обеих метрик (примерно 0.992), что свидетельствует об очень хорошей способности модели к распознаванию пикетов.

Вывод по графикам F1_curve, P_curve , R_curve, PR_curve:

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

Файл confusion_matrix

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

Файл labels

Гистограмма меток (labels.jpg) показывает, что все изображения относятся к классу «пикет». Также можно увидеть распределения меток по осям x, y, ширине и высоте объектов. Это позволяет оценить характер распределения пикетов на изображениях и их размеры.

Файл labels_correlogram

На изображении «Correlogram меток» представлена корреляционная диаграмма для меток (labels_correlogram.jpg). Здесь видны распределения координат (x, y) и размеров (width, height) объектов на изображениях. Существует явная корреляция между width и height, что говорит о том, что пикеты имеют пропорциональную форму. В координатах x и y также можно заметить ряды, что может свидетельствовать о регулярной сетке расположения пикетов на изображениях.

Примеры батчей обучающей выборки

Примеры батчей валидационной выборки

Общий вывод по результатам обучения модели YOLO для распознавания пикетов

На основе представленных изображений можно сделать следующие выводы:

1. Распределение меток и параметров объектов:

Гистограммы и корреляционная диаграмма меток показывают, что пикеты на изображениях распределены достаточно регулярно и имеют стабильные размеры. Это подтверждается корреляцией между шириной и высотой объектов, что указывает на их однородную форму и размер.

Координаты объектов по x и y также распределены относительно равномерно, с периодическими всплесками, что может свидетельствовать о регулярном расположении пикетов на изображениях.

2. Классовый дисбаланс:

На гистограмме видно, что все объекты на изображениях относятся к одному классу – "piket". Это упрощает задачу для модели, так как отсутствует необходимость различать между разными объектами. Однако это также делает модель менее гибкой в плане распознавания других классов объектов в будущем.

3. Качество классификации:

Матрица ошибок (confusion matrix) указывает на высокую точность модели. Все пикеты распознаны правильно, и отсутствуют ошибки классификации. Это говорит о том, что модель успешно научилась распознавать пикеты и минимизирует как ложные срабатывания, так и пропуски.

 Заключение по результатам обучения:

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

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

Применение модели для распознавания на новых данных

Теперь готовую обученную модель можно применить (протестировать) для распознавания пикетов на новых изображениях, которые модель ранее не «видела» ни на обучающем, ни на валидационном наборе данных.

Ниже код, который выполняет задачу детекции объектов (пикетов) на изображении с использованием обученной модели YOLOv5, состоящий из основных шагов:

1. Загрузка модели: кастомная модель YOLOv5, обученная на задаче распознавания пикетов, с использованием torch.hub.load берется из локального файла best.pt.

2. Функция detect_image загружает изображение из указанного пути и уменьшает его до заданного масштаба scale для удобного отображения результатов. Далее выполняется детекция объектов с помощью модели на оригинальном изображении (без уменьшения). Затем извлекается координаты обнаруженных объектов, их класс и уровень уверенности модели в предсказании. Если объекты найдены, они рисуются на уменьшенном изображении с учетом масштабирования. Каждый обнаруженный объект обозначается прямоугольником (bounding box), а рядом указывается его класс и уверенность.

3. Отображение результата: показывает изображение с обнаруженными объектами в новом окне.

4. Запуск функции: вызывается функция detect_image с заданным изображением для детекции пикета.

# Загрузка модели
model = torch.hub.load('ultralytics/yolov5', 'custom', path='D:/MyProjects/Piket_detection/yolov5/runs/train/MY_MODEL/weights/best.pt', force_reload=True)

# Функция для детекции на одном изображении с выводом результатов
def detect_image(image_path, scale=0.5):
    # Загрузка изображения
    img = cv2.imread(image_path)
    original_height, original_width = img.shape[:2]  # исходные размеры изображения

    # Уменьшение изображения
    img_resized = cv2.resize(img, None, fx=scale, fy=scale, interpolation=cv2.INTER_AREA)
    resized_height, resized_width = img_resized.shape[:2]  # новые размеры изображения
    
    # Запуск детекции на оригинальном изображении (без уменьшения)
    results = model(img)

    # Извлечение результатов
    detections = results.xyxy[0]  # координаты bbox в формате [xmin, ymin, xmax, ymax, confidence, class]
    
    # Проверка того, что нашла модель
    print(f'Найдено объектов: {len(detections)}')
    if len(detections) == 0:
        print('Модель не обнаружила объектов на изображении.')

    # Отображение результатов с коррекцией масштаба для уменьшенного изображения
    for *box, conf, cls in detections:
        # Масштабируем координаты в соответствии с уменьшенным изображением
        box = [coord * (resized_width / original_width if i % 2 == 0 else resized_height / original_height) for i, coord in enumerate(box)]
        
        label = f'{model.names[int(cls)]} {conf:.2f}'
        cv2.rectangle(img_resized, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), (255, 0, 0), 2)
        cv2.putText(img_resized, label, (int(box[0]), int(box[1]) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2)

    # Показать изображение
    cv2.imshow('Detection', img_resized)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

# Запуск детекции пикета на фото
detect_image('D:/MyProjects/Piket_detection/New_fotos_test/1.jpg', scale=0.7)

Результаты детекции пикета на новом фото:

Направления дальнейшего развития работы

На основе нейронных сетей, которые могут распознавать различные объекты по фотографиям с БПЛА, оснащенного навигационным оборудованием геодезического класса точности с привязкой к базовым ГНСС-станциям, возможно дальнейшее развитие в виде автоматизированного мониторинга железнодорожной инфраструктуры, с определением положения пикетных столбиков и других распознаваемых объектов, в т.ч. определять координаты центров объектов с геодезической точностью (деци- и сантиметровой) в местной, локальной или другой системе координат. Такой подход на выходе может частично автоматизировать отрисовку топографического плана (в зависимости от количества распознаваемых объектов). Это поможет улучшить точность и скорость сбора данных о состоянии зданий, пикетов, верхнего строения пути и других элементов инфраструктуры и прилегающей территории железнодорожного пути.