91 lines
2.3 KiB
Python
91 lines
2.3 KiB
Python
"""
|
|
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)
|