Использование Detectron2 для сегментации экземпляров в наборе данных Spacenet

Следуйте инструкциям в этой статье, чтобы узнать, как реализовать Detectron2 от начала до конца, используя набор данных SpaceNet.

ПРИМЕЧАНИЕ. Эта статья представляет собой небольшое введение в обнаружение объектов, поэтому, если вы уже установили Detectron2 и его зависимости, перейдите к разделу «Использование Detectron2».

Введение

Detectron2 — это платформа для обнаружения объектов, выпущенная в 2019 году командой Facebook AI Research. Основанный на платформе машинного обучения PyTorch, Detectron2 может обнаруживать объекты с помощью семантической сегментации, сегментации экземпляров и паноптической сегментации. В этой статье основное внимание будет уделено использованию сегментации экземпляров для обнаружения и выделения домов на спутниковых снимках, например, на изображении выше.

Мы будем использовать две части общедоступного набора данных AWS SpaceNet, уделяя основное внимание SpaceNet 1: Building Detection v1 и SpaceNet 2: Building Detection v2. Эти наборы данных содержат спутниковые изображения из регионов по всему миру: SpaceNet 1 содержит около 9 700 изображений из Рио-де-Жанейро, а SpaceNet 2 содержит более 14 100 изображений из комбинации Лас-Вегаса, Парижа, Шанхая и Хартума. Все изображения снабжены аннотациями, очерчивающими каждое здание многоугольником, который мы будем использовать в качестве основы во время обучения.

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

Установка Detectron2 и зависимостей

Первый шаг — убедиться, что ваше оборудование и операционная система хотя бы частично совместимы с инструкциями в этом руководстве. Я использую копию Ubuntu 20.04, которая больше подходит для такого типа работы, чем Mac OS или Windows. Если вы используете другую операционную систему, я рекомендую получить Ubuntu, выполнив двойную загрузку или запустив виртуальную машину на вашем компьютере.

Настольный компьютер, который я использовал для этой демонстрации, имеет Nvidia GeForce GTX 1060, которого, как мне показалось, более чем достаточно для написания небольших моделей. Я использовал CUDA 11.4 с версией драйвера Nvidia 470.103.01. Я бы порекомендовал что-то близкое к этому, но наличие машины с поддержкой CUDA как минимум с CUDA 11.0 позволит вам выполнить все шаги. Вы можете проверить свой CUDA и драйвер с помощью nvidia-smi. Если вам нужна помощь в решении проблемы с установкой драйверов и CUDA, нажмите здесь для получения официальной документации Nvidia.

Чтобы сохранить ваш код и зависимости в чистоте, вы, вероятно, захотите настроить либо виртуальную среду Python, либо среду Conda. Я создал среду Python с помощью команды virtualenv, которую вы можете получить, выполнив следующие строки:

sudo apt-get update && \
sudo apt-get upgrade && \
sudo apt install virtualenv

Чтобы создать виртуальную среду, выполните приведенные ниже команды. Я создам свой в домашнем каталоге ~ и назову его detectron.

virtualenv --python=/usr/bin/python3.8 ~/detectron &&
cd ~

Среда должна быть активирована, чтобы вы могли устанавливать в нее пакеты. Вы можете сделать это, выполнив команду:

source ~/detectron/bin/activate

Чтобы деактивировать среду, запустите:

deactivate

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

Linux или macOS с Python ≥ 3.6
PyTorch ≥ 1.8 и соответствующая версия torchvision
Последняя версия OpenCV

Что еще более важно, вот комбинация пакетов Python, PyTorch и torchvision, которые я установил для успешного использования предварительно собранного Detectron2:

Питон 3.8.10
PyTorch=1.10.2+cu113 и torchvision=0.11.3+cu113. Установите, перейдя по этой ссылке, или воспользуйтесь командой ниже. Возможно, вам придется удалить текущие версии этих пакетов перед установкой новых.

pip3 install \
torch==1.10.2 \
torchvision==0.11.3 \
--extra-index-url https://download.pytorch.org/whl/cu113

Последняя версия OpenCV

pip install opencv-python

Интегрированные среды разработки (IDE) очень полезны при разработке и тестировании программного обеспечения. Я предпочитаю использовать JupyterLab из-за его легкости и легкости расширения, но не стесняйтесь использовать свой любимый инструмент. Если вы хотите использовать JupyterLab, выполните следующие команды:

pip install jupyter jupyterlab

