Как сохранить этот рисунок Matplotlib как массив Numpy?

Question

У меня есть функция, которая берет изображение, хранящееся в виде массива Numpy, рисует на нем несколько прямоугольников, маркирует их, а затем отображает результат.

Форма исходного массива Numpy: (480, 640, 3) - это RGB изображение с камеры. Это, вероятно, не имеет большого значения, но я просто показываю вам пример данных, с которыми я работаю.

Это функция:

def draw_boxes(imdata, v_boxes, v_labels, v_scores):
    fig = pyplot.imshow(imdata)
    # get the context for drawing boxes
    ax = pyplot.gca()
    # plot each box
    for i in range(len(v_boxes)):
        box = v_boxes[i]
        # get coordinates
        y1, x1, y2, x2 = box.ymin, box.xmin, box.ymax, box.xmax
        # calculate width and height of the box
        width, height = x2 - x1, y2 - y1
        # create the shape
        rect = Rectangle((x1, y1), width, height, fill=False, color='white')
        # draw the box
        ax.add_patch(rect)
        # draw text and score in top left corner
        label = "%s (%.3f)" % (v_labels[i], v_scores[i])
        ax.text(x1, y1, label, color='white')
    pyplot.show()

Я бы хотелвзять аннотированное изображение (изображение с нарисованными на нем прямоугольниками и метками) и извлечь все это в виде массива Numpy. По сути, возвращаем аннотированный массив Numpy.

Я потратил пару часов, пробуя различные решения, найденные в Google, но ничего не работает. Например, я не могу сделать это ...

fig.canvas.draw()
X = np.array(fig.canvas.renderer.buffer_rgba())

... потому что fig.canvas.draw () завершается с:

AttributeError: 'AxesImage' object has no attribute 'canvas'

Answer 1

Проблема в том, что ваша переменная fig - это не цифра, а AxesImage, как указывает ошибка. Таким образом, измените первую строку вашего кода следующим образом:

fig, ax = plt.subplots()
ax = plt.imshow(imdata)

Полная функция будет:

def draw_boxes(imdata, v_boxes, v_labels, v_scores):
    fig, ax = plt.subplots()
    ax = plt.imshow(imdata)
    # get the context for drawing boxes
    ax = pyplot.gca()
    # plot each box
    for i in range(len(v_boxes)):
        box = v_boxes[i]
        # get coordinates
        y1, x1, y2, x2 = box.ymin, box.xmin, box.ymax, box.xmax
        # calculate width and height of the box
        width, height = x2 - x1, y2 - y1
        # create the shape
        rect = Rectangle((x1, y1), width, height, fill=False, color='white')
        # draw the box
        ax.add_patch(rect)
        # draw text and score in top left corner
        label = "%s (%.3f)" % (v_labels[i], v_scores[i])
        ax.text(x1, y1, label, color='white')
    fig.canvas.draw()
    X = np.array(fig.canvas.renderer.buffer_rgba(), dtype=float)
    return X