Updated to handle multiple STLs in one geometry

toolboxes/STLtoVoxel/README.rst

@@ -30,30 +30,29 @@ How to use the package
 
 The main script is ``stltovoxel.py`` which should be run at the command line and takes three arguments:
 
-* ``stlfilename`` is name of STL file to convert including the path.
-* ``-matindex`` is an integer which represents the index of the material to be used from the materials file which will accompany the generated geometry file (HDF5 format).
-* ``-dxdyz`` is the spatial discretisation of the generated voxelised mesh. It should be given as three floating point numbers.
+* ``path`` is base path to the folder containing the STL file(s) to convert.
+* ``-dxdyz`` is the spatial discretisation of the generated voxelised mesh. It should be given as a floating point number.
 
 The physical dimensions of the voxelised object will depend on the size of the object in the original STL file and the spatial discretisation chosen.
 
 Example
 -------
 
-To create a voxelised mesh (HDF5 geometry file) from the ubiquitous `Stanford bunny <https://en.wikipedia.org/wiki/Stanford_bunny>`_ STL file, using a spatial discretisation of 1mm and selecting material index 2 from a materials file:
+To create a voxelised mesh (HDF5 geometry file) from the ubiquitous `Stanford bunny <https://en.wikipedia.org/wiki/Stanford_bunny>`_ STL file, using a spatial discretisation of 1mm:
 
 .. code-block:: none
 
-    python -m toolboxes.STLtoVoxel.stltovoxel toolboxes/STLtoVoxel/examples/stl/Stanford_Bunny.stl -matindex 2 -dxdydz 0.001 0.001 0.001
+    python -m toolboxes.STLtoVoxel.stltovoxel toolboxes/STLtoVoxel/examples/stl/Stanford_Bunny.stl -dxdydz 0.001
 
 Since the number of voxels are 108 x 88 x 108 and the spatial discretisation chosen is 1mm, the physical dimensions of the Stanford bunny when imported into gprMax will be 0.108 x 0.088 x 0.108mm.
 
-The following is an example of a ``materials.txt`` file that can be used with the generated geometry file (HDF5 format) when importing into gprMax. Since ``-matindex`` is set to 2 the material with name ``hdpe``, i.e. a plastic, will be used.
+The following is an example of a ``materials.txt`` file that can be used with the generated geometry file (HDF5 format) when importing into gprMax. The material index used in the HDF5 geometry file corresponds to the number of STL files converted, e.g. it will be zero if only a single STL file is converted.
 
 .. literalinclude:: ../../toolboxes/STLtoVoxel/examples/materials.txt
     :language: none
     :linenos:
 
-The following Python script (using our Python API) can be used to import the generated geometry file ``Stanford_Bunny.h5`` and materials file ``materials.txt`` into a gprMax model:
+The following Python script (using our Python API) can be used to import the generated geometry file ``Stanford_Bunny.h5`` and materials file ``materials.txt`` into a gprMax model. The bunny material will be sand, i.e. index zero in the materials file.
 
 .. literalinclude:: ../../toolboxes/STLtoVoxel/examples/bunny.py
     :language: python

toolboxes/STLtoVoxel/convert.py

@@ -6,12 +6,13 @@ from . import slice
 
 def convert_meshes(meshes, discretization, parallel=True):
     scale, shift, shape = slice.calculate_scale_shift(meshes, discretization)
-    vol = np.zeros(shape[::-1], dtype=np.int32)
+    vol = np.zeros(shape[::-1], dtype=np.int8)
+    vol.fill(-1) # Fill array with -1 to indicate background in gprMax
 
     for mesh_ind, org_mesh in enumerate(meshes):
         slice.scale_and_shift_mesh(org_mesh, scale, shift)
         cur_vol = slice.mesh_to_plane(org_mesh, shape, parallel)
