Статистика вывода Matplotlib / Pyplot: отображение имени и количества элементов, распознаваемых с помощью более быстрого CNN - PullRequest
0 голосов
/ 13 февраля 2020

Я использую Pyplot из Matplotlib для выходной статистики приложения Faster CNN и хотел бы знать, сколько, скажем, людей или предметов было распознано на картинке.

Текущая версия Pyplot перечисляет только первый элемент, который был распознан со значением 1.

Как изменить код, чтобы включить все элементы с фактическим количеством распознанных единиц на элемент?

from matplotlib import pyplot as plt

names = pred_cls[i]
values = int(len(names.split()))

plt.figure(figsize=(9, 3))

plt.subplot(131)
plt.bar(names, values)
plt.suptitle('Categorical Plotting')
plt.show()

Вот более ранняя часть кода names и values относится к:

def get_prediction(img_path, threshold):
  img = Image.open(img_path)
  transform = T.Compose([T.ToTensor()])
  img = transform(img)
  pred = model([img])
  pred_class = [COCO_INSTANCE_CATEGORY_NAMES[i] for i in list(pred[0]['labels'].numpy())]
  pred_boxes = [[(i[0], i[1]), (i[2], i[3])] for i in list(pred[0]['boxes'].detach().numpy())]
  pred_score = list(pred[0]['scores'].detach().numpy())
  pred_t = [pred_score.index(x) for x in pred_score if x>threshold][-1]
  pred_boxes = pred_boxes[:pred_t+1]
  pred_class = pred_class[:pred_t+1]
  return pred_boxes, pred_class

def object_detection_api(img_path, threshold=0.7, rect_th=3, text_size=3, text_th=3):

  boxes, pred_cls = get_prediction(img_path, threshold)
  img = cv2.imread(img_path)
  img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
  for i in range(len(boxes)):
    cv2.rectangle(img, boxes[i][0], boxes[i][1],color=(0, 255, 0), thickness=2)
    cv2.putText(img,pred_cls[i], boxes[i][0], cv2.FONT_HERSHEY_SIMPLEX, text_size, (0,255,0),thickness=2)
  plt.figure(figsize=(20,30))
  plt.imshow(img)
  plt.xticks([])
  plt.yticks([])
  plt.show()
...