Adding dynamic plots to video - opencv

Question

I'm trying to append dynamic graph plots to an existing video.
Currently I create a pyplot figure for each frame and concatenate the video frame with the figure (using np.from_buffer to extract the frame of each figure), which is very slow. In the end I use OpenCV's VideoWriter to export the concatenated frames into a video file.
I tried to use the FuncAnimation class, but I couldn't figure how to adopt it into my pipeline.
Below is my current code. Any way I could speed things up?
class DynamicGraph(ABC):
def __init__(self, data, height, dpi):
self._data = data.copy()
self.width, self.height = int(1.5 * height), height
self.filters = filters
self.dpi = dpi
self.data = data
nona = data[~np.isnan(data)]
self.ylim = (np.min(nona), np.max(nona))
def to_numpy(self, fig):
fig.tight_layout()
fig.canvas.draw()
frame = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
frame = cv2.cvtColor(frame.reshape(fig.canvas.get_width_height()[::-1] + (3,)), cv2.COLOR_RGB2BGR)
plt.close(fig)
return frame
def __call__(self, i):
return self.plot(i)
#abstractmethod
def plot(self, i):
pass
class DynamicSignal(DynamicGraph):
def __init__(self, title, data, legend, xlabel, ylabel, window_size, height, filters=(), dpi=128):
self.title, self.legend, self.xlabel, self.ylabel = title, legend, xlabel, ylabel
self.window_size = window_size
self.radius = self.window_size // 2
self.ticks = 50
pad = np.zeros((self.radius, data.shape[1]))
super().__init__(data=np.concatenate((pad, data, pad), axis=0), height=height, filters=filters, dpi=dpi)
def plot(self, i):
fig, ax = plt.subplots(figsize=(self.width / self.dpi, self.height / self.dpi), dpi=self.dpi)
x = np.arange(i-self.radius, i+self.radius+1)
y = self.data[i:i+2*self.radius+1]
ax.plot(x, y)
mn, mx = x.min(), x.max()
xticks = np.linspace(*(np.round(np.array([mn, mx])/self.ticks) * self.ticks), (mx - mn) // self.ticks + 1)
ax.set(title=self.title, xlabel=self.xlabel, xlim=(mn, mx), xticks=xticks, ylabel=self.ylabel, ylim=self.ylim)
ax.axvline(x=i, color='r', linestyle='dotted')
ax.grid()
ax.legend(self.legend, loc='upper right')
return self.to_numpy(fig)
class VideoCreator:
def __init__(self, painters=(), graphs=()):
self.painters = painters
self.graphs = graphs
def process_frame(self, frame, i):
out_frame = frame.copy()
if any(self.graphs):
graphs = [graph(i) for graph in self.graphs]
out_frame = np.concatenate((out_frame, np.concatenate(graphs, axis=0)), axis=1)
return out_frame
def create_video(self, video_path, video_data, out_path, start=None, end=None):
fps, length, (width, height) = video_data['fps'], video_data['frame_count'], video_data['resolution']
start, end = int(0 if start is None else start), int(length if end is None else end)
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out = cv2.VideoWriter(out_path, fourcc, fps, width + self.graphs[0].width if any(self.graphs) else 0, height))
cap = cv2.VideoCapture(video_path)
cap.set(cv2.CAP_PROP_POS_FRAMES, start)
for i in tqdm(range(start, end), desc="Writing video result"):
ret, frame = cap.read()
frame = self.process_frame(frame, i)
out.write(frame)
cap.release()
out.release()