lasp/src/lasp/lasp_record.py

333 lines
9.9 KiB
Python

#!/usr/bin/python3.8
# -*- coding: utf-8 -*-
"""
Read data from stream and record sound and video at the same time
"""
import dataclasses
import logging
import os
import time
import h5py
import threading
import numpy as np
from queue import Queue
from .lasp_atomic import Atomic
from .lasp_cpp import (LASP_VERSION_MAJOR, LASP_VERSION_MINOR, InDataHandler,
StreamMgr)
@dataclasses.dataclass
class RecordStatus:
curT: float = 0
done: bool = False
class Recording:
"""
Class used to perform a recording. Recording data can come in from a
different thread, that is supposed to call the `inCallback` method, with
audio data as an argument.
"""
def __init__(
self,
fn: str,
streammgr: StreamMgr,
rectime: float = None,
wait: bool = True,
progressCallback=None,
startDelay: float = 0,
):
"""
Start a recording. Blocks if wait is set to True.
Args:
fn: Filename to record to. Extension is automatically added if not
provided.
stream: AvStream instance to record from. Should have input
channels!
rectime: Recording time [s], None for infinite, in seconds. If set
to None, or np.inf, the recording continues indefintely.
progressCallback: callable that is called with an instance of
RecordStatus instance as argument.
startDelay: Optional delay added before the recording is *actually*
started in [s].
"""
ext = ".h5"
if ext not in fn:
fn += ext
self.smgr = streammgr
self.metadata = None
if startDelay < 0:
raise RuntimeError("Invalid start delay value. Should be >= 0")
self.startDelay = startDelay
# Flag used to indicate that we have passed the start delay
self.startDelay_passed = False
# The amount of seconds (float) that is to be recorded
self.rectime = rectime
# The file name to store data to
self.fn = fn
self.curT_rounded_to_seconds = 0
# Counter of the number of blocks
self.ablockno = 0
# Stop flag, set when recording is finished.
self.stop = Atomic(False)
# Mutex, on who is working with the H5py data
self.queue_mtx = threading.Lock()
self.queue = Queue()
self.progressCallback = progressCallback
try:
# Open the file
self.f = h5py.File(self.fn, "w", 'stdio')
self.f.flush()
except Exception as e:
logging.error(f'Error creating measurement file {e}')
raise
# This flag is used to delete the file on finish(), and can be used
# when a recording is canceled.
self.deleteFile = False
# Try to obtain stream metadata
streamstatus = streammgr.getStreamStatus(StreamMgr.StreamType.input)
if not streamstatus.runningOK():
raise RuntimeError(
"Stream is not running properly. Please first start the stream"
)
self.ad = None
logging.debug("Starting record....")
self.indh = InDataHandler(
streammgr, self.inCallback, self.resetCallback)
if wait:
logging.debug("Stop recording with CTRL-C")
try:
while not self.stop():
self.handleQueue()
time.sleep(0.01)
except KeyboardInterrupt:
logging.debug("Keyboard interrupt on record")
finally:
self.finish()
def resetCallback(self, daq):
"""
Function called with initial stream data.
Important notes: we found that the h5py HDF5 library does not really
work properly in a multi-threaded context. Therefore, this method
SHOULD be called from the same thread, as the one that created self.f.
Which for the current implementation of StreamMgr and friends is true.
Args:
daq: Daq device.
"""
in_ch = daq.enabledInChannels()
blocksize = daq.framesPerBlock()
self.blocksize = blocksize
self.nchannels = daq.neninchannels()
self.fs = daq.samplerate()
f = self.f
f.attrs["LASP_VERSION_MAJOR"] = LASP_VERSION_MAJOR
f.attrs["LASP_VERSION_MINOR"] = LASP_VERSION_MINOR
# Set the bunch of attributes
f.attrs["samplerate"] = daq.samplerate()
f.attrs["nchannels"] = daq.neninchannels()
f.attrs["blocksize"] = blocksize
f.attrs["sensitivity"] = [ch.sensitivity for ch in in_ch]
f.attrs["channelNames"] = [ch.name for ch in in_ch]
# Add the start delay here, as firstFrames() is called right after the
# constructor is called. time.time() returns a floating point
# number of seconds after epoch.
f.attrs["time"] = time.time() + self.startDelay
# In V2, we do not store JSON metadata anymore, but just an enumeration
# index to a physical quantity.
f.attrs["qtys_enum_idx"] = [ch.qty.value for ch in in_ch]
# Measured physical quantity metadata
# This was how it was in LASP version < 1.0
# f.attrs['qtys'] = [ch.qty.to_json() for ch in in_ch]
f.flush()
def firstFrames(self, adata):
"""
Set up the dataset in which to store the audio data. This will create
the attribute `self.ad`
Args:
adata: Numpy array with data from DAQ
"""
# The array data type cannot
# datatype = daq.dataType()
dtype = np.dtype(adata.dtype)
self.ad = self.f.create_dataset(
"audio",
(1, self.blocksize, self.nchannels),
dtype=dtype,
maxshape=(
None, # This means, we can add blocks
# indefinitely
self.blocksize,
self.nchannels,
),
compression="gzip",
)
self.f.flush()
def inCallback(self, adata):
"""
Called from different thread.
"""
if self.stop():
return True
with self.queue_mtx:
self.queue.put(np.copy(adata))
return True
def handleQueue(self):
"""
This method is called when a block of audio data from the stream is
available. It should return either True or False.
When returning False, it will stop the stream.
"""
if self.stop():
logging.debug('Stop flag set, early return in handleQueue')
# Stop flag is raised. We do not add any data anymore.
return True
with self.queue_mtx:
while not self.queue.empty():
adata = self.queue.get()
if self.ad is None:
self.firstFrames(adata)
self.__addTimeData(adata)
def setDelete(self, val: bool):
"""
Set the delete flag. If set, measurement file is deleted at the end of
the recording. Typically used for cleaning up after canceling a
recording.
"""
self.deleteFile = val
def finish(self):
"""
This method should be called to finish and a close a recording file,
remove the queue from the stream, etc.
"""
logging.debug("Recording::finish()")
self.stop <<= True
self.f.flush()
# Remove indata handler, which also should remove callback function
# from StreamMgr. This, however does not have to happen
# instantaneously. For which we have to implement extra mutex
# guards in this class
del self.indh
self.indh = None
# Remove handle to dataset otherwise the h5 file is not closed
# properly.
del self.ad
self.ad = None
try:
# Close the recording file
self.f.close()
del self.f
except Exception as e:
logging.error(f"Error closing file: {e}")
logging.debug("Recording ended")
if self.deleteFile:
self.__deleteFile()
def __deleteFile(self):
"""
Cleanup the recording file.
"""
try:
os.remove(self.fn)
except Exception as e:
logging.error(f"Error deleting file: {self.fn}: {str(e)}")
def __addTimeData(self, indata):
"""
Called by handleQueue() and adds new time data to the storage file.
"""
# logging.debug('Recording::__addTimeData()')
curT = self.ablockno * self.blocksize / self.fs
# Increase the block counter
self.ablockno += 1
if curT < self.startDelay and not self.startDelay_passed:
# Start delay has not been passed
return
elif curT >= 0 and not self.startDelay_passed:
# Start delay passed, switch the flag!
self.startDelay_passed = True
# Reset the audio block counter and the recording time
self.ablockno = 1
curT = 0
ablockno = self.ablockno
recstatus = RecordStatus(curT=curT, done=False)
if self.progressCallback is not None:
self.progressCallback(recstatus)
curT_rounded_to_seconds = int(curT)
if curT_rounded_to_seconds > self.curT_rounded_to_seconds:
self.curT_rounded_to_seconds = curT_rounded_to_seconds
print(f"{curT_rounded_to_seconds}", end="", flush=True)
else:
print(".", end="", flush=True)
if self.rectime is not None and curT > self.rectime:
# We are done!
if self.progressCallback is not None:
recstatus.done = True
self.progressCallback(recstatus)
self.stop <<= True
return
# Add the data to the file, and resize the audio data blocks
self.ad.resize(ablockno, axis=0)
self.ad[ablockno - 1, :, :] = indata
self.f.flush()