120 lines
3.0 KiB
Python
120 lines
3.0 KiB
Python
import numpy as np
|
|
import zlib
|
|
import io
|
|
|
|
# Database structure: (all is little endian)
|
|
# HEADER:
|
|
# 0-4 "GEOMG"
|
|
# 5 Version
|
|
# 6-7 Fragmentation
|
|
# 8-9 Horizontal size in px
|
|
# 10-11 Vertical size in px
|
|
# 12 Number of layers
|
|
# LAYER1:
|
|
# HEADER:
|
|
# 0 Data type
|
|
# 1 Bytes per point (+16 if signed)
|
|
# 2-5 Length of table
|
|
# 6-7 Length of metadata
|
|
# METADATA
|
|
# TABLE:
|
|
# 4-bytes address of every chunk, zlib compressed
|
|
# DATA:
|
|
# chunk1:
|
|
# raw data, bytes per pixel depend on 'itemsize'
|
|
# chunk2:
|
|
# ...
|
|
# ...
|
|
# LAYER2:
|
|
# ...
|
|
|
|
version = b'\x01'
|
|
|
|
# Conversion to little endian
|
|
def le(n):
|
|
return n.newbyteorder("<").tobytes()
|
|
|
|
layer_count = 0
|
|
|
|
def layer(data, datamap, datatype, frag, meta=b""): # Add a layer
|
|
dmin = int(np.floor(datamap.min()))
|
|
dmax = int(np.floor(datamap.max()))
|
|
signed = dmin < 0
|
|
letter = "i" if signed else "u"
|
|
absmax = max(dmax, -dmin-1)
|
|
if absmax < 0x100:
|
|
itemsize = 1
|
|
elif absmax < 0x10000:
|
|
itemsize = 2
|
|
elif absmax < 0x100000000:
|
|
itemsize = 4
|
|
else:
|
|
itemsize = 8
|
|
|
|
datamap = datamap.astype("<"+letter+str(itemsize))
|
|
|
|
(Y, X) = datamap.shape
|
|
|
|
# Geometry stuff
|
|
table_size_x, table_size_y = int(np.ceil(X / frag)), int(np.ceil(Y / frag))
|
|
table_size = table_size_x * table_size_y
|
|
|
|
layer_table = np.zeros(table_size, dtype=np.uint32).newbyteorder("<") # Table will be a list of the position of every chunk in the data section
|
|
layer_data = io.BytesIO()
|
|
i = 0
|
|
n = 0
|
|
for y in range(0, Y, frag):
|
|
for x in range(0, X, frag):
|
|
part = datamap[y:y+frag,x:x+frag] # Take only the chunk x;y
|
|
part_raw = part.tobytes() # Convert it into binary
|
|
n += layer_data.write(zlib.compress(part_raw, 9)) # Add this to the binary buffer, and increment n by the number of bytes
|
|
layer_table[i] = n # Sets the position of the end of the chunk
|
|
i += 1
|
|
|
|
layer_table_raw = zlib.compress(layer_table.tobytes(), 9) # Compress the table too
|
|
table_length = len(layer_table_raw)
|
|
meta_length = len(meta)
|
|
layer_header = le(np.uint8(datatype)) + le(np.uint8(itemsize+signed*16)) + le(np.uint32(table_length)) + le(np.uint16(meta_length)) + meta
|
|
|
|
# Add this to the main binary
|
|
data.write(layer_header)
|
|
data.write(layer_table_raw)
|
|
data.write(layer_data.getbuffer())
|
|
|
|
global layer_count
|
|
layer_count += 1
|
|
|
|
def generate(file_output, file_conf, heightmap, rivermap=None, landmap=None, landmap_legend=None, frag=80, scale=40):
|
|
global table_size
|
|
print("Generating database")
|
|
|
|
(Y, X) = heightmap.shape
|
|
|
|
data = io.BytesIO() # This allows faster concatenation
|
|
|
|
heightmap //= scale
|
|
|
|
print("Adding heightmap")
|
|
layer(data, heightmap, 0, frag)
|
|
|
|
if type(rivermap) is not type(None):
|
|
print("Adding rivermap")
|
|
layer(data, rivermap, 1, frag)
|
|
|
|
if type(landmap) is not type(None):
|
|
print("Adding landcover")
|
|
layer(data, landmap, 2, frag, meta=landmap_legend)
|
|
|
|
print("Writing file")
|
|
# Build file header
|
|
header = b'GEOMG' + version + le(np.uint16(frag)) + le(np.uint16(X)) + le(np.uint16(Y)) + le(np.uint8(layer_count))
|
|
|
|
# Write in files
|
|
file_output.write(header + data.getbuffer())
|
|
file_output.close()
|
|
|
|
file_conf.write("scale_y = 1")
|
|
file_conf.close()
|
|
|
|
print("Done.")
|