diff --git a/worldengine/cli/main.py b/worldengine/cli/main.py index 4c6d3446..ccc9bbf9 100644 --- a/worldengine/cli/main.py +++ b/worldengine/cli/main.py @@ -356,18 +356,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() @@ -386,7 +381,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: @@ -547,7 +542,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) diff --git a/worldengine/imex/__init__.py b/worldengine/imex/__init__.py index 10b759fe..63add49e 100644 --- a/worldengine/imex/__init__.py +++ b/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)