4 from matplotlib
import cm
11 if len(sys.argv) == 3:
15 ds = gdal.Open(inname, gdal.GA_ReadOnly)
17 print(
"Could not open file ", inname)
19 raster_count = ds.RasterCount
20 x_size = ds.RasterXSize
21 y_size = ds.RasterYSize
24 max_channel = ds.RasterCount
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)
31 values = np.random.randint(1, 5, (1, 900, 7480))
32 outname =
"./colored_classification.png"
34 min_val = np.amin(values)
35 max_val = np.amax(values)
37 num_classes = max_val - min_val + 1
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)
47 colors[class_i, 3] = 255
55 print(
"Assigning data to colorized array...")
57 png_values = np.zeros((shape[1], shape[2], 4), dtype=np.uint)
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]
64 print(
"Writing PNG file...")
65 if cv2.imwrite(outname, png_values):
66 print(
"Done. Data has been written to", outname)
68 print(
"Could not write PNG file.")