Чтобы открыть JupyterLab, просто запустите jupyter-lab и скопируйте/вставьте ссылку, которую он предоставляет, в выбранный вами браузер.

Наконец, мы можем установить сам Detectron2. Если вы используете Linux и рекомендуемые драйверы и пакеты, как указано выше, вы можете запустить следующую команду:

python -m pip install detectron2 -f \
https://dl.fbaipublicfiles.com/detectron2/wheels/cu113/torch1.10/index.html

В противном случае следуйте инструкциям по установке по этой ссылке.

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

pip install \
setuptools==59.5.0 \
numpy \
pandas \
geojson \
matplotlib \
scikit-learn

Наборы данных SpaceNet состоят из изображений TIFF с цветными растрами и геопространственными метаданными. Однако Detectron2 не поддерживает этот тип изображений, поэтому сначала нам придется преобразовать изображения в формат PNG. Для изображений TIFF требуется, чтобы пакет GDAL открывался и обрабатывался, а это значит, что нам также необходимо установить его. Установка GDAL может быть неприятной, но следующие команды — лучший способ, который мы нашли до сих пор:

sudo add-apt-repository ppa:ubuntugis/ppa && \
sudo apt-get update && \
sudo apt-get install build-essential && \
sudo apt-get install gdal-bin libgdal-dev && \
export CPLUS_INCLUDE_PATH=/usr/include/gdal && \
export C_INCLUDE_PATH=/usr/include/gdal && \
pip install \
--global-option=build_ext \
--global-option="-I/usr/include/gdal" \
GDAL==`gdal-config --version`

Чтобы убедиться, что все установлено правильно, запустите новый файл Jupyter Notebook или Python и импортируйте следующие модули. Если все импортируется правильно, вы можете начать использовать Detectron2 в следующем разделе. Если возникают какие-либо ошибки, проверьте все версии, которые вы установили с помощью pip list -o и python --version.

import detectron2         # detectron2
import torch              # pytorch
import cv2 as cv          # openCV
import numpy as np        # numpy
from osgeo import gdal    # GDAL
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg

Настроить АВС

Вам нужно будет установить Amazon Web Service (AWS) и создать учетную запись, чтобы загружать наборы данных SpaceNet. Чтобы установить пакет AWS, выполните следующие команды:

cd ~ && \
curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "awscliv2.zip" && \
unzip awscliv2.zip && \
sudo ./aws/install && \
rm awscliv2.zip

Теперь, когда AWS установлен в вашем домашнем каталоге, запуск команды ls из дома должен отобразить каталог aws/. Вернитесь в каталог, в котором вы создали виртуальную среду, для следующего шага. Если у вас еще нет учетной записи AWS, вы можете создать ее здесь. Как только ваша учетная запись будет создана, запустите команду aws configure в своем терминале. Вам будет предложено ввести четыре учетных данных: идентификатор ключа доступа AWS, секретный ключ доступа AWS, имя региона по умолчанию и формат вывода по умолчанию. Выполните следующие действия, чтобы создать или получить доступ к четырем учетным данным:

Ключ доступа AWS и секретный ключ доступа AWS: имя пользователя в правом верхнем углу › Учетные данные безопасности › Ключи доступа › Создать новые ключи доступа.
Имя региона по умолчанию: вернитесь на главный экран и щелкните вкладку региона рядом с вкладкой имени пользователя. Он будет говорить что-то другое в зависимости от того, где вы находитесь. Название региона должно быть первым в списке; мой говорит «нас-восток-1».
Выходной формат по умолчанию: введите json.

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

Подготовьте набор данных SpaceNet

Наборы данных SpaceNet с открытым исходным кодом можно найти по этой ссылке. В то время как данные в SpaceNet 2 имеют единую структуру, данные Рио-де-Жанейро, содержащиеся в SpaceNet 1, организованы по-другому. Изображения Рио поставляются с .geojson файлами, в которых очерчены углы многоугольников каждого здания. Проблема в том, что местоположения точек многоугольника задаются в глобальных координатах, где нам нужны только местоположения пикселей относительно изображения. Набор данных SpaceNet 2 поставляется с .csv файлами, которые уже содержат готовые к использованию пиксельные координаты многоугольника для каждого изображения. Вы можете упорядочить данные самостоятельно или использовать уже сделанные мной скрипты. Чтобы использовать мои скрипты, клонируйте мой репозиторий:

# HTTPS
git clone https://github.com/rl02898/detectron2-spacenet.git
# SSH
git clone [email protected]:rl02898/detectron2-spacenet.git

