После тонкой настройки модели обнаружения объектов Raster RCNN, как визуализировать прогноз bbox? - PullRequest
Новая
       83

После тонкой настройки модели обнаружения объектов Raster RCNN, как визуализировать прогноз bbox?

2 голосов
/ 18 февраля 2020

Я настроил Pytorch Torchvision модель = torchvision.models.detection.fasterrcnn_resnet50_fpn (pretrained = True) в моем собственном наборе данных.

Я следовал этому руководству https://pytorch.org/tutorials/intermediate/torchvision_tutorial.html#torchvision -обнаружение объекта-finetuning- Учебное пособие , но только для обученных. Быстрее RCNN, а не Маска RCNN.

Я успешно закончил обучение без ошибок, и модель вернула сообщение, содержащее предсказанные блоки, метки и оценки.

В В руководстве, которому я следовал, они показывают, как визуализировать маски, предсказанные обученной моделью. Есть ли подобный метод для визуализации ограничительной рамки? У меня много проблем с выяснением этого.

Спасибо

1 Ответ

1 голос
/ 18 февраля 2020

Прогноз от FasterRCNN имеет вид: 

>>> predictions = model([input_img_tensor])
[{'boxes': tensor([[419.6865, 170.0683, 536.0842, 493.7452],
          [159.0727, 180.3606, 298.8194, 434.4604],
          [439.7836, 222.6208, 452.0138, 271.8359],
          [444.3562, 224.4628, 456.1511, 265.5336],
          [437.7808, 226.5965, 446.2904, 271.2691]], grad_fn=<StackBackward>),
  'labels': tensor([ 1,  1, 32, 32, 32]),
  'scores': tensor([0.9997, 0.9996, 0.5827, 0.2102, 0.0943], grad_fn=<IndexBackward>)}]

, где прогнозируемые блоки имеют формат [x1, y1, x2, y2] со значениями от 0 до H и 0 и W.

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

import cv2
img = cv2.imread('input_iamge.png', cv2.COLOR_BGR2RGB)

for i in range(len(predictions[0]['boxes'])):
    x1, x2, x3, x4 = map(int, predictions[0]['boxes'][i].tolist())
    print(x1, x2, x3, x4)
    image = cv2.rectangle(img, (x1, x2), (x3, x4), (255, 0, 0), 1)

cv2_imshow('img', image)
...