Overhauled GDAL export functionality and interface.

Author: tcld

worldengine/cli/main.py

@@ -341,18 +341,13 @@ def main():
     # -----------------------------------------------------
     export_options = parser.add_argument_group(
         "Export Options", "You can specify the formats you wish the generated output to be in. ")
-    export_options.add_argument("--export-type", dest="export_type",
-                                help="Export to a specific format such as: BMP or PNG",
-                                default="bmp")
-    export_options.add_argument("--export-bpp", dest="export_bpp", type=int,
-                                help="Bits per pixel: 8, 16 and 32",
-                                default=8)
-    export_options.add_argument("--export-signed", dest="export_signed", action="store_true",
-                                help="Used signed bits or not.",
-                                default=False)
-    export_options.add_argument("--normalize", dest="export_normalize", action="store_true",
-                                help="Normalize data to the min and max of your bpp choice.",
-                                default=False)
+    export_options.add_argument("--export-format", dest="export_format", type=str,
+                                help="Export to a specific format such as BMP or PNG. " +
+                                "See http://www.gdal.org/formats_list.html for possible formats.",
+                                default="PNG")
+    export_options.add_argument("--export-datatype", dest="export_datatype", type=str,
+                                help="Type of stored data, e.g. uint16, int32, float32 etc.",
+                                default="uint16")
 
     args = parser.parse_args()
 
@@ -371,7 +366,7 @@ def main():
     sys.setrecursionlimit(args.recursion_limit)
 
     if args.number_of_plates < 1 or args.number_of_plates > 100:
-        usage(error="Number of plates should be a in [1, 100]")
+        usage(error="Number of plates should be in [1, 100]")
 
     operation = "world"
     if args.OPERATOR is None:
@@ -498,7 +493,8 @@ def main():
     elif operation == 'export':
         world = load_world(args.FILE)
         print_world_info(world)
-        export(world, args.export_type, args.export_bpp, args.export_signed, args.export_normalize)
+        export(world, args.export_format, args.export_datatype,
+               path = '%s/%s_elevation' % (args.output_dir, world_name))
     else:
         raise Exception(
             'Unknown operation: valid operations are %s' % OPERATIONS)

worldengine/imex/__init__.py

@@ -12,65 +12,121 @@
 import sys
 
 
-def export(world, export_type, export_bpp, export_signed, export_normalize):
+'''
+Whenever a GDAL short-format (http://www.gdal.org/formats_list.html) is given
+and a unique mapping to a file suffix exists, it is looked up in gdal_mapper.
+
+Trivial ones (i.e. a call to lower() does the job) are not handled:
+   BAG, BMP, BT, ECW, ERS, FITS, GIF, GTA, PNG, RIK, VRT, XPM
+
+All other formats (>100) currently end up with their respective GDAL short-format
+converted to lower-case and might need to be renamed by the user.
+'''
+gdal_mapper = {  # TODO: Find a way to make GDAL provide this mapping.
+    "aig"     : "adf",
+    "bsb"     : "kap",
+    "doq1"    : "doq",
+    "doq2"    : "doq",
+    "esat"    : "n1",
+    "grib"    : "grb",
+    "gtiff"   : "tif",
+    "hfa"     : "img",
+    "jdem"    : "mem",
+    "jpeg"    : "jpg",
+    "msgn"    : "nat",
+    "terragen": "ter",
+    "usgsdem" : "dem",
+}
+
+
+def export(world, export_filetype = 'GTiff', export_datatype = 'float32', path = 'seed_output'):
     try:
         gdal
     except NameError:
         print("Cannot export: please install pygdal.")
         sys.exit(1)
 
-    final_driver = gdal.GetDriverByName(export_type)
+    final_driver = gdal.GetDriverByName(export_filetype)
     if final_driver is None:
-        print("%s driver not registered." % export_type)
+        print("%s driver not registered." % export_filetype)
         sys.exit(1)
 
-    if export_bpp == 8 and export_signed:
-        numpy_type = numpy.int8
-        gdal_type = gdal.GDT_Byte
-    elif export_bpp == 8 and not export_signed:
+    # try to find the proper file-suffix
+    export_filetype = export_filetype.lower()
+    if export_filetype in gdal_mapper:
+        export_filetype = gdal_mapper[export_filetype]
+
+    # Note: GDAL will throw informative errors on its own whenever file type and data type cannot be matched.
+
+    # translate export_datatype; http://www.gdal.org/gdal_8h.html#a22e22ce0a55036a96f652765793fb7a4
+    export_datatype = export_datatype.lower()
+    if export_datatype in ['gdt_byte', 'uint8', 'int8', 'byte', 'char']:  # GDAL does not support int8
+        bpp, signed, normalize = (8, False, True)
         numpy_type = numpy.uint8
