dgloss.py

import math
import os
import time
import signal
import shutil
import re

import numpy as np

import parameters
from beamforming.beamforming import beamformer
from log import Log
from post_processing.wav_handler import WavHandler
from mic_array import build_syn_mic_array, real_mic_array
from mic_array import build_syn_mic_array_custom_db
from signal_proccessing.signal_array import SignalArray
from signal_proccessing.stft_handler import STFTHandler
from simplex_mapper.simplex_handler import SimplexHandler
from post_processing.metrics import compute_and_log_metrics
from visualization import plotter
from enum import Enum
class Mode(Enum):
    Real = 0
    Synthetic = 1
    SyntheticWithCustomDB = 2
    
    
class DGLOSS:
    
    def __init__(self, mode: Mode, true_n_speakers: int = -1, rec_path: str = ""):
        self.iter = 1
        self.true_speakers = true_n_speakers
        self.mode = mode
                
        self._setup_dirs()
        if self.mode == Mode.Real:
            self.mic_array = real_mic_array.RealMicArray(parameters.DEVICE)
        elif self.mode == Mode.SyntheticWithCustomDB:
            self.mic_array, _ = build_syn_mic_array_custom_db.build_mic_array(rec_path)
        else:
            self.t_speakers_number = self.true_speakers
            self.mic_array, self.filtered_signals = build_syn_mic_array.build_mic_array(self.t_speakers_number)
            self.filtered_signals_handlers = [
                WavHandler(f"{parameters.RESULTS_PATH}/true_speaker_{i + 1}.wav") for i in range(self.filtered_signals.shape[2])
            ]
            for handler_index, handler in enumerate(self.filtered_signals_handlers):
                handler.append(self.filtered_signals[:, 0, handler_index])
        
        chunk = self.mic_array.chunk
        self.stft_handlers = [STFTHandler(chunk) for _ in range(self.mic_array.channels)]
        self.signal_array = SignalArray(self.stft_handlers, self.mic_array)
        self.Rn = np.zeros((parameters.K, parameters.M, parameters.M))
        self.Rxs = []
        self.wav_handlers: list[WavHandler] = []
        self.simplex_handler = SimplexHandler()
        self.logger = Log().logger
        self.terminate = False
        
        signal.signal(signal.SIGINT, self._signal_handler)
    
    def should_terminate(self) -> bool:
        if self.terminate:
            return True
        
        if self.mode == Mode.Real:
            return False
        
        elif self.iter > self._expected_sim_iterations():
            return True
        
        return False
    
    def fill_buffer(self):
        self.start_time = time.time()
        for _ in range(parameters.FILL_BUFFER_ITERATIONS - 1):
            self.signal_array.update_signals()
    
    def execute_iteration(self):
        self.logger.debug(f"========================================================")
        self.logger.debug(f"Iteration {self.iter}")
        time0 = time.time()
        self.signal_array.update_signals()
        
        stfts = self.signal_array.get_full_signals()
        features = self.signal_array.reshaped_features
        seperated_signals = self._seperate_signals(stfts ,features)
        
        if parameters.PLOT:
            plotter.plot_stft_frame(stfts[:,:,0], self.iter)
            plotter.plot_feature_vector(features, self.iter)

        self._write_to_wav_files(self.wav_handlers, seperated_signals)
        self.logger.debug(f"Iteration {self.iter} took {time.time() - time0:.2f} seconds")
        plotter.close_all_figures()
        self.iter += 1
    
    def excute_speaker_counting_iteration(self):
        self.signal_array.update_signals()
        features = self.signal_array.reshaped_features
        self.simplex_handler.count_speakers(features)
        self.iter += 1
        
    def evaluate_metrics(self):
        if self.true_speakers == -1 or self.terminate:
            self.logger.info("Skipping metrics evaluation for real mic array or termination signal.")
            return {}
        
        end_time = time.time()
        total_time = end_time - self.start_time
        self.logger.info(f"Total time: {total_time:.2f} seconds")
        
        est_re  = re.compile(r"estimated_speaker_(\d+)\.wav$", re.I)
        e_speakers_number = len([f for f in os.listdir(parameters.RESULTS_PATH) if est_re.match(f)])
        if e_speakers_number != self.t_speakers_number:
            self.logger.error(f"Expected {self.t_speakers_number} speakers, but found {e_speakers_number}, can't compute metrics.")
            return {}

        estimated = np.zeros((parameters.FULL_SIGNAL_LENGTH, e_speakers_number))
        reference = np.zeros((parameters.FULL_SIGNAL_LENGTH, e_speakers_number))
        for i, wav_handler in enumerate(self.wav_handlers):
            estimated[:, i] = wav_handler.read()[
                parameters.HOP_SIZE * parameters.OVERLAP_EMPTY_FRAMES:parameters.HOP_SIZE * parameters.OVERLAP_EMPTY_FRAMES + parameters.FULL_SIGNAL_LENGTH
            ]
            wav_handler.close()
        
        for i, wav_handler in enumerate(self.filtered_signals_handlers):
            reference[:, i] = wav_handler.read()
            wav_handler.close()
    
        metrics = compute_and_log_metrics(reference=reference.T, estimated=estimated.T)
        return metrics
    
    def close(self):
        for wav_handler in self.wav_handlers + self.filtered_signals_handlers:
            wav_handler.close()
        
    def _signal_handler(self, sig, frame):
        self.logger.info("\nExecution interrupted by user. Exiting gracefully...")
        self.terminate = True

    def _expected_sim_iterations(self) -> int:
        if self.mode == Mode.Real:
            return -1
        
        return 1 + math.ceil((parameters.FULL_SIGNAL_FRAMES - parameters.N) / parameters.HOPS_PER_ITER)
    
    def _seperate_signals(self, stfts: np.ndarray, features: np.ndarray) -> np.ndarray:

        time0 = time.time()
        estimated_mask, fh2 = self.simplex_handler.map_simplex(stfts, features)
        self.logger.debug(f"map_simplex: {time.time() - time0:.2f} seconds")
        
        time0 = time.time()
        self.logger.debug(f"relabel_mask: {time.time() - time0:.2f} seconds")
        
        J = self.simplex_handler.speakers_number()
        
        time0 = time.time()
        separated, self.Rn, self.Rxs = beamformer(
            signalFT=stfts,
            mask=estimated_mask,
            speakers_number=J,
            fh=fh2,
            prev_Rn=self.Rn,
            prev_Rxs=self.Rxs,
            iter=self.iter,
        )
        self.logger.debug(f"beamformer: {time.time() - time0:.2f} seconds")

        # return separated.reshape(-1, J)
        return separated
    
    def _write_to_wav_files(self, wav_handlers: list[WavHandler], separated: np.ndarray):
        self.logger.debug("Appended to wav files...")
        
        wav_handlers += [
            WavHandler(
                file_path=f"{parameters.RESULTS_PATH}/estimated_speaker_{i + 1}.wav"
            ) for i in range(len(wav_handlers), self.simplex_handler.speakers_number())
        ]
        
        for speaker, wav_handler in enumerate(wav_handlers):
            wav_handler.append(separated[:, speaker] if speaker < separated.shape[1] else np.zeros_like(separated[:, 0]))
            
    def _setup_dirs(self):
        dirs = [parameters.RESULTS_PATH, "plots"]
        plots_subdirs = ["speaker_count", "Emask_plots", "simplex_plots", "relabel_plots", "corr_mat", "stft", "istft", "feature_vec", "pe", "beamforming"]

        for dir in dirs:
            if os.path.exists(dir):
                shutil.rmtree(dir)
            os.makedirs(dir)
        
        if parameters.PLOT:
            for subdir in plots_subdirs:
                filename = os.path.join("plots", subdir)
                os.makedirs(filename)