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