После клонирования перейдите в репозиторий и запустите установочный файл:

cd detectron2-spacenet && \
source install.sh

Это может занять некоторое время, до нескольких часов, в зависимости от скорости вашего интернета и вычислительной мощности. Имейте в виду, что общий размер файлов обоих наборов данных SpaceNet составляет около 85 ГБ. Скрипт install.sh сделает следующее:

Загрузите набор данных.
Организуйте их в папки.
Создайте каталоги train/val/test.
Преобразуйте изображения в PNG.
Создайте файлы JSON для аннотаций. Они будут в тех же каталогах, что и изображения.

После завершения этого процесса у вас должна быть такая файловая структура:

Spacenet/
├─ AOInvidia-smiRio_train
...
...
...
├─AOIjupyter-labKhartoum_Test_public
├─ train/
│  ├─ images.png
│  ├─ via_region_data.json
├─ val/
│  ├─ images.png
│  ├─ via_region_data.json
├─ test/
│  ├─ images.png

Поздравляем, теперь вы готовы начать использовать Detectron2!

Использование Дететрон2

В документации здесь есть официальное руководство по ноутбуку Google Colab, но я остановлюсь на нем подробнее. Начните свой скрипт Python с импорта:

Далее нам нужно сделать то, что Detectron2 называет «зарегистрировать ваш набор данных». По сути, мы должны сообщить Detectron2, где находится наш набор данных и где искать информацию. Для этого нам нужно создать функцию, которая может считывать наши файлы JSON ранее. Я основывал свою функцию на той, что была в учебнике.

Зарегистрируйте набор данных:

Мы определяем классы и цвета, которые для нашего набора данных только одни. Если ваш набор данных содержит больше классов, просто добавьте их в список и добавьте больше цветов в формате RGB. Этот раздел будет работать, только если вы создали файлы JSON в правильном формате. Если есть какие-либо проблемы, проверьте файлы JSON или используйте функции из моего репозитория.

Нанесение основных истин:

Этот фрагмент кода будет определять метаданные и считывать информацию с помощью созданной нами функции get_dataset_dicts(). Затем он выберет три случайных изображения из нашего тренировочного набора и построит изображения вместе с экземплярами истинной наземной информации. Параметр instance_mode=ColorMode.IMAGE определяет, хотим ли мы, чтобы каждый класс был разного цвета или одного и того же цвета. Мы обсудим это позже с другими параметрами и тем, как их настроить.

Настройка конфигураций модели:

Следующее, что нужно сделать, это настроить параметры. Базовая конфигурация, предоставляемая Detectron2, представлена в формате YAML, что упрощает их экспорт и сохранение. Это именно то, что мы будем делать. Чтобы увидеть все доступные параметры, вы можете запустить print(get_cfg()).

Некоторые дополнительные комментарии и пояснения к приведенному выше коду:

Для cfg.merge_from_file() мы используем готовую модель, предоставленную Detectron2. Список моделей, которые вы можете использовать, представлен здесь. Просто замените mask_rcnn_X_101_32x8d_FPN_3x.yaml на модель, которая вам подходит. Большие модели будут работать лучше, но также будут более дорогими в вычислительном отношении и будут медленнее делать выводы.
Для cfg.MODEL.WEIGHTS мы используем предварительно обученную модель в качестве контрольной точки. Если вы хотите инициализировать случайные веса, просто используйте пустую строку "" .
Настройте остальные параметры в соответствии с вашими потребностями и возможностями графического процессора.

Теперь мы готовы к обучению.

Обучение работе с одним GPU:

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

Результаты будут сохранены в ваш аргумент cfg.OUTPUT_DIR.

Обучение работе с несколькими графическими процессорами:

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

В моем репозитории, наряду с Detectron2, есть файл с именем train_net.py. Мы добавляем следующие строки кода в начало функции main():

Наряду с этим вам также необходимо переопределить функцию get_dataset_dicts() в файле, желательно вверху. Если вы воспользуетесь моим скриптом, он уже будет содержать эти изменения. После внесения этих изменений мы можем обучить нашу модель. Ниже приведен пример обучения Detectron2 на 4 GPU с помощью командной строки.

python train_net.py \
--num-gpus 4 \
--config-file Spacenet/SpacenetD2cfg.yaml

