You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
DroneDetector/src/core/waterfall.py

168 lines
7.3 KiB
Python

This file contains ambiguous Unicode characters!

This file contains ambiguous Unicode characters that may be confused with others in your current locale. If your use case is intentional and legitimate, you can safely ignore this warning. Use the Escape button to highlight these characters.

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