Spaces:
Runtime error
Runtime error
File size: 2,402 Bytes
1a79cb6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 |
"""
Plot and save images
"""
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import os.path
import numpy as np
import cv2
def bgr2rgb(img):
# OpenCV's BGR to RGB
rgb = np.copy(img)
rgb[..., 0], rgb[..., 2] = img[..., 2], img[..., 0]
return rgb
def check_do_plot(func):
def inner(self, *args, **kwargs):
if self.do_plot:
func(self, *args, **kwargs)
return inner
def check_do_save(func):
def inner(self, *args, **kwargs):
if self.do_save:
func(self, *args, **kwargs)
return inner
class Plotter(object):
def __init__(self, plot=True, rows=0, cols=0, num_images=0, out_folder=None, out_filename=None):
self.save_counter = 1
self.plot_counter = 1
self.do_plot = plot
self.do_save = out_filename is not None
self.out_filename = out_filename
self.set_filepath(out_folder)
if (rows + cols) == 0 and num_images > 0:
# Auto-calculate the number of rows and cols for the figure
self.rows = np.ceil(np.sqrt(num_images / 2.0))
self.cols = np.ceil(num_images / self.rows)
else:
self.rows = rows
self.cols = cols
def set_filepath(self, folder):
if folder is None:
self.filepath = None
return
if not os.path.exists(folder):
os.makedirs(folder)
self.filepath = os.path.join(folder, 'frame{0:03d}.png')
self.do_save = True
@check_do_save
def save(self, img, filename=None):
if self.filepath:
filename = self.filepath.format(self.save_counter)
self.save_counter += 1
elif filename is None:
filename = self.out_filename
mpimg.imsave(filename, bgr2rgb(img))
print(filename + ' saved')
@check_do_plot
def plot_one(self, img):
p = plt.subplot(self.rows, self.cols, self.plot_counter)
p.axes.get_xaxis().set_visible(False)
p.axes.get_yaxis().set_visible(False)
plt.imshow(bgr2rgb(img))
self.plot_counter += 1
@check_do_plot
def show(self):
plt.gcf().subplots_adjust(hspace=0.05, wspace=0,
left=0, bottom=0, right=1, top=0.98)
plt.axis('off')
#plt.show()
plt.savefig('result.png')
@check_do_plot
def plot_mesh(self, points, tri, color='k'):
""" plot triangles """
for tri_indices in tri.simplices:
t_ext = [tri_indices[0], tri_indices[1], tri_indices[2], tri_indices[0]]
plt.plot(points[t_ext, 0], points[t_ext, 1], color)
|