From 3c09c5a39f4fabed91f67096df22f8bb920a8002 Mon Sep 17 00:00:00 2001 From: lynnschmittwilken Date: Sun, 7 May 2023 08:26:11 -0400 Subject: [PATCH] feat(export): to image, npy, mat, pickle, json --- stimupy/utils/export.py | 187 +++++++++++++++++++++++++++++++++++++--- 1 file changed, 177 insertions(+), 10 deletions(-) diff --git a/stimupy/utils/export.py b/stimupy/utils/export.py index adef9bc..0b66461 100644 --- a/stimupy/utils/export.py +++ b/stimupy/utils/export.py @@ -1,14 +1,22 @@ import json from hashlib import md5 +import copy import numpy as np from PIL import Image +from scipy.io import savemat +import pickle __all__ = [ "array_to_checksum", "array_to_image", + "array_to_npy", + "array_to_mat", + "array_to_pickle", "arrays_to_checksum", "to_json", + "to_mat", + "to_pickle", ] @@ -19,6 +27,7 @@ def array_to_checksum(arr): Parameters ---------- arr : np.ndarray + Array to be hashed. Returns ---------- @@ -27,22 +36,117 @@ def array_to_checksum(arr): return md5(np.ascontiguousarray(arr.round(8))).hexdigest() -def array_to_image(arr, filename): +def array_to_image(arr, filename, norm=True): """ Save a 2D numpy array as a grayscale image file. Parameters ---------- arr : np.ndarray - The data to be stored in the image. Values will be cropped to [0,255]. + Array to be saved. Values will be cropped to [0,255]. + filename : str + full path to the file to be created. + norm : bool + if True (default), multiply array by 255 + """ + if filename[-4:] != ".png" and filename[-4:] != ".jpg": + filename += ".png" + + if isinstance(arr, (np.ndarray, list)): + arr = np.array(arr) + + if norm: + arr = arr * 255 + + if Image: + imsize = arr.shape + im = Image.new("L", (imsize[1], imsize[0])) + im.putdata(arr.flatten()) + im.save(filename) + else: + raise ValueError("arr should be a np.ndarray") + + +def array_to_npy(arr, filename): + """ + Save a numpy array to npy-file. + + Parameters + ---------- + arr : np.ndarray + Array to be saved. + filename : str + full path to the file to be creaated. + """ + if isinstance(arr, (np.ndarray, list)): + np.save(filename, arr) + else: + raise ValueError("arr should be a np.ndarray") + + +def array_to_mat(arr, filename): + """ + Save a numpy array to a mat-file. + + Parameters + ---------- + arr : np.ndarray + Array to be saved. + filename : str + full path to the file to be creaated. + """ + if filename[-4:] != ".mat": + filename += ".mat" + + if isinstance(arr, (np.ndarray, list)): + savemat(filename, {"arr": arr}) + else: + raise ValueError("arr should be a np.ndarray") + + +def array_to_pickle(arr, filename): + """ + Save a numpy array to a pickle-file. + + Parameters + ---------- + arr : np.ndarray + Array to be saved. + filename : str + full path to the file to be creaated. + """ + if filename[-7:] != ".pickle": + filename += ".pickle" + + if isinstance(arr, (np.ndarray, list)): + with open(filename, 'wb') as handle: + pickle.dump({"arr": arr}, handle, protocol=pickle.HIGHEST_PROTOCOL) + else: + raise ValueError("arr should be a np.ndarray") + + +def array_to_json(arr, filename): + """ + Save a numpy array to a (pretty) JSON. + + Parameters + ---------- + arr : np.ndarray + Array to be saved. filename : str full path to the file to be creaated. """ - if Image: - imsize = arr.shape - im = Image.new("L", (imsize[1], imsize[0])) - im.putdata(arr.flatten()) - im.save(filename) + if filename[-5:] != ".json": + filename += ".json" + + if isinstance(arr, np.ndarray): + with open(filename, "w", encoding="utf-8") as f: + json.dump(arr.tolist(), f, ensure_ascii=False, indent=4) + elif isinstance(arr, list): + with open(filename, "w", encoding="utf-8") as f: + json.dump(arr, f, ensure_ascii=False, indent=4) + else: + raise ValueError("arr should be a np.ndarray") def arrays_to_checksum(stim, keys=["img", "mask"]): @@ -74,7 +178,7 @@ def arrays_to_checksum(stim, keys=["img", "mask"]): def to_json(stim, filename): """ - Stimulus-dict(s) as (pretty) JSON + Save stimulus-dict(s) as (pretty) JSON Parameters ---------- @@ -84,6 +188,69 @@ def to_json(stim, filename): full path to the file to be creaated. """ + if filename[-5:] != ".json": + filename += ".json" + # stimulus-dict(s) as (pretty) JSON - with open(filename, "w", encoding="utf-8") as f: - json.dump(stim, f, ensure_ascii=False, indent=4) + if isinstance(stim, dict): + stim2 = copy.deepcopy(stim) + + for key in stim2.keys(): + # np.ndarrays are not serializable; change to list + if isinstance(stim2[key], np.ndarray): + stim2[key] = stim2[key].tolist() + + with open(filename, "w", encoding="utf-8") as f: + json.dump(stim2, f, ensure_ascii=False, indent=4) + else: + raise ValueError("stim should be a dict") + + +def to_mat(stim, filename): + """ + Save stimulus-dict(s) as mat-file + + Parameters + ---------- + stim : dict + stimulus dictionary containing keys + filename : str + full path to the file to be creaated. + + """ + if filename[-4:] != ".mat": + filename += ".mat" + + if isinstance(stim, dict): + savemat(filename, stim) + else: + raise ValueError("stim should be a dict") + + +def to_pickle(stim, filename): + """ + Save stimulus-dict(s) as pickle-file + + Parameters + ---------- + stim : dict + stimulus dictionary containing keys + filename : str + full path to the file to be creaated. + + """ + if filename[-7:] != ".pickle": + filename += ".pickle" + + if isinstance(stim, dict): + stim2 = copy.deepcopy(stim) + + for key in stim2.keys(): + # certain classes can cause problems for pickles; change to list + if key in ["visual_size", "ppd", "shape"]: + stim2[key] = list(stim2[key]) + + with open(filename, 'wb') as handle: + pickle.dump(stim2, handle, protocol=pickle.HIGHEST_PROTOCOL) + else: + raise ValueError("stim should be a dict")