Привет, Хабр! С вами Кирилл Иванов, участник профессионального сообщества NTA. Модели машинного обучения компьютерного зрения стали крайне актуальной задачей в современном мире. Компьютерные системы, способные «видеть», могут применяться во многих областях жизни. Одна из самых популярных областей применения моделей компьютерного зрения — распознавание объектов на изображениях и видео. Это полезно, к примеру, для систем видеонаблюдения, автоматической сортировки на производстве, диагностирования медицинских изображений. Также модели машинного обучения используются в дополненной и виртуальной реальности. Они позволяют создавать интерактивные пользовательские интерфейсы, а также визуализируют информацию на основе видео и изображений.

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

Как же создать такую модель?

Создание модели машинного обучения состоит из нескольких основных этапов:

Новички скажут, что самый сложный этап — это обучение модели. И, конечно же, ошибутся, на то они и новички. Любой data scientist скажет, что самый сложный и трудоёмкий этап — это сбор и подготовка данных. Зачастую на это уходит до 80 % времени. Для высокой точности работы модели необходимо не только собрать огромное количество изображений, но и правильно их разметить. Сбор и подготовка данных для классификации и распознавания образов подразумевает выделение частей изображения и добавление текстовой метки.

Задача классификации позволяет по заданному изображению определить, к классу с какой текстовой меткой можно отнести данное изображений. Задача распознавания образов, наоборот, по заданной метке класса пытается найти изображение или его фрагмент. Для решения подобных задач «вручную» можно использовать Paint. Он позволяет просматривать изображения и выписывать координаты пикселей, выделять объекты прямоугольниками и многоугольниками, сохранять изменённые изображения.

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

ImgLab — веб‑приложение, позволяющее загружать папки с изображениями, добавлять прямоугольники, многоугольники, точки и окружности.

Интерфейс приложения. Недостатки: использование веб-приложений недопустимо при работе с изображениями, не предназначенными для размещения в интернете.

LabelImg — кроссплатформенное приложение. Позволяет выбирать папки с файлами, выделять объекты прямоугольниками и добавлять метки к изображениям.

Оба приложения имеют схожие инструменты, но не позволяют легко расширить возможности автоматизированной разметки графических данных для задач машинного обучения. Кроме того, у меня на работе эти инструменты недоступны в защищённом контуре. А мне нужно было повысить качество и скорость разметки графических данных, не предназначенных для размещения в сети. Пришлось разработать лёгкое в использовании и расширяемое c помощью дополнительных модулей приложение, которое позволит:

1. Обрабатывать последовательно изображения из выбранной директории.

2. Выделять объекты прямоугольниками и многоугольниками.

3. Добавлять текстовые метки к элементам изображения.

Создание приложения

Я воспользовался Tkinter — библиотекой для кроссплатформенных приложений с графическим интерфейсом. Она входит в стандартную поставку Python.

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

def walk_by_directory (self, root_dir_path):
    for dirpath, dirnames, filenames in os.walk(root_dir_path):
        for file in filenames:
            if any([file.endswith(end) for end in ('jpg', 'jpeg', 'png', 'svg', 'bmp')]):
                yield os.path.join(root_dir_path, dirpath, file)

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

def choose_directory(self):
    directory = filedialog.askdirectory(title='Открыть папку')
    return directory

Перемещение по списку файлов в обе стороны:

def get_next_image_handler(self):
    if len(self.list_path_files) > 0 and self.now_num_file_generator + 1 < len(self.list_path_files):
         self.now_num_file_generator += 1
         img_path = self.list_path_files[self.now_num_file_generator]
         self.selected_file = img_path
         self.load_image_handler()

Проверяю номер файла в сформированном списке файлов директории. Если есть ещё файлы, то и добавляю в форму следующее изображение:

def get_before_image_handler(self):
    if len(self.list_path_files) > 0 and self.now_num_file_generator > 0 and self.now_num_file_generator < len(self.list_path_files):
         self.now_num_file_generator -= 1
         self.selected_file = self.list_path_files[self.now_num_file_generator]
         self.load_image_handler()

Таким же образом работает и получение предыдущего изображения.

Рассмотрим добавление в приложение картинки и рисование фигуры поверх него. Для работы с изображениями в Python используется библиотека Pillow. С её помощью можно открывать, обрабатывать и сохранять изображения. Отрисовывать картинку будем с помощью Canvas — его можно представить как холст, на который я приклею изображение и поверх него нанесу фигуры, чтобы выделить объект.

Добавлю изображение в графический интерфейс.

def draw_image(self, new_image_path):
    self.im = Image.open(new_image_path)
    self.im2 = ImageTk.PhotoImage(self.im)
    self.imgtag = self.canv.create_image(0,0,anchor="nw",image=self.im2)

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

def create_rectangle(self, x0, y0):
    if self.first_point is not None:
        self.canvas.canv.create_rectangle(self.first_point[0],
            self.first_point[1],
            x0,
            y0,
            width=3,
            outline="#004D40",
            tags="rectangle")
        self.first_point = None
    else:
        self.canvas.canv.delete("rectangle")
        self.first_point = x0, y0

Сначала я проверяю, стоит ли уже первая точка. Если да, то получаю координаты второй точки и добавляю прямоугольник. Если первой точки не было, тогда запоминаю координаты первой вершины и удаляю существующий прямоугольник (если он был создан ранее).

Многоугольник создаётся похожим образом:

def create_polygon(self, x0, y0):
    if len(self.points) > 0 and self.full_figure_flag:
        self.canvas.canv.delete("many_pointes")
        self.full_figure_flag = False
        self.points = []
    elif len(self.points) > 0:
        self.canvas.canv.create_line(*self.points[-1], x0, y0, 
            fill="red", width=2, tags="many_pointes")
    if len(self.points) > 2 and self.points[0][0] - x0 < 5 and self.points[0][1] - y0 < 5:
        self.full_figure_flag = True
    else:
        self.points.append((x0, y0))

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

Далее проверяю, создана ли фигура. Самая простая — треугольник, поэтому проверяю, что в списке точек не менее двух записей. Буду считать фигуру законченной, если новая точка находится не дальше 5 пикселей по каждой оси от первой точки. Если новая точка не является финальной, то добавляю её координаты в список точек.

После добавления к изображению необходимых меток нужно сохранить результат:

def save_res(self, file_path, label):
    row = {'path': file_path, 'label': label}
    self.data = self.data.append(row, ignore_index=True)

Я использую dataframe Pandas для хранения табличных данных. Передаю путь к файлу и метки, которые были добавлены — координаты фигуры или текстовую метку.

Все необходимые функции реализованы, осталось протестировать приложение.

Готовое приложение

Выберу папку с данными:

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

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

Выделю муравья на изображении прямоугольником:

Обведу человечка многоугольником:

Посмотрю, как выглядит файл с метками после обработки папки с данными:

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

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