-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathwave_visualizer.py
55 lines (42 loc) · 13.6 KB
/
wave_visualizer.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
import numpy as np
import cupy as cp
import cv2
import matplotlib.pyplot
colormap_icefire = [[179, 224, 216], [178, 223, 216], [176, 222, 215], [175, 221, 215], [173, 219, 214], [171, 218, 214], [169, 217, 214], [167, 215, 213], [165, 214, 213], [162, 212, 212], [160, 210, 212], [157, 209, 211], [154, 207, 211], [151, 205, 210], [148, 203, 210], [146, 201, 209], [143, 199, 209], [140, 198, 208], [137, 196, 208], [134, 194, 208], [131, 192, 207], [128, 190, 207], [125, 188, 207], [122, 187, 207], [119, 185, 206], [116, 183, 206], [113, 181, 206], [110, 179, 206], [108, 177, 206], [105, 176, 205], [102, 174, 205], [99, 172, 205], [97, 170, 205], [94, 168, 205], [91, 166, 205], [89, 164, 205], [86, 162, 205], [84, 161, 205], [82, 159, 205], [79, 157, 205], [77, 155, 205], [75, 153, 206], [73, 151, 206], [71, 149, 206], [69, 147, 206], [68, 145, 206], [66, 143, 206], [65, 140, 206], [64, 138, 206], [63, 136, 206], [62, 134, 206], [61, 132, 206], [61, 130, 205], [61, 127, 205], [60, 125, 205], [60, 123, 204], [60, 121, 203], [60, 118, 203], [61, 116, 202], [61, 114, 201], [61, 112, 200], [62, 109, 198], [62, 107, 197], [63, 105, 195], [64, 103, 194], [65, 100, 192], [65, 98, 190], [66, 96, 187], [67, 94, 185], [67, 92, 183], [68, 90, 180], [68, 88, 177], [69, 86, 174], [69, 85, 171], [69, 83, 168], [70, 81, 165], [70, 79, 162], [70, 78, 158], [69, 76, 155], [69, 75, 151], [69, 73, 148], [68, 72, 144], [68, 70, 141], [67, 69, 137], [66, 67, 134], [66, 66, 130], [65, 65, 127], [64, 63, 123], [63, 62, 120], [62, 61, 116], [61, 60, 113], [60, 59, 109], [59, 57, 106], [58, 56, 103], [57, 55, 99], [55, 54, 96], [54, 53, 93], [53, 52, 90], [52, 50, 87], [51, 49, 84], [50, 48, 81], [48, 47, 78], [47, 46, 75], [46, 45, 72], [45, 44, 70], [44, 43, 67], [43, 42, 65], [42, 41, 62], [41, 40, 60], [40, 39, 57], [39, 38, 55], [38, 37, 53], [37, 37, 51], [37, 36, 49], [36, 35, 47], [35, 35, 45], [35, 34, 44], [34, 33, 42], [34, 33, 41], [33, 32, 39], [33, 32, 38], [33, 32, 37], [33, 31, 36], [33, 31, 35], [33, 31, 35], [34, 30, 34], [34, 30, 33], [34, 30, 33], [35, 30, 32], [36, 30, 32], [36, 30, 32], [37, 30, 32], [38, 30, 32], [39, 30, 32], [40, 30, 32], [41, 30, 32], [42, 30, 33], [44, 31, 33], [46, 31, 34], [47, 31, 34], [49, 31, 35], [51, 32, 35], [53, 32, 36], [55, 32, 37], [57, 33, 38], [59, 33, 38], [61, 33, 39], [63, 34, 40], [65, 34, 41], [67, 35, 42], [70, 35, 43], [72, 36, 44], [74, 36, 45], [77, 37, 46], [79, 37, 47], [82, 38, 48], [84, 38, 49], [87, 39, 50], [90, 39, 51], [92, 40, 52], [95, 40, 53], [98, 40, 54], [100, 41, 55], [103, 41, 56], [106, 42, 57], [109, 42, 58], [111, 42, 59], [114, 43, 60], [117, 43, 60], [120, 43, 61], [123, 44, 62], [126, 44, 63], [129, 44, 63], [131, 44, 64], [134, 45, 64], [137, 45, 65], [140, 45, 65], [143, 46, 65], [146, 46, 65], [149, 46, 66], [152, 47, 66], [155, 47, 66], [158, 48, 66], [160, 48, 66], [163, 49, 65], [166, 49, 65], [169, 50, 65], [172, 51, 64], [174, 52, 64], [177, 53, 63], [180, 54, 63], [182, 55, 62], [185, 56, 62], [187, 57, 61], [190, 58, 61], [192, 60, 60], [195, 61, 59], [197, 63, 59], [199, 65, 58], [201, 66, 57], [203, 68, 57], [206, 70, 56], [208, 72, 55], [209, 74, 55], [211, 76, 54], [213, 78, 54], [215, 81, 54], [217, 83, 53], [218, 85, 53], [220, 88, 53], [221, 90, 53], [223, 93, 54], [224, 95, 54], [225, 98, 55], [227, 101, 55], [228, 103, 56], [229, 106, 57], [230, 109, 58], [231, 111, 60], [232, 114, 61], [233, 117, 62], [234, 120, 64], [235, 123, 66], [236, 125, 68], [237, 128, 70], [237, 131, 73], [238, 134, 75], [239, 137, 78], [240, 139, 80], [240, 142, 83], [241, 145, 86], [242, 148, 89], [242, 151, 93], [243, 153, 96], [243, 156, 99], [244, 159, 103], [245, 162, 106], [245, 165, 110], [246, 167, 113], [246, 170, 117], [247, 173, 120], [247, 176, 124], [248, 178, 127], [248, 181, 131], [249, 184, 134], [249, 186, 138], [250, 188, 141], [250, 190, 144], [251, 192, 147], [251, 194, 149], [251, 196, 152], [252, 198, 154], [252, 200, 156], [252, 201, 158], [253, 203, 160]]
colormap_wave1 = [[255, 255, 255], [254, 254, 253], [254, 253, 252], [253, 252, 250], [253, 250, 248], [252, 249, 246], [252, 248, 244], [251, 246, 242], [251, 245, 240], [250, 243, 237], [250, 242, 235], [249, 240, 232], [248, 238, 230], [248, 237, 227], [247, 235, 224], [247, 233, 221], [246, 231, 218], [245, 229, 215], [245, 227, 212], [244, 225, 209], [243, 223, 206], [242, 221, 203], [242, 219, 200], [241, 217, 196], [240, 215, 193], [239, 213, 190], [239, 211, 186], [238, 208, 183], [237, 206, 179], [236, 204, 176], [235, 202, 172], [234, 199, 169], [233, 197, 165], [232, 195, 162], [231, 192, 158], [230, 190, 155], [230, 188, 151], [228, 185, 148], [227, 183, 144], [226, 181, 141], [225, 178, 137], [224, 176, 134], [223, 174, 130], [222, 171, 127], [221, 169, 124], [219, 167, 120], [218, 164, 117], [217, 162, 114], [216, 160, 111], [214, 157, 108], [213, 155, 105], [212, 153, 102], [210, 151, 99], [209, 149, 96], [207, 146, 93], [206, 144, 91], [204, 142, 88], [203, 140, 85], [201, 138, 83], [199, 136, 80], [198, 134, 78], [196, 132, 76], [194, 130, 74], [193, 128, 72], [191, 127, 70], [189, 125, 69], [187, 123, 67], [185, 121, 65], [183, 119, 63], [180, 117, 62], [178, 115, 60], [176, 113, 58], [173, 111, 56], [170, 109, 55], [168, 107, 53], [165, 105, 52], [162, 102, 50], [159, 100, 49], [156, 98, 47], [154, 96, 46], [151, 94, 44], [148, 92, 43], [144, 90, 41], [141, 87, 40], [138, 85, 39], [135, 83, 37], [132, 81, 36], [129, 79, 35], [125, 76, 34], [122, 74, 33], [119, 72, 31], [115, 70, 30], [112, 68, 29], [109, 66, 28], [105, 64, 27], [102, 62, 27], [99, 60, 26], [96, 58, 25], [93, 56, 25], [89, 54, 25], [86, 52, 25], [83, 51, 25], [80, 49, 25], [77, 47, 25], [74, 45, 25], [71, 44, 25], [68, 42, 25], [65, 41, 25], [62, 39, 25], [60, 38, 25], [57, 37, 25], [54, 35, 25], [52, 34, 25], [49, 33, 25], [47, 32, 25], [45, 31, 25], [43, 30, 25], [40, 29, 25], [39, 28, 25], [37, 28, 25], [35, 27, 25], [33, 27, 25], [32, 26, 25], [30, 26, 25], [29, 25, 25], [28, 25, 25], [27, 25, 25], [26, 25, 25], [26, 25, 26], [26, 26, 27], [26, 26, 28], [26, 26, 30], [26, 27, 31], [26, 27, 33], [26, 28, 34], [26, 29, 36], [26, 30, 38], [26, 31, 40], [26, 32, 42], [26, 33, 44], [26, 34, 47], [26, 35, 49], [26, 37, 51], [26, 38, 54], [26, 40, 56], [26, 41, 59], [26, 43, 62], [26, 44, 64], [26, 46, 67], [26, 48, 70], [26, 50, 73], [27, 51, 76], [28, 53, 79], [28, 55, 82], [29, 57, 85], [30, 59, 88], [31, 61, 91], [32, 64, 94], [33, 66, 97], [35, 68, 101], [36, 70, 104], [37, 72, 107], [38, 74, 110], [40, 77, 113], [41, 79, 117], [42, 81, 120], [44, 84, 123], [45, 86, 126], [47, 88, 130], [48, 91, 133], [50, 93, 136], [51, 95, 139], [53, 98, 142], [54, 100, 145], [56, 102, 148], [58, 104, 151], [59, 107, 154], [61, 109, 157], [63, 111, 160], [64, 114, 163], [66, 116, 165], [68, 118, 168], [70, 120, 171], [71, 122, 173], [73, 125, 176], [75, 127, 178], [77, 129, 181], [78, 131, 183], [80, 133, 185], [82, 135, 187], [84, 136, 189], [86, 138, 191], [87, 140, 193], [89, 142, 194], [91, 144, 196], [93, 146, 198], [96, 147, 199], [98, 149, 201], [100, 151, 203], [103, 153, 204], [105, 155, 206], [108, 157, 207], [110, 160, 209], [113, 162, 210], [116, 164, 212], [118, 166, 213], [121, 168, 214], [124, 170, 216], [127, 172, 217], [130, 174, 218], [133, 176, 219], [136, 178, 221], [139, 180, 222], [142, 183, 223], [145, 185, 224], [148, 187, 225], [152, 189, 226], [155, 191, 227], [158, 193, 228], [161, 195, 230], [164, 197, 231], [168, 200, 231], [171, 202, 232], [174, 204, 233], [177, 206, 234], [180, 208, 235], [183, 210, 236], [187, 212, 237], [190, 214, 238], [193, 216, 239], [196, 218, 239], [199, 220, 240], [202, 222, 241], [205, 223, 242], [208, 225, 242], [211, 227, 243], [214, 229, 244], [217, 231, 245], [219, 232, 245], [222, 234, 246], [225, 236, 247], [227, 237, 247], [230, 239, 248], [232, 240, 248], [234, 242, 249], [237, 243, 250], [239, 245, 250], [241, 246, 251], [243, 247, 251], [245, 249, 252], [247, 250, 252], [249, 251, 253], [251, 252, 253], [252, 253, 254], [254, 254, 254]]
colormap_wave2 = [[255, 255, 255], [253, 254, 254], [252, 254, 253], [250, 253, 252], [248, 253, 252], [246, 252, 251], [244, 252, 250], [242, 251, 249], [240, 251, 247], [237, 250, 246], [235, 250, 245], [232, 249, 244], [230, 248, 243], [227, 248, 242], [224, 247, 240], [221, 247, 239], [218, 246, 238], [215, 245, 236], [212, 245, 235], [209, 244, 234], [206, 243, 232], [203, 242, 231], [200, 242, 229], [196, 241, 228], [193, 240, 226], [190, 239, 225], [186, 239, 223], [183, 238, 222], [179, 237, 220], [176, 236, 218], [172, 235, 217], [169, 234, 215], [165, 233, 213], [162, 232, 212], [158, 231, 210], [155, 231, 208], [151, 230, 206], [148, 228, 205], [144, 227, 203], [141, 226, 201], [137, 225, 199], [134, 224, 198], [130, 223, 196], [127, 222, 194], [124, 221, 192], [120, 219, 190], [117, 218, 188], [114, 217, 187], [111, 216, 185], [108, 214, 183], [105, 213, 181], [102, 212, 179], [99, 210, 177], [96, 209, 176], [93, 207, 174], [91, 206, 172], [88, 204, 170], [85, 203, 168], [83, 201, 166], [80, 199, 164], [78, 198, 163], [76, 196, 161], [74, 194, 159], [72, 193, 157], [70, 191, 155], [69, 189, 153], [67, 187, 152], [65, 185, 149], [63, 183, 147], [62, 180, 145], [60, 178, 143], [59, 175, 141], [57, 173, 138], [56, 170, 136], [54, 167, 133], [53, 164, 131], [52, 161, 128], [50, 158, 125], [49, 155, 123], [48, 152, 120], [47, 149, 117], [46, 146, 115], [45, 143, 112], [43, 140, 109], [42, 136, 106], [41, 133, 103], [40, 130, 101], [39, 126, 98], [39, 123, 95], [38, 119, 92], [37, 116, 89], [36, 112, 86], [35, 109, 84], [35, 106, 81], [34, 102, 78], [33, 99, 75], [33, 95, 73], [32, 92, 70], [31, 89, 67], [31, 86, 65], [30, 82, 62], [30, 79, 60], [29, 76, 57], [29, 73, 55], [28, 70, 52], [28, 67, 50], [28, 64, 48], [27, 61, 46], [27, 58, 44], [27, 55, 42], [26, 53, 40], [26, 50, 38], [26, 48, 37], [26, 46, 35], [26, 43, 34], [26, 41, 32], [26, 39, 31], [26, 37, 30], [26, 35, 29], [26, 34, 28], [26, 32, 27], [26, 31, 26], [26, 29, 26], [26, 28, 25], [26, 27, 25], [26, 27, 25], [26, 26, 25], [26, 25, 25], [26, 25, 26], [26, 25, 26], [26, 25, 27], [26, 25, 27], [26, 25, 28], [27, 25, 30], [27, 25, 31], [27, 26, 32], [28, 27, 34], [28, 27, 35], [28, 28, 37], [29, 29, 39], [29, 30, 41], [30, 31, 43], [30, 32, 45], [31, 34, 48], [31, 35, 50], [32, 36, 53], [32, 38, 55], [33, 40, 58], [33, 41, 61], [34, 43, 64], [35, 45, 66], [36, 47, 69], [36, 49, 73], [37, 51, 76], [38, 53, 79], [39, 55, 82], [39, 57, 85], [40, 59, 89], [41, 61, 92], [42, 64, 95], [43, 66, 99], [44, 68, 102], [45, 71, 105], [46, 73, 109], [47, 76, 112], [48, 78, 116], [49, 80, 119], [50, 83, 122], [52, 86, 126], [53, 88, 129], [54, 91, 133], [55, 93, 136], [57, 96, 139], [58, 98, 143], [59, 100, 146], [60, 103, 149], [62, 105, 152], [63, 108, 155], [65, 110, 158], [66, 113, 161], [68, 115, 164], [69, 117, 167], [71, 120, 170], [72, 122, 173], [74, 124, 175], [75, 126, 178], [77, 128, 180], [79, 131, 183], [80, 133, 185], [82, 134, 187], [84, 136, 189], [86, 138, 191], [87, 140, 193], [89, 142, 194], [91, 144, 196], [93, 146, 198], [96, 147, 199], [98, 149, 201], [100, 151, 203], [103, 153, 204], [105, 155, 206], [108, 157, 207], [110, 160, 209], [113, 162, 210], [116, 164, 212], [118, 166, 213], [121, 168, 214], [124, 170, 216], [127, 172, 217], [130, 174, 218], [133, 176, 219], [136, 178, 221], [139, 180, 222], [142, 183, 223], [145, 185, 224], [148, 187, 225], [152, 189, 226], [155, 191, 227], [158, 193, 228], [161, 195, 230], [164, 197, 231], [168, 200, 231], [171, 202, 232], [174, 204, 233], [177, 206, 234], [180, 208, 235], [183, 210, 236], [187, 212, 237], [190, 214, 238], [193, 216, 239], [196, 218, 239], [199, 220, 240], [202, 222, 241], [205, 223, 242], [208, 225, 242], [211, 227, 243], [214, 229, 244], [217, 231, 245], [219, 232, 245], [222, 234, 246], [225, 236, 247], [227, 237, 247], [230, 239, 248], [232, 240, 248], [234, 242, 249], [237, 243, 250], [239, 245, 250], [241, 246, 251], [243, 247, 251], [245, 249, 252], [247, 250, 252], [249, 251, 253], [251, 252, 253], [252, 253, 254], [254, 254, 254]]
class WaveVisualizer:
def __init__(self, field_colormap, intensity_colormap):
self.field_colormap = field_colormap
self.intensity_colormap = intensity_colormap
self.intensity = None
self.intensity_exp_average_factor = 0.98
self.field = None
def update(self, wave_sim):
self.field = wave_sim.get_field()
if self.intensity is None:
self.intensity = cp.zeros_like(self.field)
t = self.intensity_exp_average_factor
self.intensity = self.intensity*t + (self.field**2)*(1.0-t)
def render_field(self, brightness_scale=1.0):
gray = (cp.clip(self.field*brightness_scale, -1.0, 1.0) * 127 + 127).astype(np.uint8)
img = self.field_colormap[gray].get() if self.field_colormap is not None else gray.get()
return cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
def render_intensity(self, brightness_scale=1.0, exp=0.5):
gray = (cp.clip((self.intensity**exp)*brightness_scale, 0.0, 1.0) * 254.0).astype(np.uint8)
img = self.intensity_colormap[gray].get() if self.intensity_colormap is not None else gray.get()
return cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
def get_colormap_lut(name, invert, black_level=0.0):
if name == 'colormap_icefire':
color_values = np.array(colormap_icefire)/255
elif name == 'colormap_wave1':
color_values = np.array(colormap_wave1)/255
elif name == 'colormap_wave2':
color_values = np.array(colormap_wave2) / 255
else:
colormap = matplotlib.pyplot.get_cmap(name)
color_values = colormap(np.linspace(0, 1, 255))
if invert:
color_values = 1.0-color_values
color_values = color_values*(1.0-black_level)+black_level
return cp.asarray((color_values*255).astype(np.uint8))