def save(self, grid, output_dir, output_name):
"""
Save the output to a specific dir.
Args:
grid (numpy.ndarray): Grid of all heatmaps
output_dir (str): Output directory path
output_name (str): Output name
"""
save_rgb(grid, output_dir, output_name)
def test_should_save_rgb_image(output_dir):
input_image = np.ones((28, 28, 3), dtype="uint8")
save_rgb(input_image, output_dir, "rgb.png")
assert len(list(Path(output_dir).iterdir())) == 1
