Source code for pyparadigm.extras
"""Using functionality from this file requires extra dependencies."""
import numpy as np
from matplotlib import cm
import contextlib
with contextlib.redirect_stdout(None):
import pygame
def _normalize(mat: np.ndarray):
"""rescales a numpy array, so that min is 0 and max is 255"""
return ((mat - mat.min()) * (255 / mat.max())).astype(np.uint8)
[docs]def to_24bit_gray(mat: np.ndarray):
"""returns a matrix that contains RGB channels, and colors scaled
from 0 to 255"""
return np.repeat(np.expand_dims(_normalize(mat), axis=2), 3, axis=2)
[docs]def apply_color_map(name: str, mat: np.ndarray = None):
"""returns an RGB matrix scaled by a matplotlib color map"""
def apply_map(mat):
return (cm.get_cmap(name)(_normalize(mat))[:, :, :3] * 255).astype(np.uint8)
return apply_map if mat is None else apply_map(mat)
[docs]def mat_to_surface(mat: np.ndarray, transformer=to_24bit_gray):
"""Can be used to create a pygame.Surface from a 2d numpy array.
By default a grey image with scaled colors is returned, but using the
transformer argument any transformation can be used.
:param mat: the matrix to create the surface of.
:type mat: np.ndarray
:param transformer: function that transforms the matrix to a valid color
matrix, i.e. it must have 3dimension, were the 3rd dimension are the color
channels. For each channel a value between 0 and 255 is allowed
:type transformer: Callable[np.ndarray[np.ndarray]]"""
return pygame.pixelcopy.make_surface(transformer(mat.transpose())
if transformer is not None else mat.transpose())