-        vol[cur_vol] = mesh_ind + 1
+        vol[cur_vol] = mesh_ind # Removed plus 1 to work with gprMax material indexing
 
     return vol, scale, shift
 

toolboxes/STLtoVoxel/stltovoxel.py

@@ -1,10 +1,11 @@
 import argparse
+import glob
 import logging
-from pathlib import Path
+import os
 
 import h5py
 
-from .convert import convert_file
+from .convert import convert_files
 
 logger = logging.getLogger(__name__)
 logging.basicConfig(format="%(message)s", level=logging.INFO)
@@ -13,29 +14,42 @@ if __name__ == "__main__":
     # Parse command line arguments
     parser = argparse.ArgumentParser(
         description="Allows the user to convert a STL files to voxelized mesh.",
-        usage="cd gprMax; python -m user_libs.STLtoVoxel.stltovoxel stlfilename matindex dx_dy_dz",
+        usage="cd gprMax; python -m toolboxes.STLtoVoxel.stltovoxel stlfilename -matindex -dxdydz",
     )
-    parser.add_argument("stlfilename", help="name of STL file to convert including path")
-    parser.add_argument("-matindex", type=int, required=True, help="index of material to extract from STL file")
+    parser.add_argument("stlfiles", help="can be the filename of a single STL file, or the path to folder containing multiple STL files")
     parser.add_argument(
-        "-dxdydz", nargs="+", type=float, required=True, help="discretisation to use in voxelisation process"
+        "-dxdydz", type=float, required=True, help="discretisation to use in voxelisation process"
     )
     args = parser.parse_args()
 
-    filename_stl = Path(args.stlfilename)
-    dxdydz = tuple(args.dxdydz)
+    if os.path.isdir(args.stlfiles):
+        path = args.stlfiles
+        files = sorted(glob.glob(path + "/*.stl"))
+        filename_hdf5 = os.path.join(path, os.path.basename(path) + "_geo.h5")
+        filename_mats = os.path.join(path, os.path.basename(path) + "_mats.txt")
+    elif os.path.isfile(args.stlfiles):
+        path = os.path.dirname(args.stlfiles)
+        files = args.stlfiles
+        filename_hdf5 = os.path.join(path, os.path.split(os.path.basename(path))[0] + "_geo.h5")
+        filename_mats = os.path.join(path, os.path.split(os.path.basename(path))[0] + "_mats.txt")
 
-    logger.info(f"\nConverting STL file: {filename_stl.name}")
-    model_array = convert_file(filename_stl, dxdydz)
-    model_array[model_array == 0] = -1
-    model_array[model_array == 1] = args.matindex
+    dxdydz = (args.dxdydz, args.dxdydz, args.dxdydz)
+
+    newline = "\n\t"
+    logger.info(f"\nConverting STL file(s): {newline.join(files)}")
+    model_array = convert_files(files, dxdydz)
     logger.info(f"Number of voxels: {model_array.shape[0]} x {model_array.shape[1]} x {model_array.shape[2]}")
     logger.info(f"Spatial discretisation: {dxdydz[0]} x {dxdydz[1]} x {dxdydz[2]}m")
 
     # Write HDF5 file for gprMax using voxels
-    filename_hdf5 = filename_stl.with_suffix(".h5")
     with h5py.File(filename_hdf5, "w") as f:
         f.create_dataset("data", data=model_array)
         f.attrs["dx_dy_dz"] = (dxdydz[0], dxdydz[1], dxdydz[2])
+    logger.info(f"Written geometry object file: {filename_hdf5}")
 
-    logger.info(f"Written geometry object file: {filename_hdf5.name}")
+    # Write materials file for gprMax
+    # with open(filename_mats, 'w') as f:
+    #     for i, file in enumerate(files):
+    #         name = os.path.splitext(os.path.basename(file))[0]
+    #         f.write(f"#material: {i + 1} 0 1 0 {name}" + "\n")
+    # logger.info(f"Written materials file: {filename_mats}")