|
|
import gc
|
|
|
import os
|
|
|
import copy
|
|
|
from typing import Tuple
|
|
|
import matplotlib.cm as cm
|
|
|
import matplotlib.pyplot as plt
|
|
|
import matplotlib.colors as mcolors
|
|
|
import DroneScanner.utils.utils as utils
|
|
|
from matplotlib.ticker import AutoMinorLocator
|
|
|
|
|
|
|
|
|
class Waterfall:
|
|
|
"""
|
|
|
Класс для вывода графика типа "Водопад"
|
|
|
|
|
|
Атрибуты:
|
|
|
_freqs: Список частот.
|
|
|
_delay: Количество чтений, после которых обновляем водопад.
|
|
|
_size: Размер водопада по оси Y.
|
|
|
_counter_til_init: Счетчик от _size до 0. Отвечает за инициализацию сolors и valuest_to_plot.
|
|
|
_skip_update_counter:
|
|
|
_points_scale: Масштаб точек графика в зависимости от количества частот.
|
|
|
debug_flag: Флаг для отладки класса.
|
|
|
_save_plots:
|
|
|
_cur_plot_file_idx:
|
|
|
_max_plot_files:
|
|
|
_dir_to_save:
|
|
|
colors: Массив цветов точек графика.
|
|
|
values_to_plot: Значения по оси Y.
|
|
|
|
|
|
"""
|
|
|
|
|
|
def __init__(self, freqs: list, delay: int, size: int, debug_flag: bool, save_plots=False, max_plot_files=1000,
|
|
|
dir_to_save='/home/orangepi/plots'):
|
|
|
"""
|
|
|
Инициализирует новый водопад.
|
|
|
|
|
|
:param freqs: Список частот.
|
|
|
:param delay: Количество чтений, после которых обновляем водопад.
|
|
|
:param size: Размер водопада по оси Y.
|
|
|
:param debug_flag: Флаг для отладки класса.
|
|
|
:param save_plots:
|
|
|
:param max_plot_files:
|
|
|
:param dir_to_save:
|
|
|
"""
|
|
|
self._freqs = freqs
|
|
|
self._delay = delay
|
|
|
self._size = size
|
|
|
self._counter_til_init = size - 1
|
|
|
self._skip_update_counter = 1
|
|
|
self._points_scale = 500 / len(self._freqs)
|
|
|
self.debug_flag = debug_flag
|
|
|
self._save_plots = save_plots
|
|
|
self._cur_plot_file_idx = 0
|
|
|
self._max_plot_files = max_plot_files
|
|
|
self._dir_to_save = dir_to_save
|
|
|
self.colors = []
|
|
|
self.values_to_plot = []
|
|
|
|
|
|
if not os.path.exists(self._dir_to_save):
|
|
|
os.makedirs(self._dir_to_save)
|
|
|
|
|
|
def interpolate_color(self, start_color: tuple, end_color: tuple, factor: float) -> tuple:
|
|
|
"""
|
|
|
Интерполяция цвета между start_color и end_color в зависимости от factor.
|
|
|
Чем больше factor, тем ближе end_color, меньше - ближе к start_color.
|
|
|
:param start_color: Тройка (R, G, B), где R, G, B из 0..255.
|
|
|
:param end_color: Тройка (R, G, B), где R, G, B из 0..255.
|
|
|
:param factor: Число от 0 до 1.
|
|
|
:return: Тройка (a, b, c), где 0 <= a,b,c => 1
|
|
|
"""
|
|
|
return tuple((start + (end - start) * factor) / 255 for start, end in zip(start_color, end_color))
|
|
|
|
|
|
def decorate(self, data: list) -> list:
|
|
|
"""
|
|
|
По строчке с датой строим строчку соответствующих цветов в зависимости от элементов из data.
|
|
|
:param data: Строка длиной len(freqs) с характеристиками сигнала (предобработанным сигналом).
|
|
|
:return: Строка длиной len(freqs) с тройками, соответствующие некоторым цветам из (R, G, B).
|
|
|
"""
|
|
|
green = (0, 255, 0)
|
|
|
red = (255, 0, 0)
|
|
|
colored_data = []
|
|
|
|
|
|
for elem in data:
|
|
|
colored_data.append(self.interpolate_color(green, red, elem))
|
|
|
|
|
|
return colored_data
|
|
|
|
|
|
def transform_value_structure(self) -> Tuple[list, list, list]:
|
|
|
"""
|
|
|
Приведение одномерного freqs и двумерных values_to_plot и colors в одномерные длиной len(freqs) x size
|
|
|
для построения по ним водопада.
|
|
|
:return: Координаты точек по х и y и соответствующие цвета, в которые нужно покрасить точки.
|
|
|
"""
|
|
|
x = []
|
|
|
y = []
|
|
|
z = []
|
|
|
for i in range(utils.get_num_columns_of_array(self.values_to_plot)):
|
|
|
for j in range(utils.get_num_rows_of_array(self.values_to_plot) - 1, -1, -1):
|
|
|
x.append(self._freqs[i])
|
|
|
y.append(self.values_to_plot[j][i])
|
|
|
z.append(self.colors[j][i])
|
|
|
return x, y, z
|
|
|
|
|
|
def plot(self) -> None:
|
|
|
"""
|
|
|
Построение водопада при помощи plt.scatter по x,y, с=z, s=_points_scale и автомасштабируемой сеткой.
|
|
|
:return: None.
|
|
|
"""
|
|
|
x, y, colors = self.transform_value_structure()
|
|
|
# colors_hex = [mcolors.to_hex(color) for color in colors]
|
|
|
plt.scatter(x, y, c=colors, s=self._points_scale, marker='s')
|
|
|
|
|
|
# Добавление цветовой шкалы.
|
|
|
plt.colorbar(cm.ScalarMappable(norm=mcolors.Normalize(vmin=0, vmax=1), cmap='RdYlGn_r'), label='Значение')
|
|
|
|
|
|
# Настройка ограничения по ОY
|
|
|
plt.ylim(-self._size / 10, (self._size - 1) + (self._size / 10))
|
|
|
|
|
|
# Настройка заголовков
|
|
|
plt.xlabel('Частоты')
|
|
|
plt.ylabel('Итерации')
|
|
|
plt.title('Waterfall')
|
|
|
|
|
|
# Добавление автомасштабируемой сетки.
|
|
|
plt.grid(which='major', color='gray', linestyle='-', linewidth=0.5)
|
|
|
plt.minorticks_on()
|
|
|
plt.gca().xaxis.set_minor_locator(AutoMinorLocator())
|
|
|
plt.gca().yaxis.set_minor_locator(AutoMinorLocator())
|
|
|
plt.grid(which='minor', color='gray', linestyle=':', linewidth=0.5)
|
|
|
|
|
|
if self._save_plots:
|
|
|
plt.savefig(f'{self._dir_to_save}/waterfall_{self._cur_plot_file_idx}.png')
|
|
|
if self._cur_plot_file_idx == self._max_plot_files:
|
|
|
self._cur_plot_file_idx = 0
|
|
|
else:
|
|
|
self._cur_plot_file_idx += 1
|
|
|
else:
|
|
|
plt.show()
|
|
|
|
|
|
plt.pause(0.001)
|
|
|
plt.close()
|
|
|
gc.collect()
|
|
|
|
|
|
def update(self, data: list) -> None:
|
|
|
"""
|
|
|
Конструирование массивов values_to_plot (значения по Y) и colors (цвета точек) с последующим
|
|
|
обновлением colors (см. документацию).
|
|
|
:param data:
|
|
|
:return: None.
|
|
|
"""
|
|
|
if self._skip_update_counter == self._delay:
|
|
|
data_data = copy.deepcopy(data)
|
|
|
colored_data = self.decorate(data_data)
|
|
|
|
|
|
if self._counter_til_init != -1:
|
|
|
self.values_to_plot.append([self._counter_til_init for _ in range(len(self._freqs))])
|
|
|
self._counter_til_init -= 1
|
|
|
else:
|
|
|
self.colors.pop()
|
|
|
|
|
|
self.colors.insert(0, colored_data)
|
|
|
self.plot()
|
|
|
self._skip_update_counter = 1
|
|
|
|
|
|
else:
|
|
|
self._skip_update_counter += 1
|