graphics.py

import numpy as np
from PIL import Image
import time
import threading

def save_image(x, path):
    im = Image.fromarray(x)
    im.save(path, optimize = True)
    return

# Assumes [NCHW] format
def save_raster(x, path, rescale=False, width=None):
    t = threading.Thread(target=_save_raster, args=(x, path, rescale, width))
    t.start()

def _save_raster(x, path, rescale, width):
    x = to_raster(x, rescale, width)
    save_image(x, path)

# Shape: (n_patches,rows,columns,channels)
def to_raster_old(x, rescale=False, width=None):
    x = np.transpose(x, (0,3,1,2))

    #x = x.swapaxes(2, 3)
    if len(x.shape) == 3:
        x = x.reshape((x.shape[0],1,x.shape[1],x.shape[2]))
    if x.shape[1] == 1:
        x = np.repeat(x, 3, axis=1)
    if rescale:
        x = (x - x.min()) / (x.max() - x.min()) * 255.
    x = np.clip(x, 0, 255)
    assert len(x.shape) == 4
    assert x.shape[1] == 3
    n_patches = x.shape[0]
    if width is None:
        width = int(np.ceil(np.sqrt(n_patches))) #result width
    height = int(n_patches/width) #result height
    tile_height = x.shape[2]
    tile_width = x.shape[3]
    result = np.zeros((3,int(height*tile_height),int(width*tile_width)), dtype='uint8')
    for i in range(height):
        for j in range(width):
            result[:, i*tile_height:(i+1)*tile_height, j*tile_width:(j+1)*tile_width] = x[i]
    return result


# Shape: (n_patches,rows,columns,channels)
def to_raster(x, rescale=False, width=None):
    if len(x.shape) == 3:
        x = x.reshape((x.shape[0], x.shape[1], x.shape[2], 1))
    if x.shape[3] == 1:
        x = np.repeat(x, 3, axis=3)
    if rescale:
        x = (x - x.min()) / (x.max() - x.min()) * 255.
    x = np.clip(x, 0, 255)
    assert len(x.shape) == 4
    assert x.shape[3] == 3
    n_batch = x.shape[0]
    if width is None:
        width = int(np.ceil(np.sqrt(n_batch)))  # result width
    height = int(n_batch / width)  # result height
    tile_height = x.shape[1]
    tile_width = x.shape[2]
    result = np.zeros((int(height * tile_height), int(width * tile_width), 3), dtype='uint8')
    for i in range(height):
        for j in range(width):
            result[i * tile_height:(i + 1) * tile_height, j * tile_width:(j + 1) * tile_width] = x[width*i+j]
    return result