create_png.py
Go to the documentation of this file.
1 import gdal
2 import sys
3 import numpy as np
4 from matplotlib import cm
5 import cv2
6 
7 # Takes a one channel classification TIFF file as input and converts it to a colorized PNG file
8 # Input: Classification as .tif
9 # Output: Colorized classification panorama as PNG
10 
11 if len(sys.argv) == 3:
12  inname = sys.argv[1]
13  outname = sys.argv[2]
14 
15  ds = gdal.Open(inname, gdal.GA_ReadOnly)
16  if not ds:
17  print("Could not open file ", inname)
18 
19  raster_count = ds.RasterCount
20  x_size = ds.RasterXSize
21  y_size = ds.RasterYSize
22 
23  min_channel = 1
24  max_channel = ds.RasterCount
25 
26  values = np.empty((raster_count, y_size, x_size), dtype=np.int)
27  for i in range(min_channel-1, max_channel):
28  values[i] = np.array(ds.GetRasterBand(i+1).ReadAsArray()).astype(np.int)
29 else:
30 
31  values = np.random.randint(1, 5, (1, 900, 7480))
32  outname = "./colored_classification.png"
33 
34 min_val = np.amin(values)
35 max_val = np.amax(values)
36 
37 num_classes = max_val - min_val + 1
38 
39 #print("Loading color map...")
40 color_map = np.array(cm.get_cmap('viridis').colors)
41 colors = np.empty((num_classes, 4), dtype=np.uint)
42 step = int((len(color_map)-1) / (num_classes-1))
43 for class_i in range(num_classes):
44  for rgb_channel in range(3):
45  colors[class_i, rgb_channel] = int(color_map[step * class_i, rgb_channel] * 256)
46  # alpha channel
47  colors[class_i, 3] = 255
48 
49 # demo
50 #colors[0] = [0, 255, 0, 255] # greenVeg
51 #colors[2] = [191, 117, 0, 255] # Gestein
52 #colors[1] = [0, 0, 255, 255] # Wasser
53 #colors[3] = [102, 102, 102, 255] # Schatten
54 
55 print("Assigning data to colorized array...")
56 shape = values.shape
57 png_values = np.zeros((shape[1], shape[2], 4), dtype=np.uint)
58 # OpenCV requires BGRA channel order
59 for rgba_channel in reversed(range(4)):
60  for y in range(shape[1]):
61  for x in range(shape[2]):
62  png_values[y, x, rgba_channel] = colors[values[0, y, x] - 1, rgba_channel]
63 
64 print("Writing PNG file...")
65 if cv2.imwrite(outname, png_values):
66  print("Done. Data has been written to", outname)
67 else:
68  print("Could not write PNG file.")


lvr2
Author(s): Thomas Wiemann , Sebastian Pütz , Alexander Mock , Lars Kiesow , Lukas Kalbertodt , Tristan Igelbrink , Johan M. von Behren , Dominik Feldschnieders , Alexander Löhr
autogenerated on Wed Mar 2 2022 00:37:23