utils.py 851 Bytes
Newer Older
1
2
3
import os

import cv2 as cv
Chen Sirui's avatar
Chen Sirui committed
4
import matplotlib
5
import matplotlib.animation as manimation
Chen Sirui's avatar
Chen Sirui committed
6
7
8
import matplotlib.pyplot as plt
import numpy as np

9
matplotlib.use("agg")
Chen Sirui's avatar
Chen Sirui committed
10
11
12
13
14
15

# Make video can be used to visualize test data


def make_video(xy, filename):
    os.system("rm -rf pics/*")
16
17
18
19
    FFMpegWriter = manimation.writers["ffmpeg"]
    metadata = dict(
        title="Movie Test", artist="Matplotlib", comment="Movie support!"
    )
Chen Sirui's avatar
Chen Sirui committed
20
21
22
23
24
    writer = FFMpegWriter(fps=15, metadata=metadata)
    fig = plt.figure()
    plt.xlim(-200, 200)
    plt.ylim(-200, 200)
    fig_num = len(xy)
25
    color = ["ro", "bo", "go", "ko", "yo", "mo", "co"]
Chen Sirui's avatar
Chen Sirui committed
26
27
28
29
30
    with writer.saving(fig, filename, len(xy)):
        for i in range(len(xy)):
            for j in range(len(xy[0])):
                plt.plot(xy[i, j, 1], xy[i, j, 0], color[j % len(color)])
            writer.grab_frame()