GIMP Python plugin to load and save HDF5 files

file-hdf5.py

#!/usr/bin/env python2
#
# GIMP Plug-in for HDF5 files
# https://www.hdfgroup.org/solutions/hdf5/
#
# Copyright 2018 by Alessandro Re <[email protected]>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# This plugin depends on the h5py library and I used Jon Nordby's OpenRaster
# plugin as an inspiration for building a Python saver/loader.
#
# Version 0.1
# - Basic support for HDF5 files, images are loaded and saved as RGB32 and
#   there is no choice of which images to load from the file (see TODOs)

from gimpfu import *

import h5py
import numpy as np


def load_hdf5(filename, raw_filename):
    with h5py.File(filename, 'r') as hf:
        # Find layers that can be images (by inspecting their size)
        layers_2d_or_3d = {}
        for k in hf:
            if len(hf[k].shape) in [2, 3]:
                layers_2d_or_3d.setdefault(hf[k].shape[:2], []).append(k)
        
        # Select a group of layers
        if len(layers_2d_or_3d) == 1:
            sel = sorted(layers_2d_or_3d.keys())[0]
        else:
            # TODO ask here to user what set of groups to use, if necessary
            # FIXME I'm using the first key
            sel = sorted(layers_2d_or_3d.keys())[0]

        # Create the image (FIXME this is always RGB)
        h, w = sel
        img = gimp.Image(w, h, RGB)
        img.filename = filename
            
        # Add layers
        for i, k in enumerate(layers_2d_or_3d[sel]):
            layer = gimp.Layer(img, k, w, h, RGB_IMAGE, 100)
            layer.set_offsets(0, 0)

            # Get data and map them in ubyte
            data = np.asarray(hf[k][()]).astype(np.float64)
            data = (data - data.min()) / data.ptp()
            data = (data * 255).astype(np.uint8)
            # Make RGB out of grayscale
            if len(data.shape) == 2:
                data = np.stack([data, data, data], axis=2)

            # Write data to pixel region
            pr = layer.get_pixel_rgn(0, 0, w, h)
            pr[:,:] = data.tobytes(order='C')

            img.add_layer(layer, i)

    return img


def save_hdf5(img, drawable, filename, raw_filename):
    # TODO this function saves all layers as RGB
    with h5py.File(filename, 'w') as hf:
        for layer in img.layers:
            reg = layer.get_pixel_rgn(0, 0, layer.width, layer.height)
            data = np.frombuffer(reg[:,:], dtype=np.uint8)
            data = data.reshape(layer.height, layer.width, -1)
            hf[layer.name] = data      


def register_load_handlers():
    gimp.register_load_handler('file-hdf5-load', 'h5', '')
    pdb['gimp-register-file-handler-mime']('file-hdf5-load', 'application/x-hdf5')


def register_save_handlers():
    gimp.register_save_handler('file-hdf5-save', 'h5', '')


register(
    'file-hdf5-load',  # Name
    'Load a HDF5 (.h5) file',  # Description
    'Load a HDF5 (.h5) file',  # Help
    'Alessandro Re',  # Author
    'Alessandro Re',  # Copyright
    '2018',  # Year
    'HDF5',  # Label
    None,  # Image types
    # Params
    [   # Input args, type, name, description, default[, extra]
        (PF_STRING, 'filename', 'The name of the file to load', None),
        (PF_STRING, 'raw-filename', 'The name entered', None),
    ],
    # Result
    [(PF_IMAGE, 'image', 'Output image')], # Results, type, name, description
    load_hdf5, # Function callback
    on_query=register_load_handlers,
    menu="<Load>",
)


register(
    'file-hdf5-save',
    'Save a HDF5 (.h5) file',
    'Save a HDF5 (.h5) file',
    'Alessandro Re',
    'Alessandro Re',
    '2018',
    'HDF5',
    '*',
    [
        (PF_IMAGE, "image", "Input image", None),
        (PF_DRAWABLE, "drawable", "Input drawable", None),
        (PF_STRING, "filename", "The name of the file", None),
        (PF_STRING, "raw-filename", "The name of the file", None),
    ],
    [],  # No results
    save_hdf5,
    on_query=register_save_handlers,
    menu="<Save>"
)

main()

TODOs

• Support for loading and saving grayscale layers (datasets), preliminary support can be added by saving bpp in the layer's name
• bpp is fixed to 32, it would be nice to understand how to use GEGL and support more data types (floats, bools, 8, 16 and 32 bits)
• Datasets which are not loaded are lost!