Added a module to the tools sub-package that enables a PNG image file to be converted into a HDF5 file that can then be used to import geometry.

tools/convert_png2h5.py

@@ -0,0 +1,126 @@
+# Copyright (C) 2015-2017: The University of Edinburgh
+#                 Authors: Craig Warren and Antonis Giannopoulos
+#
+# This file is part of gprMax.
+#
+# gprMax 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.
+#
+# gprMax is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with gprMax.  If not, see <http://www.gnu.org/licenses/>.
+
+import argparse
+import os
+
+import h5py
+import matplotlib.image as mpimg
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+class Cursor(object):
+    """Get RGB(A) value of pixel at x,y coordinate of button press and store in a list."""
+
+    def __init__(self, im, materials):
+        """
+        Args:
+            im (ndarray): Pixels of the image.
+            materials (list): To store selected RGB(A) values of selected pixels.
+        """
+        self.im = im
+        self.materials = materials
+        plt.connect('button_press_event', self)
+
+    def __call__(self, event):
+        """
+        Args:
+            event (MouseEvent): matplotlib mouse event.
+        """
+        if not event.dblclick:
+            x, y = event.xdata, event.ydata
+            if x is not None and y is not None:
+                pixel = self.im[int(y), int(x)]
+                pixel = np.floor(pixel * 255).astype(np.int16) # Convert pixel values from float (0-1) to integer (0-255)
+                match = pixel_match(materials, pixel)
+                if match is False:
+                    print('x, y: {} {} px; RGB: {}; material ID: {}'.format(int(x), int(y), pixel[:-1], len(self.materials)))
+                    materials.append(pixel)
+
+def pixel_match(pixellist, pixeltest):
+    """Checks if the RGB(A) value of a pixel already exists in a list of pixel values.
+
+    Args:
+        pixellist (list): List of numpy arrays of pixels to test against.
+        pixeltest (ndarray): RGB(A) value of test pixel.
+
+    Returns:
+        match (boolean): True if pixel is matched in pixel list or False if not found.
+    """
+    match = False
+    for pixel in pixellist:
+        if np.all(pixel == pixeltest):
+            match = True
+            break
+    return match
+
+
+if __name__ == "__main__":
+    
+    # Parse command line arguments
+    parser = argparse.ArgumentParser(description='Convert a PNG image to a HDF5 file that can be used to import geometry (#geometry_objects_read) into a 2D gprMax model. Colours from the image are selected which correspond to a list of materials that should be supplied in a separate text file.', usage='python convert_png2h5.py imagefile dx dy dz')
+    parser.add_argument('imagefile', help='name of image file including path')
+    parser.add_argument('dxdydz', type=float, action='append', nargs=3, help='spatial resolution of model, e.g. dx dy dz')
+    parser.add_argument('-zcells', default=1, type=int, help='number of cells for domain in z-direction (infinite direction)')
+    args = parser.parse_args()
+
+    # Open image file
+    im = mpimg.imread(args.imagefile)
+    
+    # Store image data to use for creating geometry
+    imdata = np.rot90(im, k=3) # Rotate 90CW
+    imdata = np.floor(imdata * 255).astype(np.int16) # Convert pixel values from float (0-1) to integer (0-255)
+
+    print('Reading PNG image file: {}'.format(os.path.split(args.imagefile)[1]))
+    print(' 1. Select discrete material colours by clicking on parts of the image.\n 2. When all materials have been selected close the image.')
+
+    # List to hold selected RGB values from image
+    materials = []
+
+    # Plot image and record rgb values from mouse clicks
+    fig = plt.figure(num=os.path.split(args.imagefile)[1], facecolor='w', edgecolor='w')
+    im = np.flipud(im) # Flip image for viewing with origin in lower left
+    plt.imshow(im, interpolation='nearest', aspect='equal', origin='lower')
+    Cursor(im, materials)
+    plt.show()
+
+    # Format spatial resolution into tuple
+    dxdydz = (args.dxdydz[0][0], args.dxdydz[0][1], args.dxdydz[0][2])
+    
+    # Filename for geometry (HDF5) file
+    hdf5file = os.path.splitext(args.imagefile)[0] + '.h5'
+    
+    # Array to store geometry data (initialised as background, i.e. -1)
+    data = np.ones((imdata.shape[0], imdata.shape[1], args.zcells), dtype=np.int16) * -1
+
+    # Write geometry (HDF5) file
+    with h5py.File(hdf5file, 'w') as fout:
+
+        # Add attribute with name 'dx, dy, dz' for spatial resolution
+        fout.attrs['dx, dy, dz'] = dxdydz
+
+        # Use a boolean mask to match selected pixel values with position in image
+        for i, material in enumerate(materials):
+            mask = np.all(imdata == material, axis=-1)
+            data[mask,:] = i
+        
+        # Write data to file
+        fout.create_dataset('data', data=data)
+
+    print('Written HDF5 file: {}'.format(os.path.split(hdf5file)[1]))