Add raw image generation

This commit is contained in:
LilyRose2798 2024-04-21 01:05:35 +10:00
parent 7221172e54
commit eebcbae626
2 changed files with 79 additions and 85 deletions

1
.gitignore vendored
View File

@ -1,3 +1,4 @@
/data /data
/public/assets/tiles /public/assets/tiles
/raws
**/*.bin **/*.bin

163
ipmap.py
View File

@ -81,18 +81,17 @@ default_colormap_names = ["viridis"]
default_quantile = 0.995 default_quantile = 0.995
def make_tiles(coords_path: Path, input_path: Path, tiles_dir: Path, *, def make_tiles(coords_path: Path, input_path: Path, tiles_dir: Path, *,
tile_size = default_tile_size, alpha = False, tile_size = default_tile_size, alpha = False, variant_names: list[str] = default_variant_names,
variant_names: list[str] = default_variant_names, colormap_names: list[str] = default_colormap_names, raws_path: Path | None = None,
colormap_names: list[str] = default_colormap_names, quantile = default_quantile, num_rows: int | None = None, skip_iters: int | None = None,
quantile = default_quantile, num_rows: int | None = None, json_path: Path | None = None):
skip_iters: int | None = None, json_path: Path | None = None):
if not 64 <= tile_size <= num_ips_sqrt or tile_size & (tile_size - 1) != 0: if not 64 <= tile_size <= num_ips_sqrt or tile_size & (tile_size - 1) != 0:
raise ValueError(f"tile size must be a power of 2 between 64 and {num_ips_sqrt}") raise ValueError(f"tile size must be a power of 2 between 64 and {num_ips_sqrt}")
if len(variant_names) == 0: if len(variant_names) == 0:
raise ValueError("must specify at least one variant") raise ValueError("must specify at least one variant")
if len(colormap_names) == 0: if len(colormap_names) == 0 and raws_path is None:
raise ValueError("must specify at least one colormap") raise ValueError("must specify at least one colormap or a path to save raws to")
if not 0 <= quantile <= 1: if not 0 <= quantile <= 1:
raise ValueError(f"quantile must be between 0 and 1") raise ValueError(f"quantile must be between 0 and 1")
@ -134,7 +133,7 @@ def make_tiles(coords_path: Path, input_path: Path, tiles_dir: Path, *,
else: else:
tiles_dir_parts = None tiles_dir_parts = None
def create_images(data: np.ndarray, colormap: Colormap, num_colors: int, path: Path): def create_tile_images(data: np.ndarray, colormap: Colormap, num_colors: int, path: Path):
print(f"creating {num_colors} color stop(s) of {colormap.name} colormap...", end = " ", flush = True) print(f"creating {num_colors} color stop(s) of {colormap.name} colormap...", end = " ", flush = True)
colors = np.concatenate(([empty_color], ((colormap([0.0]) if num_colors == 1 else colormap.lut(num_colors))[:, 0:channels] * 255).astype(np.uint8))) colors = np.concatenate(([empty_color], ((colormap([0.0]) if num_colors == 1 else colormap.lut(num_colors))[:, 0:channels] * 255).astype(np.uint8)))
print("done") print("done")
@ -152,17 +151,24 @@ def make_tiles(coords_path: Path, input_path: Path, tiles_dir: Path, *,
y_path = z_path / f"{y}" y_path = z_path / f"{y}"
y_path.mkdir(exist_ok = True) y_path.mkdir(exist_ok = True)
for x in range(tiles_per_side): for x in range(tiles_per_side):
path = y_path / f"{x}.png" x_path = y_path / f"{x}.png"
rows = image_data[ rows = image_data[
y * tile_size : y * tile_size + tile_size, y * tile_size : y * tile_size + tile_size,
x * tile_size : x * tile_size + tile_size, x * tile_size : x * tile_size + tile_size,
] ]
Writer(tile_size, tile_size, greyscale = False, alpha = alpha).write_packed(path.open("wb"), rows) Writer(tile_size, tile_size, greyscale = False, alpha = alpha).write_packed(x_path.open("wb"), rows)
print("done")
def create_raw_image(data: np.ndarray, path: Path):
path.mkdir(exist_ok = True, parents = True)
z_path = path / f"{(data.shape[0] // tile_size).bit_length() - 1}.png"
print(f"writing {data.shape[1]}x{data.shape[0]} raw image to '{path}'...", end = " ", flush = True)
Writer(data.shape[1], data.shape[0], greyscale = False, alpha = True).write_packed(z_path.open("wb"), data)
print("done") print("done")
def get_scan_data() -> tuple[NDArray[np.uint32], NDArray[np.uint32]]: def get_scan_data() -> tuple[NDArray[np.uint32], NDArray[np.uint32]]:
print(f"reading scan data from file '{input_path}'...", end = " ", flush = True) print(f"reading scan data from file '{input_path}'...", end = " ", flush = True)
data = np.fromfile(input_path, dtype = np.uint32).reshape(-1, 2) data = np.fromfile(input_path, count = num_rows * 2 if num_rows else -1, dtype = np.uint32).reshape(-1, 2)
ip_arr = np.copy(data.T[0]) ip_arr = np.copy(data.T[0])
rtt_arr = np.copy(data.T[1]) rtt_arr = np.copy(data.T[1])
print("done") print("done")
@ -182,6 +188,7 @@ def make_tiles(coords_path: Path, input_path: Path, tiles_dir: Path, *,
def generate_density(): def generate_density():
possible_overlaps = 1 possible_overlaps = 1
variant_name = "density"
print(f"allocating empty {num_ips_sqrt}x{num_ips_sqrt} array of density data...", end = " ", flush = True) print(f"allocating empty {num_ips_sqrt}x{num_ips_sqrt} array of density data...", end = " ", flush = True)
density_data = np.zeros((num_ips_sqrt, num_ips_sqrt), dtype = np.uint32) density_data = np.zeros((num_ips_sqrt, num_ips_sqrt), dtype = np.uint32)
@ -208,7 +215,9 @@ def make_tiles(coords_path: Path, input_path: Path, tiles_dir: Path, *,
def write_all_colormaps(): def write_all_colormaps():
for colormap_name, colormap in colormaps: for colormap_name, colormap in colormaps:
create_images(density_data, colormap, possible_overlaps, tiles_dir / "density" / colormap_name) create_tile_images(density_data, colormap, possible_overlaps, tiles_dir / variant_name / colormap_name)
if raws_path is not None:
create_raw_image(density_data, raws_path / variant_name)
write_all_colormaps() write_all_colormaps()
while density_data.shape[0] > tile_size: while density_data.shape[0] > tile_size:
@ -216,64 +225,56 @@ def make_tiles(coords_path: Path, input_path: Path, tiles_dir: Path, *,
write_all_colormaps() write_all_colormaps()
def generate_rtt(): def generate_rtt():
nonlocal rtt_arr
num_colors = (1 << 16) - 1 num_colors = (1 << 16) - 1
multiplier = num_colors - 1 variant_name = "rtt"
def get_rtt_data(): print(f"retrieving {quantile:.1%} quantile for rtt data...", end = " ", flush = True)
nonlocal rtt_arr rtt_quantile = int(np.quantile(rtt_arr, quantile))
print(f"retrieving {quantile:.1%} quantile for rtt data...", end = " ", flush = True) divisor = (rtt_quantile - 1) / (num_colors - 1)
rtt_quantile = np.quantile(rtt_arr, quantile) print("done")
print("done") print("clipping rtt data between 0 and quantile...", end = " ", flush = True)
print(f"scaling rtt data using rtt quantile...", end = " ", flush = True) rtt_arr.clip(0, rtt_quantile, out = rtt_arr)
rtt_arr_f = rtt_arr / rtt_quantile print("done")
print("done") print(f"allocating empty {num_ips_sqrt}x{num_ips_sqrt} array for rtt data...", end = " ", flush = True)
del rtt_arr rtt_data = np.zeros((num_ips_sqrt, num_ips_sqrt), dtype = np.uint32)
collect() print("done")
print("clipping rtt data between 0 and 1...", end = " ", flush = True) print(f"assigning values to rtt data array...", end = " ", flush = True)
rtt_arr_f.clip(0, 1, out = rtt_arr_f) rtt_data[coords] = rtt_arr
print("done") print("done")
del rtt_arr
print(f"allocating empty {num_ips_sqrt}x{num_ips_sqrt} array for rtt data...", end = " ", flush = True) collect()
rtt_data = np.full((num_ips_sqrt, num_ips_sqrt), np.nan, dtype = np.float32)
print("done")
print(f"assigning values to rtt data array...", end = " ", flush = True)
rtt_data[coords] = rtt_arr_f
print("done")
return rtt_data
rtt_data = get_rtt_data()
def squish(): def squish():
nonlocal rtt_data nonlocal rtt_data
print(f"sorting rtt values for median calculation...", end = " ", flush = True) print(f"sorting rtt values for median calculation...", end = " ", flush = True)
rtt_data = np.swapaxes(rtt_data.reshape(rtt_data.shape[0] // 2, 2, rtt_data.shape[1] // 2, 2), 1, 2) rtt_data = np.swapaxes(rtt_data.reshape(rtt_data.shape[0] // 2, 2, rtt_data.shape[1] // 2, 2), 1, 2)
rtt_data[np.isnan(rtt_data)] = np.inf # convert NaNs to Inf so comparisons work correctly
mask = np.empty((rtt_data.shape[0], rtt_data.shape[1]), dtype = np.bool_) mask = np.empty((rtt_data.shape[0], rtt_data.shape[1]), dtype = np.bool_)
np.greater(rtt_data[:, :, 0, 0], rtt_data[:, :, 0, 1], out = mask) # sort first row np.less(rtt_data[:, :, 0, 0], rtt_data[:, :, 0, 1], out = mask) # sort first row
rtt_data[mask, 0] = rtt_data[mask, 0, ::-1] rtt_data[mask, 0] = rtt_data[mask, 0, ::-1]
np.greater(rtt_data[:, :, 1, 0], rtt_data[:, :, 1, 1], out = mask) # sort second row np.less(rtt_data[:, :, 1, 0], rtt_data[:, :, 1, 1], out = mask) # sort second row
rtt_data[mask, 1] = rtt_data[mask, 1, ::-1] rtt_data[mask, 1] = rtt_data[mask, 1, ::-1]
np.greater(rtt_data[:, :, 0, 0], rtt_data[:, :, 1, 0], out = mask) # sort first column np.less(rtt_data[:, :, 0, 1], rtt_data[:, :, 1, 1], out = mask) # sort second column
rtt_data[mask, :, 0] = rtt_data[mask, ::-1, 0]
np.less(rtt_data[:, :, 0, 1], rtt_data[:, :, 1, 1], out = mask) # sort second column in reverse order
rtt_data[mask, :, 1] = rtt_data[mask, ::-1, 1] rtt_data[mask, :, 1] = rtt_data[mask, ::-1, 1]
np.less(rtt_data[:, :, 1, 0], rtt_data[:, :, 1, 1], out = mask) # sort second row in reverse order np.greater(rtt_data[:, :, 0, 0], rtt_data[:, :, 1, 0], out = mask) # sort first column in reverse order
rtt_data[mask, 1] = rtt_data[mask, 1, ::-1] rtt_data[mask, :, 0] = rtt_data[mask, ::-1, 0]
# rtt_data[:, :, :, 1] = rtt_data[:, :, ::-1, 1] # swap second column (not entirely necessary, just makes indices below nicer) np.greater(rtt_data[:, :, 0, 0], rtt_data[:, :, 0, 1], out = mask) # sort first row in reverse order
rtt_data[np.isinf(rtt_data)] = np.nan # restore NaNs rtt_data[mask, 0] = rtt_data[mask, 0, ::-1]
rtt_data[:, :, :, 0] = rtt_data[:, :, ::-1, 0] # swap first column
print("done") print("done")
print("calculating median rtt values...", end = " ", flush = True) print("calculating median rtt values...", end = " ", flush = True)
mask2 = np.empty((rtt_data.shape[0], rtt_data.shape[1]), dtype = np.bool_) # need second mask for binary ops mask2 = np.empty((rtt_data.shape[0], rtt_data.shape[1]), dtype = np.bool_) # need second mask for binary ops
np.invert(np.isnan(rtt_data[:, :, 0, 1], out = mask), out = mask) # four nums populated np.not_equal(rtt_data[:, :, 1, 1], 0, out = mask) # four nums populated
rtt_data[mask, 0, 0] = rtt_data[mask, 1, 1] # take average of index 1 and 2 rtt_data[mask, 0, 1] //= 2
rtt_data[mask, 0, 0] += rtt_data[mask, 1, 0] rtt_data[mask, 1, 0] //= 2
rtt_data[mask, 0, 0] /= 2 rtt_data[mask, 0, 0] = rtt_data[mask, 0, 1]
np.logical_and(np.invert(np.isnan(rtt_data[:, :, 1, 0], out = mask), out = mask), np.isnan(rtt_data[:, :, 0, 1], out = mask2), out = mask) # three nums populated rtt_data[mask, 0, 0] += rtt_data[mask, 1, 0] # take average of middle two nums
rtt_data[mask, 0, 0] = rtt_data[mask, 1, 1] # take index 1 np.logical_and(np.not_equal(rtt_data[:, :, 1, 0], 0, out = mask), np.equal(rtt_data[:, :, 1, 1], 0, out = mask2), out = mask) # three nums populated
np.logical_and(np.invert(np.isnan(rtt_data[:, :, 1, 1], out = mask), out = mask), np.isnan(rtt_data[:, :, 1, 0], out = mask2), out = mask) # two nums populated rtt_data[mask, 0, 0] = rtt_data[mask, 0, 1] # take middle of three nums
rtt_data[mask, 0, 0] = rtt_data[mask, 0, 0] # take average of index 0 and 1 np.logical_and(np.not_equal(rtt_data[:, :, 0, 1], 0, out = mask), np.equal(rtt_data[:, :, 1, 0], 0, out = mask2), out = mask) # two nums populated
rtt_data[mask, 0, 0] += rtt_data[mask, 1, 1] rtt_data[mask, 0, 0] //= 2
rtt_data[mask, 0, 0] /= 2 rtt_data[mask, 0, 1] //= 2
rtt_data[mask, 0, 0] += rtt_data[mask, 0, 1] # take average of first two nums
# everything else (1 or 0 nums populated) don't need any modifications # everything else (1 or 0 nums populated) don't need any modifications
print(f"done (shrunk rtt data from {rtt_data.shape[0] * 2}x{rtt_data.shape[1] * 2} -> {rtt_data.shape[0]}x{rtt_data.shape[1]})") print(f"done (shrunk rtt data from {rtt_data.shape[0] * 2}x{rtt_data.shape[1] * 2} -> {rtt_data.shape[0]}x{rtt_data.shape[1]})")
rtt_data = rtt_data[:, :, 0, 0] rtt_data = rtt_data[:, :, 0, 0]
@ -283,22 +284,30 @@ def make_tiles(coords_path: Path, input_path: Path, tiles_dir: Path, *,
squish() squish()
def get_normalized_data(): def get_normalized_data():
print(f"normalizing rtt data: multiplying...", end = " ", flush = True) print("normalizing rtt data: getting non-zero...", end = " ", flush = True)
rtt_data_f = rtt_data * multiplier non_zero = rtt_data != 0
print(f"incrementing...", end = " ", flush = True) print("converting to floating point...", end = " ", flush = True)
rtt_data_f += 1 rtt_data_f = rtt_data.astype(np.float32)
# print(f"replacing NaNs...", end = " ", flush = True) print("decrementing non-zero...", end = " ", flush = True)
# rtt_data_f[np.isnan(rtt_data_f)] = 0.0 rtt_data_f[non_zero] -= 1
print(f"converting to ints...", end = " ", flush = True) print("dividing...", end = " ", flush = True)
rtt_data_f /= divisor
print("incrementing non-zero...", end = " ", flush = True)
rtt_data_f[non_zero] += 1
del non_zero
collect()
print("converting to ints...", end = " ", flush = True)
with catch_warnings(action = "ignore"): with catch_warnings(action = "ignore"):
rtt_data_norm = rtt_data_f.astype(np.uint16) rtt_data_norm = rtt_data_f.astype(np.uint16)
print("done") print("done")
return rtt_data_norm return rtt_data_norm
def write_all_colormaps(): def write_all_colormaps():
if raws_path is not None:
create_raw_image(rtt_data, raws_path / variant_name)
rtt_data_norm = get_normalized_data() rtt_data_norm = get_normalized_data()
for colormap_name, colormap in colormaps: for colormap_name, colormap in colormaps:
create_images(rtt_data_norm, colormap, num_colors, tiles_dir / "rtt" / colormap_name) create_tile_images(rtt_data_norm, colormap, num_colors, tiles_dir / variant_name / colormap_name)
write_all_colormaps() write_all_colormaps()
while rtt_data.shape[0] > tile_size: while rtt_data.shape[0] > tile_size:
@ -376,24 +385,6 @@ def remove_tiles(tiles_dir: Path, *, json_path: Path | None = None):
rmtree(tiles_dir) rmtree(tiles_dir)
print("done") print("done")
@dataclass
class IpMapArgs:
command: Literal["mkcoords", "convert", "mktiles", "rmtiles"]
quiet: bool
coords: str
input: str
output: str
batches: int
processes: int
tile_size: int
alpha: bool
variants: str
colormaps: str
quantile: float
num_rows: int | None
skip_iters: int | None
json: str | None
def parse_list_arg(arg: str): def parse_list_arg(arg: str):
return [x.strip().lower() for x in arg.split(",") if x.strip()] return [x.strip().lower() for x in arg.split(",") if x.strip()]
@ -413,6 +404,7 @@ def main():
mktiles_parser.add_argument("-a", "--alpha", action = "store_true", help = "use alpha channel instead of black") mktiles_parser.add_argument("-a", "--alpha", action = "store_true", help = "use alpha channel instead of black")
mktiles_parser.add_argument("-v", "--variants", default = ",".join(default_variant_names), help = "a comma separated list of variants to generate (default: %(default)s)") mktiles_parser.add_argument("-v", "--variants", default = ",".join(default_variant_names), help = "a comma separated list of variants to generate (default: %(default)s)")
mktiles_parser.add_argument("-c", "--colormaps", default = ",".join(default_colormap_names), help = "a comma separated list of colormaps to generate (default: %(default)s)") mktiles_parser.add_argument("-c", "--colormaps", default = ",".join(default_colormap_names), help = "a comma separated list of colormaps to generate (default: %(default)s)")
mktiles_parser.add_argument("-r", "--raws", help = "generate images containing the raw data for each selected variant and save them to the provided path (default: none)")
mktiles_parser.add_argument("-q", "--quantile", type = float, default = default_quantile, help = "the quantile to use for scaling data such as rtt (default: %(default)s)") mktiles_parser.add_argument("-q", "--quantile", type = float, default = default_quantile, help = "the quantile to use for scaling data such as rtt (default: %(default)s)")
mktiles_parser.add_argument("-n", "--num-rows", type = int, help = "how many rows to read from the scan data (default: all)") mktiles_parser.add_argument("-n", "--num-rows", type = int, help = "how many rows to read from the scan data (default: all)")
mktiles_parser.add_argument("-s", "--skip-iters", type = int, help = "how many iterations to skip generating images for (default: none)") mktiles_parser.add_argument("-s", "--skip-iters", type = int, help = "how many iterations to skip generating images for (default: none)")
@ -423,7 +415,7 @@ def main():
rmtiles_parser = subparsers.add_parser("rmtiles", help = "remove tile images") rmtiles_parser = subparsers.add_parser("rmtiles", help = "remove tile images")
rmtiles_parser.add_argument("-j", "--json", help = "the path for the json file to store metadata about the tile images (default: none)") rmtiles_parser.add_argument("-j", "--json", help = "the path for the json file to store metadata about the tile images (default: none)")
rmtiles_parser.add_argument("input", help = "the path containing tile images to remove") rmtiles_parser.add_argument("input", help = "the path containing tile images to remove")
args = parser.parse_args(namespace = IpMapArgs) args = parser.parse_args()
try: try:
with redirect_stdout(open(devnull, "w") if args.quiet else stdout): with redirect_stdout(open(devnull, "w") if args.quiet else stdout):
@ -435,8 +427,9 @@ def main():
case "mktiles": case "mktiles":
make_tiles(coords_path = Path(args.coords), input_path = Path(args.input), tiles_dir = Path(args.output), make_tiles(coords_path = Path(args.coords), input_path = Path(args.input), tiles_dir = Path(args.output),
tile_size = args.tile_size, alpha = args.alpha, variant_names = parse_list_arg(args.variants), tile_size = args.tile_size, alpha = args.alpha, variant_names = parse_list_arg(args.variants),
colormap_names = parse_list_arg(args.colormaps), quantile = args.quantile, colormap_names = parse_list_arg(args.colormaps), raws_path = Path(args.raws) if args.raws else None,
num_rows = args.num_rows, skip_iters = args.skip_iters, json_path = Path(args.json) if args.json else None) quantile = args.quantile, num_rows = args.num_rows, skip_iters = args.skip_iters,
json_path = Path(args.json) if args.json else None)
case "rmtiles": case "rmtiles":
remove_tiles(tiles_dir = Path(args.input), json_path = Path(args.json) if args.json else None) remove_tiles(tiles_dir = Path(args.input), json_path = Path(args.json) if args.json else None)
case _: case _: