fix: properly read samples in different datatypes

pull/628/head
Luciano Righetti 2023-07-13 11:06:25 +02:00
parent e26bfef477
commit df2183ce54
1 changed files with 41 additions and 33 deletions

View File

@ -30,52 +30,60 @@ moduleinfo = {'version': '0.1', 'author': 'Luciano Righetti',
'module-type': ['expansion']}
def generate_plots(recording, meta_filename):
def get_samples(data_bytes, data_type) -> np.ndarray:
"""
Get samples from bytes.
Source: https://github.com/IQEngine/IQEngine/blob/main/api/rf/samples.py
Parameters
----------
data_bytes : bytes
The bytes to convert to samples.
data_type : str
The data type of the bytes.
Returns
-------
np.ndarray
The samples.
"""
if data_type == "ci16_le" or data_type == "ci16":
samples = np.frombuffer(data_bytes, dtype=np.int16)
samples = samples[::2] + 1j * samples[1::2]
elif data_type == "cf32_le":
samples = np.frombuffer(data_bytes, dtype=np.complex64)
elif data_type == "ci8" or data_type == "i8":
samples = np.frombuffer(data_bytes, dtype=np.int8)
samples = samples[::2] + 1j * samples[1::2]
else:
raise ("Datatype " + data_type + " not implemented")
return samples
def generate_plots(recording, meta_filename, data_bytes):
# FFT plot
filename = meta_filename.replace('.sigmf-data', '')
# snippet from https://gist.github.com/daniestevez/0d519fd4044f3b9f44e170fd619fbb40
NFFT = 2048
N = NFFT * 4096
fs = recording.get_global_info()['core:sample_rate']
x = np.fromfile(recording.data_file, 'int16', count=2*N)
x = x[::2] + 1j * x[1::2]
# f = np.fft.fftshift(np.average(
# np.abs(np.fft.fft(x.reshape(-1, NFFT)))**2, axis=0))
# freq = np.fft.fftshift(np.fft.fftfreq(NFFT, 1/fs))
# plt.figure(figsize=(10, 4))
# plt.plot(1e-6 * freq, 10*np.log10(f))
# plt.title(filename)
# plt.ylabel('PSD (dB)')
# plt.xlabel('Baseband frequency (MHz)')
# fft_buff = io.BytesIO()
# plt.savefig(fft_buff, format='png')
# fft_buff.seek(0)
# fft_png = base64.b64encode(fft_buff.read()).decode('utf-8')
# fft_attr = {
# 'type': 'attachment',
# 'value': filename + '-fft.png',
# 'data': fft_png,
# 'comment': 'FFT plot of the recording'
# }
samples = get_samples(
data_bytes, recording.get_global_info()['core:datatype'])
sample_rate = recording.get_global_info()['core:sample_rate']
# Waterfall plot
# snippet from https://pysdr.org/content/frequency_domain.html#fast-fourier-transform-fft
fft_size = 1024
# // is an integer division which rounds down
num_rows = len(x) // fft_size
num_rows = len(samples) // fft_size
spectrogram = np.zeros((num_rows, fft_size))
for i in range(num_rows):
spectrogram[i, :] = 10 * \
np.log10(np.abs(np.fft.fftshift(
np.fft.fft(x[i*fft_size:(i+1)*fft_size])))**2)
np.fft.fft(samples[i*fft_size:(i+1)*fft_size])))**2)
plt.figure(figsize=(10, 4))
plt.title(filename)
plt.imshow(spectrogram, aspect='auto', extent=[
fs/-2/1e6, fs/2/1e6, 0, len(x)/fs])
sample_rate/-2/1e6, sample_rate/2/1e6, 0, len(samples)/sample_rate])
plt.xlabel("Frequency [MHz]")
plt.ylabel("Time [ms]")
plt.savefig(filename + '-spectrogram.png')
@ -91,7 +99,6 @@ def generate_plots(recording, meta_filename):
'comment': 'Waterfall plot of the recording'
}
# return [fft_attr, waterfall_attr]
return [{'relation': 'waterfall-plot', 'attribute': waterfall_attr}]
@ -176,7 +183,8 @@ def handler(q=False):
# add FFT and waterfall plot
try:
plots = generate_plots(recording, sigmf_data_attr['value'])
plots = generate_plots(
recording, sigmf_data_attr['value'], sigmf_data_bin)
except Exception as e:
logging.exception(e)
return {"error": "Could not generate plots"}