Независимо от того, тренируетесь ли вы на одном или четырех графических процессорах, ваши результаты будут сохранены в месте вывода. Результаты будут включать журналы, конфигурации, метрики и контрольные точки модели. Detectron2 не сохраняет лучшую контрольную точку автоматически, поэтому вам придется отслеживать метрики и выбирать ту, которая лучше всего соответствует вашим потребностям. Я обычно брал окончательную модель для простоты. Контрольные точки Detectron2 сохраняются в виде .pth файлов. Далее мы покажем, как использовать эту контрольную точку для вывода изображений и видео.

Вывод на изображениях:

Делать выводы по изображениям проще всего при использовании сценариев, предоставляемых Detectron2. Для этого я предоставил два необходимых скрипта в каталоге inference/ моего репозитория. Перед выводом убедитесь, что файл конфигурации YAML содержит вашу новую модель в аргументе cfg.MODEL_WEIGHTS. В конце файла detectron_demo.py я изменил аргумент о том, куда выводить файл модели. Если вы сделали то же самое, это должно выглядеть примерно так:

Сделайте вывод, запустив файл demo.py и используя следующие аргументы в командной строке. Входными данными могут быть одно изображение (image.png) или список изображений (*.png). На выходе может быть одно изображение или каталог. Порог достоверности отсекает прогнозы ниже порога. Аргумент файла конфигурации — это расположение вашего файла конфигурации с вашей новой моделью, сохраненной под аргументом cfg.MODEL_WEIGHTS.

python inference/demo.py \
--config-file Spacenet/SpacenetD2cfg.yaml \
--input Spacenet/test/*.png \
--output Spacenet/predictions \
--confidence-threshold 0.5

Вывод на видео:

Выводы на видео делаются так же, как и на изображениях. Будет работать следующая команда:

python inference/demo.py \
--config-file Spacenet/SpacenetD2cfg.yaml \
--video-input example.mp4 \
--output example_output.mp4 \
--confidence-threshold 0.5

Некоторые распространенные проблемы

Проблема 1:

При выводе изображений с помощью скрипта demo.py каждая сегментация экземпляра будет иметь случайный цвет. Если вы хотите, чтобы каждый класс использовал определенный вами цвет, выполните следующие действия. Другими словами, если вы хотите изменить это…

в это…

и то, в это…

затем выполните следующие действия:

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

Указание instance_mode=ColorMode.SEGMENTATION устанавливает для каждого класса свой собственный цвет. Когда instance_mode=ColorMode.IMAGE цвета класса случайны. Если вы хотите избежать небольших вариаций цветов сегментации (как показано на втором рисунке), выполните следующие несколько шагов:

Найдите файл visualizer.py в исходном коде. Это будут сайты-пакеты из вашей виртуальной среды. Если вы внимательно прочитали эту статью, вы найдете файл по следующему пути: ~/detectron/lib/python3.8/site-packages/detectron2/utils/...
Мы используем функцию draw_instance_predictions(), поэтому найдите эту функцию в файле visualizer.py. Он будет находиться в классе Visualizer.
В строке, начинающейся с if self._instance_mode == ColorMode.SEGMENTATION, удалите часть self._jitter(). Для меня это строка 409.
Я также рекомендую изменить значение альфа с 0,8 до 0,5, чтобы сегментация была более прозрачной, но отрегулируйте это значение по своему усмотрению.
Если вы хотите, чтобы наземные изображения не имели джиттера, удалите часть self.jitter() из функции draw_dataset_dict() в файле visualizer.py. Строка, которую вы ищете, начинается с if self._instance_mode == ColorMode.SEGMENTATION. Для меня это была линия 571.

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

Проблема 2:

Когда вы начинаете обучать свою модель, мы несколько раз сталкивались с проблемой с модулем setuptools. Получаем такую ошибку:

AttributeError: модуль «distutils» не имеет атрибута «версия»

Мы смогли исправить эту ошибку, понизив нашу версию до 59.5.0. Это решение было найдено на странице StackOverflow здесь. Выполните следующую команду:

pip uninstall setuptools && \
pip install setuptools==59.5.0

Проблема 3:

При запуске файла inference/demo.py в новой среде и новом пользователе мы столкнулись с этой ошибкой:

ImportError: нет модуля с именем «Tkinter»

Мы не уверены, является ли это проблемой virtualenv или проблемой Ubuntu, но если вы столкнетесь с этой проблемой, вы можете исправить ее, установив модуль с помощью:

sudo apt-get update && \
sudo apt-get install python3-tk

Заключение

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

Если у вас есть какие-либо вопросы по этой статье, пожалуйста, оставьте комментарий ниже, и мы будем рады ответить вам.

-Рассел Лэнд, Джастин Холл, Майкл Фокс