node_spectrogram.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
import time
import cv2
import numpy as np
import dearpygui.dearpygui as dpg
import librosa
import traceback

from node_editor.util import dpg_get_value, dpg_set_value
from node.node_abc import DpgNodeABC
from node.basenode import Node as BaseNode
from src.utils.logging import get_logger

# Import STFT-based functions from spectrogram_utils
from node.InputNode.spectrogram_utils import (
    fourier_transformation,
    make_logscale,
    create_spectrogram_from_audio,
    apply_colormap_to_spectrogram,
    REFERENCE_AMPLITUDE
)

logger = get_logger(__name__)


def create_mel_spectrogram(audio_data, sample_rate=22050):
    """Create mel spectrogram using librosa"""
    mel_spec = librosa.feature.melspectrogram(y=audio_data, sr=sample_rate, n_fft=2048, hop_length=512, n_mels=128)
    mel_spec_db = librosa.power_to_db(mel_spec, ref=np.max)
    mel_spec_db_transposed = np.transpose(mel_spec_db)
    spec_image = apply_colormap_to_spectrogram(mel_spec_db_transposed, method='cv2', cmap='INFERNO')
    spec_image = np.flipud(spec_image)
    return spec_image


def create_stft_spectrogram(audio_data, sample_rate=22050):
    """Create STFT spectrogram using librosa"""
    stft = librosa.stft(audio_data, n_fft=2048, hop_length=512)
    stft_db = librosa.amplitude_to_db(np.abs(stft), ref=np.max)
    stft_db_transposed = np.transpose(stft_db)
    spec_image = apply_colormap_to_spectrogram(stft_db_transposed, method='cv2', cmap='VIRIDIS')
    spec_image = np.flipud(spec_image)
    return spec_image


def create_chromagram(audio_data, sample_rate=22050):
    """Create chromagram using librosa"""
    chroma = librosa.feature.chroma_stft(y=audio_data, sr=sample_rate, n_fft=2048, hop_length=512)
    chroma_transposed = np.transpose(chroma)
    spec_image = apply_colormap_to_spectrogram(chroma_transposed, method='cv2', cmap='PLASMA')
    spec_image = np.flipud(spec_image)
    return spec_image


def create_mfcc(audio_data, sample_rate=22050):
    """Create MFCC using librosa"""
    mfcc = librosa.feature.mfcc(y=audio_data, sr=sample_rate, n_fft=2048, hop_length=512, n_mfcc=20)
    mfcc_transposed = np.transpose(mfcc)
    spec_image = apply_colormap_to_spectrogram(mfcc_transposed, method='cv2', cmap='JET')
    spec_image = np.flipud(spec_image)
    return spec_image


def create_stft_custom(audio_data, sample_rate=22050, binsize=1024, colormap="jet"):
    """Create STFT spectrogram using custom fourier_transformation method"""
    return create_spectrogram_from_audio(audio_data, sample_rate, binsize, colormap)


class FactoryNode:
    node_label = 'Spectrogram'
    node_tag = 'Spectrogram'
    
    def __init__(self):
        pass

    def add_node(
        self,
        parent,
        node_id,
        pos=[0, 0],
        opencv_setting_dict=None,
        callback=None,
    ):
        node = Node()
        node.tag_node_name = str(node_id) + ':' + self.node_tag
        node.tag_node_input01_name = node.tag_node_name + ':' + node.TYPE_AUDIO + ':Input01'
        node.tag_node_input01_value_name = node.tag_node_name + ':' + node.TYPE_AUDIO + ':Input01Value'
        node.tag_node_input02_name = node.tag_node_name + ':' + node.TYPE_TEXT + ':Input02'
        node.tag_node_input02_value_name = node.tag_node_name + ':' + node.TYPE_TEXT + ':Input02Value'
        node.tag_node_output01_name = node.tag_node_name + ':' + node.TYPE_IMAGE + ':Output01'
        node.tag_node_output01_value_name = node.tag_node_name + ':' + node.TYPE_IMAGE + ':Output01Value'
        node.tag_node_output02_name = node.tag_node_name + ':' + node.TYPE_TIME_MS + ':Output02'
        node.tag_node_output02_value_name = node.tag_node_name + ':' + node.TYPE_TIME_MS + ':Output02Value'

        node._opencv_setting_dict = opencv_setting_dict
        small_window_w = node._opencv_setting_dict['process_width']
        small_window_h = node._opencv_setting_dict['process_height']
        use_pref_counter = node._opencv_setting_dict['use_pref_counter']

        # Create black texture for initial display
        black_image = np.zeros((small_window_h, small_window_w, 3))
        black_texture = node.convert_cv_to_dpg(
            black_image,
            small_window_w,
            small_window_h,
        )

        # Register texture
        with dpg.texture_registry(show=False):
            dpg.add_raw_texture(
                small_window_w,
                small_window_h,
                black_texture,
                tag=node.tag_node_output01_value_name,
                format=dpg.mvFormat_Float_rgb,
            )

        # Create node UI
        with dpg.node(
                tag=node.tag_node_name,
                parent=parent,
                label=self.node_label,
                pos=pos,
        ):
            # Audio input
            with dpg.node_attribute(
                    tag=node.tag_node_input01_name,
                    attribute_type=dpg.mvNode_Attr_Input,
            ):
                dpg.add_text(
                    tag=node.tag_node_input01_value_name,
                    default_value='Input Audio',
                )

            # Method selector
            with dpg.node_attribute(
                    tag=node.tag_node_input02_name,
                    attribute_type=dpg.mvNode_Attr_Static,
            ):
                dpg.add_combo(
                    items=['mel', 'stft', 'stft_custom', 'chromagram', 'mfcc'],
                    default_value='mel',
                    width=small_window_w,
                    label="Method",
                    tag=node.tag_node_input02_value_name,
                    callback=callback,
                )

            # Image output
            with dpg.node_attribute(
                    tag=node.tag_node_output01_name,
                    attribute_type=dpg.mvNode_Attr_Output,
            ):
                dpg.add_image(node.tag_node_output01_value_name)

            # Performance counter
            if use_pref_counter:
                with dpg.node_attribute(
                        tag=node.tag_node_output02_name,
                        attribute_type=dpg.mvNode_Attr_Output,
                ):
                    dpg.add_text(
                        tag=node.tag_node_output02_value_name,
                        default_value='elapsed time(ms)',
                    )

        return node