-        gdal_type = gdal.GDT_Byte
-    elif export_bpp == 16 and export_signed:
-        numpy_type = numpy.int16
-        gdal_type = gdal.GDT_Int16
-    elif export_bpp == 16 and not export_signed:
+        gdal_type  = gdal.GDT_Byte
+    elif export_datatype in ['gdt_uint16', 'uint16']:
+        bpp, signed, normalize = (16, False, True)
         numpy_type = numpy.uint16
-        gdal_type = gdal.GDT_UInt16
-    elif export_bpp == 32 and export_signed:
-        numpy_type = numpy.int32
-        gdal_type = gdal.GDT_Int32
-    elif export_bpp == 32 and not export_signed:
+        gdal_type  = gdal.GDT_UInt16
+    elif export_datatype in ['gdt_uint32', 'uint32']:
+        bpp, signed, normalize = (32, False, True)
         numpy_type = numpy.uint32
-        gdal_type = gdal.GDT_UInt32
+        gdal_type  = gdal.GDT_UInt32
+    elif export_datatype in ['gdt_int16', 'int16']:
+        bpp, signed, normalize = (16, True, True)
+        numpy_type = numpy.int16
+        gdal_type  = gdal.GDT_Int16
+    elif export_datatype in ['gdt_int32', 'int32', 'int']:  # fallback for 'int'
+        bpp, signed, normalize = (32, True, True)
+        numpy_type = numpy.int32
+        gdal_type  = gdal.GDT_Int32
+    elif export_datatype in ['gdt_float32', 'float32', 'float']:  # fallback for 'float'
+        bpp, signed, normalize = (32, True, False)
+        numpy_type = numpy.float32
+        gdal_type  = gdal.GDT_Float32
+    elif export_datatype in ['gdt_float64', 'float64']:
+        bpp, signed, normalize = (64, True, False)
+        numpy_type = numpy.float64
+        gdal_type  = gdal.GDT_Float64
     else:
-        print ("BPP %d is not valid, we only support 8, 16 and 32." % export_bpp)
-        sys.exit(1)
+        raise TypeError("Type of data not recognized or not supported by GDAL: %s" % export_datatype)
 
     # massage data to scale between the absolute min and max
-    elevation = numpy.array(world.elevation['data'])
+    elevation = numpy.copy(world.elevation['data'])
 
-    if not export_signed and elevation.min() < 0.0:
-        elevation += abs(elevation.min())  # TODO: need better way to handle negative numbers
+    # shift data according to minimum possible value
+    if signed:
+        elevation = elevation - world.sea_level()  # elevation 0.0 now refers to sea-level
+    else:
+        elevation -= elevation.min()  # lowest point at 0.0
 
-    if export_normalize:
-        if export_signed:
-            elevation *= (((2**export_bpp)/2)-1)/elevation.max()
+    # rescale data (currently integer-types only)
+    if normalize:
+        # elevation maps usually have a range of 0 to 10, maybe 15 - rescaling for integers is essential
+        if signed:
+            elevation *= (2**(bpp - 1) - 1) / max(abs(elevation.min(), abs(elevation.max())))
         else:
-            elevation *= (2**export_bpp)/elevation.max()
+            elevation *= (2**bpp - 1) / abs(elevation.max())
+
+    # round data (integer-types only)
+    if numpy_type != numpy.float32 and numpy_type != numpy.float64:
+        elevation = elevation.round()
 
+    # switch to final data type; no rounding performed
     elevation = elevation.astype(numpy_type)
 
-    # take elevation data and push it into an intermediate GTiff format
+    # take elevation data and push it into an intermediate GTiff format (some formats don't support being written by Create())
     inter_driver = gdal.GetDriverByName("GTiff")
-    _, inter_file = tempfile.mkstemp()
-    initial_ds = inter_driver.Create(inter_file, world.height, world.width, 1, gdal_type)
+    _, inter_file = tempfile.mkstemp()  # returns: (file-handle, absolute path)
+    initial_ds = inter_driver.Create(inter_file, world.width, world.height, 1, gdal_type)
     band = initial_ds.GetRasterBand(1)
+    
     band.WriteArray(elevation)
     band = None  # dereference band
     initial_ds = None  # save/flush and close
 
     # take the intermediate GTiff format and convert to final format
     initial_ds = gdal.Open(inter_file)
-    final_driver.CreateCopy('seed_output-%d.%s' % (export_bpp, export_type), initial_ds)
+    final_driver.CreateCopy('%s-%d.%s' % (path, bpp, export_filetype), initial_ds)
 
     os.remove(inter_file)