Fast Python plotting library to draw plots directly on 2D numpy array image buffer?

Question

I often draw 2D plots directly on 2D numpy array image buffer coming from opencv webcam stream using opencv drawing functions. And, I send the numpy array to imshow and video writer to monitor and create a video.

import cv2
import numpy as np

cap = cv2.VideoCapture(0)
ret, frame = cap.read()  # frame is a 2D numpy array w640 h480
h,w,_ = frame.shape # (480,640,3)
x = np.arange(w)

writer = cv2.VideoWriter( 'out.avi', cv2.cv.FOURCC('D','I','V','3'),
            fps=30, frameSize=(w,h), isColor=True )

while True:
    ret, frame = cap.read()  # frame is a 2D numpy array w640 h480

    B = frame[:,:,0].sum(axis=0)
    B = h - h * B / B.max()
    G = frame[:,:,1].sum(axis=0)
    G = h - h * G / G.max()
    R = frame[:,:,2].sum(axis=0)
    R = h - h * R / R.max()

    pts = np.vstack((x,B)).astype(np.int32).T
    cv2.polylines(frame, [pts], isClosed=False, color=(255,0,0))
    pts = np.vstack((x,G)).astype(np.int32).T
    cv2.polylines(frame, [pts], isClosed=False, color=(0,255,0))
    pts = np.vstack((x,R)).astype(np.int32).T
    cv2.polylines(frame, [pts], isClosed=False, color=(0,0,255))

    writer.write(frame)

    cv2.imshow('frame', frame)
    key = cv2.waitKey(33) & 0xFF # for 64 bit PC
    if key in 27: # ESC key
        break

cap.release()
writer.release()

This works great but I wonder if I can do more stuff like what matplotlib can do such as axes, ticks, grid, title, bar graphs etc without rolling out my own plotting library based on basic cv2 drawing functions, which will be possible but I don't want reinventing the wheel.

Looking into https://wiki.python.org/moin/NumericAndScientific/Plotting, there are so many plotting libraries. So, I feel that one of them might already do this.

I thought about using matplotlib and export the plot as image by savefig. But this will be slow for video capture.

(edit) I could embed a matplotlib plot into frame using mplfig_to_npimage as suggested in the accepted answer! It seems fast enough for video rate.

import cv2
from pylab import *
from moviepy.video.io.bindings import mplfig_to_npimage

fp = r'C:/Users/Public/Videos/Sample Videos/Wildlife.wmv'

cap = cv2.VideoCapture(fp)
ret, frame = cap.read()  # frame is a 2D numpy array
h,w,_ = frame.shape
writer = cv2.VideoWriter( 'out.avi', cv2.cv.FOURCC('D','I','V','3'),
                fps=30, frameSize=(w,h), isColor=True )

# prepare a small figure to embed into frame
fig, ax = subplots(figsize=(4,3), facecolor='w')
B = frame[:,:,0].sum(axis=0)
line, = ax.plot(B, lw=3)
xlim([0,w])
ylim([40000, 130000])  # setup wide enough range here
box('off')
tight_layout()

graphRGB = mplfig_to_npimage(fig)
gh, gw, _ = graphRGB.shape

while True:
    ret, frame = cap.read()  # frame is a 2D numpy array
    B = frame[:,:,0].sum(axis=0)
    line.set_ydata(B)
    frame[:gh,w-gw:,:] = mplfig_to_npimage(fig)

    cv2.imshow('frame', frame)
    writer.write(frame)

    key = cv2.waitKey(33) & 0xFF # for 64 bit
    if key in 27: # ESC key
        break

cap.release()
writer.release()

Answer 1

So, if I'm getting this right, you want:

• Plot conceptual figures (paths, polygons), with out-of-the-box indicators (axes, enclosing plots automagically) over an image
• Video dump and hopefully realtime streaming.

If so, I'd recommend using matplotlib with moviepy.

Indeed doing savefig to stream video is not the best way, but you can make these two work rather easily.

Including small example from the link above for the record (mind licenses):

import matplotlib.pyplot as plt
import numpy as np
from moviepy.video.io.bindings import mplfig_to_npimage
import moviepy.editor as mpy

# DRAW A FIGURE WITH MATPLOTLIB

duration = 2

fig_mpl, ax = plt.subplots(1,figsize=(5,3), facecolor='white')
xx = np.linspace(-2,2,200) # the x vector
zz = lambda d: np.sinc(xx**2)+np.sin(xx+d) # the (changing) z vector
ax.set_title("Elevation in y=0")
ax.set_ylim(-1.5,2.5)
line, = ax.plot(xx, zz(0), lw=3)

# ANIMATE WITH MOVIEPY (UPDATE THE CURVE FOR EACH t). MAKE A GIF.

def make_frame_mpl(t):
    line.set_ydata( zz(2*np.pi*t/duration))  # <= Update the curve
    return mplfig_to_npimage(fig_mpl) # RGB image of the figure

animation =mpy.VideoClip(make_frame_mpl, duration=duration)
animation.write_gif("sinc_mpl.gif", fps=20)