127127
128128class HeatMapGenerator (object ):
129129
130- def __init__ (self , image_path , title , ignore_ssids = [], aps = None ):
130+ graphs = {
131+ 'rssi' : 'RSSI (level)' ,
132+ 'quality' : 'iwstats Quality' ,
133+ 'tcp_upload_Mbps' : 'TCP Upload Mbps' ,
134+ 'tcp_download_Mbps' : 'TCP Download Mbps' ,
135+ 'udp_Mbps' : 'UDP Upload Mbps' ,
136+ 'jitter' : 'UDP Jitter (ms)'
137+ }
138+
139+ def __init__ (
140+ self , image_path , title , ignore_ssids = [], aps = None , thresholds = None
141+ ):
131142 self ._ap_names = {}
132143 if aps is not None :
133144 with open (aps , 'r' ) as fh :
134145 self ._ap_names = {
135146 x .upper (): y for x , y in json .loads (fh .read ()).items ()
136147 }
137148 self ._image_path = image_path
149+ self ._layout = None
150+ self ._image_width = 0
151+ self ._image_height = 0
152+ self ._corners = [(0 , 0 ), (0 , 0 ), (0 , 0 ), (0 , 0 )]
138153 self ._title = title
139154 if not self ._title .endswith ('.json' ):
140155 self ._title += '.json'
@@ -143,22 +158,17 @@ def __init__(self, image_path, title, ignore_ssids=[], aps=None):
143158 'Initialized HeatMapGenerator; image_path=%s title=%s' ,
144159 self ._image_path , self ._title
145160 )
146- self ._layout = imread (self ._image_path )
147- self ._image_width = len (self ._layout [0 ])
148- self ._image_height = len (self ._layout ) - 1
149- self ._corners = [
150- (0 , 0 ), (0 , self ._image_height ),
151- (self ._image_width , 0 ), (self ._image_width , self ._image_height )
152- ]
153- logger .debug (
154- 'Loaded image with width=%d height=%d' ,
155- self ._image_width , self ._image_height
156- )
157161 with open (self ._title , 'r' ) as fh :
158162 self ._data = json .loads (fh .read ())
159163 logger .info ('Loaded %d measurement points' , len (self ._data ))
160-
161- def generate (self ):
164+ self .thresholds = {}
165+ if thresholds is not None :
166+ logger .info ('Loading thresholds from: %s' , thresholds )
167+ with open (thresholds , 'r' ) as fh :
168+ self .thresholds = json .loads (fh .read ())
169+ logger .debug ('Thresholds: %s' , self .thresholds )
170+
171+ def load_data (self ):
162172 a = defaultdict (list )
163173 for row in self ._data :
164174 a ['x' ].append (row ['x' ])
@@ -179,6 +189,24 @@ def generate(self):
179189 a ['ap' ].append (ap + '_2.4' )
180190 else :
181191 a ['ap' ].append (ap + '_5G' )
192+ return a
193+
194+ def _load_image (self ):
195+ self ._layout = imread (self ._image_path )
196+ self ._image_width = len (self ._layout [0 ])
197+ self ._image_height = len (self ._layout ) - 1
198+ self ._corners = [
199+ (0 , 0 ), (0 , self ._image_height ),
200+ (self ._image_width , 0 ), (self ._image_width , self ._image_height )
201+ ]
202+ logger .debug (
203+ 'Loaded image with width=%d height=%d' ,
204+ self ._image_width , self ._image_height
205+ )
206+
207+ def generate (self ):
208+ self ._load_image ()
209+ a = self .load_data ()
182210 for x , y in self ._corners :
183211 a ['x' ].append (x )
184212 a ['y' ].append (y )
@@ -195,14 +223,7 @@ def generate(self):
195223 y = np .linspace (0 , self ._image_height , num_y )
196224 gx , gy = np .meshgrid (x , y )
197225 gx , gy = gx .flatten (), gy .flatten ()
198- for k , ptitle in {
199- 'rssi' : 'RSSI (level)' ,
200- 'quality' : 'iwstats Quality' ,
201- 'tcp_upload_Mbps' : 'TCP Upload Mbps' ,
202- 'tcp_download_Mbps' : 'TCP Download Mbps' ,
203- 'udp_Mbps' : 'UDP Upload Mbps' ,
204- 'jitter' : 'UDP Jitter (ms)'
205- }.items ():
226+ for k , ptitle in self .graphs .items ():
206227 self ._plot (
207228 a , k , '%s - %s' % (self ._title , ptitle ), gx , gy , num_x , num_y
208229 )
@@ -300,6 +321,7 @@ def _add_inner_title(self, ax, title, loc, size=None, **kwargs):
300321 return at
301322
302323 def _plot (self , a , key , title , gx , gy , num_x , num_y ):
324+ logger .debug ('Plotting: %s' , key )
303325 pp .rcParams ['figure.figsize' ] = (
304326 self ._image_width / 300 , self ._image_height / 300
305327 )
@@ -313,15 +335,26 @@ def _plot(self, a, key, title, gx, gy, num_x, num_y):
313335 # Render the interpolated data to the plot
314336 pp .axis ('off' )
315337 # begin color mapping
316- norm = matplotlib .colors .Normalize (
317- vmin = min (a [key ]), vmax = max (a [key ]), clip = True
318- )
338+ if 'min' in self .thresholds .get (key , {}):
339+ vmin = self .thresholds [key ]['min' ]
340+ logger .debug ('Using min threshold from thresholds: %s' , vmin )
341+ else :
342+ vmin = min (a [key ])
343+ logger .debug ('Using calculated min threshold: %s' , vmin )
344+ if 'max' in self .thresholds .get (key , {}):
345+ vmax = self .thresholds [key ]['max' ]
346+ logger .debug ('Using max threshold from thresholds: %s' , vmax )
347+ else :
348+ vmax = max (a [key ])
349+ logger .debug ('Using calculated max threshold: %s' , vmax )
350+ norm = matplotlib .colors .Normalize (vmin = vmin , vmax = vmax , clip = True )
319351 mapper = cm .ScalarMappable (norm = norm , cmap = 'RdYlBu_r' )
320352 # end color mapping
321353 image = pp .imshow (
322354 z ,
323355 extent = (0 , self ._image_width , self ._image_height , 0 ),
324- cmap = 'RdYlBu_r' , alpha = 0.5 , zorder = 100
356+ cmap = 'RdYlBu_r' , alpha = 0.5 , zorder = 100 ,
357+ vmin = vmin , vmax = vmax
325358 )
326359 pp .colorbar (image )
327360 pp .imshow (self ._layout , interpolation = 'bicubic' , zorder = 1 , alpha = 1 )
@@ -359,6 +392,8 @@ def parse_args(argv):
359392 help = 'verbose output. specify twice for debug-level output.' )
360393 p .add_argument ('-i' , '--ignore' , dest = 'ignore' , action = 'append' ,
361394 default = [], help = 'SSIDs to ignore from channel graph' )
395+ p .add_argument ('-t' , '--thresholds' , dest = 'thresholds' , action = 'store' ,
396+ type = str , help = 'thresholds JSON file path' )
362397 p .add_argument ('-a' , '--ap-names' , type = str , dest = 'aps' , action = 'store' ,
363398 default = None ,
364399 help = 'If specified, a JSON file mapping AP MAC/BSSID to '
@@ -412,7 +447,8 @@ def main():
412447 set_log_info ()
413448
414449 HeatMapGenerator (
415- args .IMAGE , args .TITLE , ignore_ssids = args .ignore , aps = args .aps
450+ args .IMAGE , args .TITLE , ignore_ssids = args .ignore , aps = args .aps ,
451+ thresholds = args .thresholds
416452 ).generate ()
417453
418454
0 commit comments