from PIL import Image
from typing import List
from machin.parallel import get_context
import os
import numpy as np
import moviepy.editor as mpy
import matplotlib.pyplot as plt
[docs]def show_image(image: np.array,
show_normalized: bool = True,
pause_time: float = 0.01,
title: str = ""):
"""
Use matplotlib to show a single image. You may repeatedly call this method
with the same ``title`` argument to show a video or a dynamically changing
image.
Args:
image: A numpy array of shape (H, W, C) or (H, W), and with ``dtype``
= any float or any int.
When a frame is float type, its value range should be [0, 1].
When a frame is integer type, its value range should be [0, 255].
show_normalized: Show normalized image alongside the original one.
pause_time: Pause time between displaying current image and the next
one.
title: Title of the display window.
"""
if np.issubdtype(image.dtype, np.integer):
image = image.astype(np.floating) / 255
fig = plt.figure(title, clear=True)
fig.canvas.set_window_title(title)
if show_normalized:
ax = fig.add_subplot("121")
ax.set_facecolor((0.0, 0.0, 0.0))
ax.imshow(image, vmin=np.min(image), vmax=np.max(image))
ax2 = fig.add_subplot("122")
ax2.set_facecolor((0.0, 0.0, 0.0))
pix_range = (np.max(image) - np.min(image)) + 1e-6
ax2.imshow((image - np.min(image)) / pix_range, vmin=0, vmax=1)
plt.pause(pause_time)
else:
ax = fig.add_subplot("111")
ax.set_facecolor((0.0, 0.0, 0.0))
ax.imshow(image, vmin=np.min(image), vmax=np.max(image))
plt.pause(pause_time)
[docs]def create_video(frames: List[np.array],
path: str,
filename: str,
extension: str = ".gif",
fps: int = 15):
"""
Args:
frames: A list of numpy arrays of shape (H, W, C) or (H, W), and with
``dtype`` = any float or any int.
When a frame is float type, its value range should be [0, 1].
When a frame is integer type, its value range should be [0, 255].
path: Directory to save the video.
filename: File name.
extension: File extension.
fps: frames per second.
"""
if frames:
for f in range(len(frames)):
if np.issubdtype(frames[f].dtype, np.integer):
frames[f] = frames[f].astype(np.uint8)
elif np.issubdtype(frames[f].dtype, np.floating):
frames[f] = (frames[f] * 255).astype(np.uint8)
if frames[f].ndim == 2:
# consider as a grey scale image
frames[f] = np.repeat(frames[f][:, :, np.newaxis], 3, axis=2)
clip = mpy.ImageSequenceClip(frames, fps=fps)
if extension.lower() == ".gif":
clip.write_gif(os.path.join(path, filename + extension),
fps=fps, verbose=False, logger=None)
else:
clip.write_videofile(os.path.join(path, filename + extension),
fps=fps, verbose=False, logger=None)
[docs]def create_video_subproc(frames: List[np.array],
path: str,
filename: str,
extension: str = ".gif",
fps: int = 15,
daemon: bool = True):
"""
Create video with a subprocess, since it takes a lot of time for ``moviepy``
to encode the video file.
See Also:
:func:`.create_video`
Note:
if ``daemon`` is true, then this function cannot be used in a
daemonic subprocess.
Args:
frames: A list of numpy arrays of shape (H, W, C) or (H, W), and with
``dtype`` = any float or any int.
When a frame is float type, its value range should be [0, 1].
When a frame is integer type, its value range should be [0, 255].
path: Directory to save the video.
filename: File name.
extension: File extension.
fps: frames per second.
daemon: Whether launching the saving process as a daemonic process.
Returns:
A wait function, once called, block until creation has finished.
"""
def wait():
pass
if frames:
p = get_context("spawn").Process(target=create_video,
args=(frames,
path,
filename,
extension,
fps))
p.daemon = daemon
p.start()
def wait():
p.join()
return wait
[docs]def create_image(image: np.array,
path: str,
filename: str,
extension: str = ".png"):
"""
Args:
image: A numpy array of shape (H, W, C) or (H, W), and with
``dtype`` = any float or any int.
When a frame is float type, its value range should be [0, 1].
When a frame is integer type, its value range should be [0, 255].
path: Directory to save the image.
filename: File name.
extension: File extension.
"""
if np.issubdtype(image.dtype, np.integer):
image = image.astype(np.uint8)
elif np.issubdtype(image.dtype, np.floating):
image = (image * 255).astype(np.uint8)
if image.ndim == 2:
# consider as a grey scale image
image = np.repeat(image[:, :, np.newaxis], 3, axis=2)
image = Image.fromarray(image)
image.save(os.path.join(path, filename + extension))
[docs]def create_image_subproc(image: np.array,
path: str,
filename: str,
extension: str = ".png",
daemon: bool = True):
"""
Create image with a subprocess.
See Also:
:func:`.create_image`
Note:
if ``daemon`` is true, then this function cannot be used in a
daemonic subprocess.
Args:
image: A numpy array of shape (H, W, C) or (H, W), and with
``dtype`` = any float or any int.
When a frame is float type, its value range should be [0, 1].
When a frame is integer type, its value range should be [0, 255].
path: Directory to save the image.
filename: File name.
extension: File extension.
daemon: Whether launching the saving process as a daemonic process.
Returns:
A wait function, once called, block until creation has finished.
"""
p = get_context("spawn").Process(target=create_image,
args=(image,
path,
filename,
extension))
p.daemon = daemon
p.start()
def wait():
p.join()
return wait