class Node(BaseNode):
    _ver = '0.0.1'

    node_label = 'Spectrogram'
    node_tag = 'Spectrogram'

    _opencv_setting_dict = None

    def __init__(self):
        pass

    def update(
        self,
        node_id,
        connection_list,
        node_image_dict,
        node_result_dict,
        node_audio_dict,
    ):
        tag_node_name = str(node_id) + ':' + self.node_tag
        input_value02_tag = tag_node_name + ':' + self.TYPE_TEXT + ':Input02Value'
        output_value01_tag = tag_node_name + ':' + self.TYPE_IMAGE + ':Output01Value'
        output_value02_tag = tag_node_name + ':' + self.TYPE_TIME_MS + ':Output02Value'

        # Handle case when _opencv_setting_dict is None
        if self._opencv_setting_dict is None:
            small_window_w = 240
            small_window_h = 135
            use_pref_counter = False
        else:
            small_window_w = self._opencv_setting_dict['process_width']
            small_window_h = self._opencv_setting_dict['process_height']
            use_pref_counter = self._opencv_setting_dict['use_pref_counter']

        # Get the selected method
        try:
            method = dpg_get_value(input_value02_tag)
        except Exception as e:
            logger.debug(f"Could not get method value from DPG: {e}")
            method = 'mel'  # Default method if dpg is not available

        # Get audio input
        audio_data = None
        sample_rate = 22050  # Default sample rate
        
        for connection_info in connection_list:
            connection_type = connection_info[0].split(':')[2]
            if connection_type == self.TYPE_AUDIO:
                connection_info_src = ':'.join(connection_info[0].split(':')[:2])
                audio_dict_entry = node_audio_dict.get(connection_info_src, None)
                if audio_dict_entry is not None:
                    # Handle dictionary format from Video node
                    if isinstance(audio_dict_entry, dict):
                        audio_data = audio_dict_entry.get('data', None)
                        if audio_data is None:
                            logger.warning("Audio dictionary missing 'data' key")
                        sample_rate = audio_dict_entry.get('sample_rate', 22050)
                    # Handle legacy tuple format for backward compatibility
                    elif isinstance(audio_dict_entry, (list, tuple)) and len(audio_dict_entry) == 2:
                        audio_data, sample_rate = audio_dict_entry
                    else:
                        logger.warning(f"Unexpected audio data format: {type(audio_dict_entry)}, expected dict or tuple")
                break

        frame = None
        
        if audio_data is not None and use_pref_counter:
            start_time = time.monotonic()

        if audio_data is not None:
            try:
                # Create spectrogram based on selected method
                if method == 'mel':
                    frame = create_mel_spectrogram(audio_data, sample_rate)
                elif method == 'stft':
                    frame = create_stft_spectrogram(audio_data, sample_rate)
                elif method == 'chromagram':
                    frame = create_chromagram(audio_data, sample_rate)
                elif method == 'mfcc':
                    frame = create_mfcc(audio_data, sample_rate)
                elif method == 'stft_custom':
                    frame = create_stft_custom(audio_data, sample_rate, binsize=1024, colormap="jet")
                else:
                    # Default to mel
                    frame = create_mel_spectrogram(audio_data, sample_rate)
            except Exception as e:
                logger.error(f"Error creating {method} spectrogram: {e}", exc_info=True)
                frame = None

        if frame is not None and use_pref_counter:
            elapsed_time = time.monotonic() - start_time
            elapsed_time = int(elapsed_time * 1000)
            try:
                dpg_set_value(output_value02_tag, str(elapsed_time).zfill(4) + 'ms')
            except Exception as e:
                logger.debug(f"Could not set performance counter value: {e}")

        if frame is not None:
            try:
                texture = self.convert_cv_to_dpg(
                    frame,
                    small_window_w,
                    small_window_h,
                )
                dpg_set_value(output_value01_tag, texture)
            except Exception as e:
                logger.debug(f"Could not set output texture: {e}")

        return {"image": frame, "json": None, "audio": None}

    def close(self, node_id):
        pass

    def get_setting_dict(self, node_id):
        tag_node_name = str(node_id) + ':' + self.node_tag
        input_value02_tag = tag_node_name + ':' + self.TYPE_TEXT + ':Input02Value'

        pos = dpg.get_item_pos(tag_node_name)
        method = dpg_get_value(input_value02_tag)

        setting_dict = {}
        setting_dict['ver'] = self._ver
        setting_dict['pos'] = pos
        setting_dict['method'] = method

        return setting_dict

    def set_setting_dict(self, node_id, setting_dict):
        tag_node_name = str(node_id) + ':' + self.node_tag
        input_value02_tag = tag_node_name + ':' + self.TYPE_TEXT + ':Input02Value'

        method = setting_dict.get('method', 'mel')
        dpg_set_value(input_value02_tag, method)