Обнаружение объектов всегда является актуальной темой в компьютерном зрении и применяется во многих областях, таких как безопасность, наблюдение, системы автономных транспортных средств и осмотр машин. Широко используемые алгоритмы обнаружения объектов представляют собой либо алгоритмы обнаружения на основе области (Faster R-CNN, R-FCN, FPN), либо алгоритмы однократного обнаружения (SSD и YOLO).
Модель детектирования на основе региона или двухэтапного детектирования состоит из двух этапов:
- Предложение региона
- Классификация этих регионов и уточнение прогноза местоположения.
Одноразовое обнаружение пропускает этап предложения региона и сразу дает окончательную локализацию и прогноз содержания. Более быстрый RCNN более популярен в детекторах на основе регионов. Теперь мы увидим, как реализовать собственный детектор объектов с помощью Faster RCNN с PyTorch.
В рамках этой статьи мы будем обнаруживать лица на изображении. На одном изображении может быть несколько лиц. Я получил изображения из Google и изменил их размер до 512 X 512. А для аннотации я использовал инструмент labelImg. Файл аннотации - это отдельный файл CSV, содержащий image_id и координаты ограничивающего прямоугольника. Давайте импортируем необходимые библиотеки.
import pandas as pd import numpy as np import cv2 import torch import torchvision from torchvision.models.detection.faster_rcnn import FastRCNNPredictor from torchvision.models.detection import FasterRCNN from torch.utils.data import DataLoader, Dataset
Затем мы можем определить набор данных для обучения. Для этого мы можем унаследовать класс PyTorch Dataset и создать наш собственный класс TrainDataset. Согласно документации PyTorch, наш класс должен реализовывать методы __len__ и __getitem__.
Для нашего класса Dataset мы получаем файл аннотации в качестве входных данных. Все наши обучающие изображения находятся в папке «лицо». Внутри конструктора мы конвертируем файл аннотации в фрейм данных и получаем все уникальные image_ids для дальнейшей обработки. В методе __getitem__ мы можем прочитать изображение, используя image_id, который есть в кадре данных, а также мы можем получить все ограничивающие прямоугольники, связанные с этим изображением.
Затем мы инициализируем dict с именем target, который будет передан модели для обучения. Эта цель будет иметь метаданные аннотации, такие как фактические координаты ограничивающей рамки, соответствующие метки, image_id, область ограничивающих рамок. Параметр площади используется во время оценки с помощью метрики COCO, чтобы разделить метрические оценки между маленькими, средними и большими ячейками. Если мы установим iscrowd как True, эти экземпляры будут игнорироваться во время оценки. Метод __len__ даст размер набора данных.
Следующее, что нужно сделать, это установить загрузчик данных, который будет загружать обучающие данные пакетами в модель для обучения. Для этого мы также будем использовать утилиту PyTorch DataLoader.
# Initialize Dataset
train_dataset = TrainDataset('face/face_annotation.csv')
def collate_fn(batch):
return tuple(zip(*batch))
train_data_loader = DataLoader(
train_dataset,
batch_size=2,
shuffle=True,
num_workers=2,
collate_fn=collate_fn
)
В DataLoader мы можем передать инициализированный train_dataset. batch_size и num_workers будут произвольными числами, основанными на доступной памяти, с которой мы можем поэкспериментировать.
Если у нас есть графический процессор, мы можем установить устройство как cuda, иначе как CPU.
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
Теперь загрузим некоторые обучающие данные и посмотрим.
мы можем запустить приведенный выше фрагмент и построить изображения с его ограничивающими рамками. Вот несколько образцов.
Теперь инициализируем нашу модель. Я собираюсь использовать FasterRCNN torchvision с магистралью resnet50.
Мы установили pretrained как true, поэтому функция вернет модель, предварительно обученную в ImageNet. Здесь мы устанавливаем num_classes равным 2, считая фон одним классом.
Прежде чем мы начнем обучение, мы можем объявить количество эпох, которые мы хотим обучить, а также настроить оптимизатор и планировщик скорости обучения.
params = [p for p in model.parameters() if p.requires_grad] optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005) lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1) num_epochs = 40
Оптимизатор, который мы здесь используем, - это SGD (стохастический градиентный спуск), у нас есть много других оптимизаторов, доступных в torch.optim, мы можем поэкспериментировать с этим. Планировщик скорости обучения помогает регулировать скорость обучения в ходе обучения для достижения большей точности и ускорения сходимости. Мы - планировщик StepLR, который снижает скорость обучения каждой группы параметров на гамму каждые эпохи step_size. Гиперпараметры gamma и step_size будут определять распад lr, мы можем поэкспериментировать с этими значениями. Наконец, мы обучим эту модель 40 эпохам.
Приведенный выше фрагмент кода будет запускать обучение для 40 эпох. optimizer.zero_grad () очищает старые градиенты с последнего шага, иначе мы просто накапливаем градиенты от всех потерь. обратные вызовы ().
Вызов loss.backward () вычислит градиент потерь по всем параметрам потери, которые имеют requires_grad = True, и сохранит их в parameter.grad атрибут для каждого параметра.
optimizer.step () обновит все параметры на основе parameter.grad. lr_scheduler.step () отрегулирует скорость обучения. Мы должны убедиться, что lr.scheduler.step () был вызван после optimizer.step (), иначе первое значение расписания скорости обучения будет пропущено. Мы сохранили модель для вывода и контрольной точки на случай, если нам понадобится возобновить обучение для других эпох. После завершения обучения у нас будет файл model.pth.
Для вывода нам нужно создать test_dataset и test_data_loader.
TestDataset похож на TrainDataset, в нем также будут реализованы методы __getitem__ и __len__. Но это просто вернет изображения вместо изображения и ограничивающей рамки.
test_dataset = TestDataset('test/')
test_data_loader = DataLoader(
test_dataset,
batch_size=2,
shuffle=True,
num_workers=2,
collate_fn=collate_fn
)
мы можем повторно использовать код инициализации модели, который мы использовали для обучения с предварительно обученным как Fase, а затем мы можем загрузить наш файл веса, как показано ниже.
model.load_state_dict(torch.load('./model.pth'))
model.to(device)
Давайте теперь построим наши прогнозы.
У нас было 1500 изображений для обучения, и это заняло около 5 часов на графическом процессоре NVIDIA Tesla K80, и точность довольно приличная. Мы все еще можем улучшить эту модель с помощью различных методов, таких как увеличение изображения, вывод ансамбля и ТТА. Этот блог подробно объяснит эти вещи.
Спасибо за прочтение!
Francium Tech - это технологическая компания, специализирующаяся на поставке программного обеспечения высочайшего качества и масштабируемости на экстремальных скоростях. Цифры и размер данных нас не пугают. Если у вас есть какие-либо требования или вы хотите бесплатно проверить работоспособность своих систем или архитектуры, напишите письмо по адресу [email protected], мы свяжемся с вами!
