Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+test/audio.wav filter=lfs diff=lfs merge=lfs -text

dataset_preprocess.py

@@ -0,0 +1,231 @@
+"""
+Audio Converter
+
+This script scans a directory for audio files and converts them to WAV format.
+It only processes audio files and skips all other file types.
+
+Usage:
+    python audio_converter.py --input_dir /path/to/audio/files --output_dir /path/to/output
+"""
+
+import os
+import sys
+import argparse
+from pathlib import Path
+from typing import List, Tuple
+import subprocess
+
+
+def print_info(message):
+    print(f"INFO: {message}")
+
+def print_error(message):
+    print(f"ERROR: {message}")
+
+def print_debug(message):
+    if VERBOSE:
+        print(f"DEBUG: {message}")
+
+VERBOSE = False
+
+# Audio formats that can be converted
+AUDIO_FORMATS = {
+    '.mp3', '.m4a', '.aac', '.flac', '.ogg', '.wma', '.aiff', '.ape', '.opus'
+}
+
+
+def check_dependencies() -> bool:
+    """
+    Check if required dependencies are installed.
+    
+    Returns:
+        bool: True if dependencies are met, False otherwise
+    """
+    # Check for ffmpeg
+    try:
+        subprocess.run(
+            ["ffmpeg", "-version"], 
+            stdout=subprocess.PIPE, 
+            stderr=subprocess.PIPE
+        )
+        print_info("ffmpeg is installed.")
+        return True
+    except FileNotFoundError:
+        print_error("ffmpeg is not installed. Please install it before running this script.")
+        return False
+
+
+def scan_directory(directory: str) -> Tuple[List[Path], List[Path]]:
+    """
+    Scan directory for audio files.
+    
+    Args:
+        directory: Path to the directory to scan
+        
+    Returns:
+        Tuple containing lists of audio files and files to skip
+    """
+    audio_files = []
+    skip_files = []
+    
+    dir_path = Path(directory)
+    if not dir_path.exists():
+        raise FileNotFoundError(f"Directory not found: {directory}")
+    
+    for file_path in dir_path.glob('**/*'):
+        if file_path.is_file():
+            file_ext = file_path.suffix.lower()
+            
+            if file_ext in AUDIO_FORMATS:
+                audio_files.append(file_path)
+            elif file_ext == '.wav':
+                # Skip existing WAV files
+                print_debug(f"Skipping existing WAV file: {file_path}")
+                skip_files.append(file_path)
+            else:
+                # Skip non-audio files
+                print_debug(f"Skipping non-audio file: {file_path}")
+                skip_files.append(file_path)
+    
+    print_info(f"Found {len(audio_files)} audio files to convert")
+    print_info(f"Skipping {len(skip_files)} files (WAV or non-audio)")
+    
+    return audio_files, skip_files
+
+
+def convert_audio_to_wav(input_file: Path, output_file: Path) -> bool:
+    """
+    Convert audio file to WAV format using ffmpeg.
+    
+    Args:
+        input_file: Path to input audio file
+        output_file: Path to output WAV file
+        
+    Returns:
+        bool: True if conversion was successful, False otherwise
+    """
+    try:
+        # Ensure output directory exists
+        output_file.parent.mkdir(parents=True, exist_ok=True)
+        
+        cmd = [
+            "ffmpeg",
+            "-y",  # Overwrite output file if it exists
+            "-i", str(input_file),  # Input file
+            "-acodec", "pcm_s16le",  # Output codec (16-bit PCM)
+            "-ar", "44100",  # Sample rate (44.1kHz)
+            "-ac", "1",  # Mono audio (1 channel)
+            str(output_file)  # Output file
+        ]
+        
+        process = subprocess.run(
+            cmd,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE
+        )
+        
+        if process.returncode != 0:
+            print_error(f"Error converting {input_file}: {process.stderr.decode()}")
+            return False
+            
+        print_info(f"Successfully converted {input_file} to WAV")
+        return True
+        
+    except Exception as e:
+        print_error(f"Error converting {input_file}: {str(e)}")
+        return False
+
+
+def process_files(audio_files: List[Path], input_dir: str, output_dir: str, 
+                  preserve_structure: bool = True) -> Tuple[int, int]:
+    """
+    Process all identified audio files for conversion.
+    
+    Args:
+        audio_files: List of audio files to convert
+        input_dir: Input directory path 
+        output_dir: Output directory path
+
+    Returns:
+        Tuple of successful conversions, failed conversions
+    """
+    input_base = Path(input_dir)
+    output_base = Path(output_dir)
+    
+    success_count = 0
+    failure_count = 0
+    
+    # Process audio files
+    for audio_file in audio_files:
+        if preserve_structure:
+            rel_path = audio_file.relative_to(input_base)
+            output_file = output_base / rel_path.with_suffix('.wav')
+        else:
+          
+            output_file = output_base / f"{audio_file.stem}.wav"
+        
+        if convert_audio_to_wav(audio_file, output_file):
+            success_count += 1
+        else:
+            failure_count += 1
+    
+    return success_count, failure_count
+
+
+def parse_arguments() -> argparse.Namespace:
+    """Parse command-line arguments."""
+    parser = argparse.ArgumentParser(description="Convert audio files to WAV format")
+    parser.add_argument('--input_dir', type=str, required=True, 
+                        help='Directory containing files to convert')
+    parser.add_argument('--output_dir', type=str, required=True, 
+                        help='Directory for output WAV files')
+    parser.add_argument('--flat', action='store_true',
+                        help='Don\'t preserve directory structure')
+    parser.add_argument('--verbose', '-v', action='store_true',
+                        help='Enable verbose output')
+    
+    return parser.parse_args()
+
+
+def main():
+    """Main function to run the script."""
+    global VERBOSE
+    
+    try:
+        args = parse_arguments()
+        
+        VERBOSE = args.verbose
+        
+        print_info(f"Input directory: {args.input_dir}")
+        print_info(f"Output directory: {args.output_dir}")
+        
+        if not check_dependencies():
+            print_error("Missing dependencies. Please install required packages.")
+            sys.exit(1)
+        
+        # Create output directory if it doesn't exist
+        os.makedirs(args.output_dir, exist_ok=True)
+        
+        audio_files, skip_files = scan_directory(args.input_dir)
+        
+        # Process files
+        preserve_structure = not args.flat
+        success_count, failure_count = process_files(
+            audio_files, 
+            args.input_dir, 
+            args.output_dir, 
+            preserve_structure
+        )
+        
+        # Print summary
+        print_info(f"Conversion complete!")
+        print_info(f"Successfully converted: {success_count} files")
+        print_info(f"Failed conversions: {failure_count} files")
+        
+    except Exception as e:
+        print_error(f"Error during execution: {str(e)}")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

features.py

@@ -0,0 +1,150 @@
+"""Feature computation for YAMNet."""
+
+import numpy as np
+import tensorflow as tf
+
+
+def waveform_to_log_mel_spectrogram_patches(waveform, params):
+  """Compute log mel spectrogram patches of a 1-D waveform."""
+  with tf.name_scope('log_mel_features'):
+    # waveform has shape [<# samples>]
+
+    # Convert waveform into spectrogram using a Short-Time Fourier Transform.
+    # Note that tf.signal.stft() uses a periodic Hann window by default.
+    window_length_samples = int(
+      round(params.sample_rate * params.stft_window_seconds))
+    hop_length_samples = int(
+      round(params.sample_rate * params.stft_hop_seconds))
+    fft_length = 2 ** int(np.ceil(np.log(window_length_samples) / np.log(2.0)))
+    num_spectrogram_bins = fft_length // 2 + 1
+    if params.tflite_compatible:
+      magnitude_spectrogram = _tflite_stft_magnitude(
+          signal=waveform,
+          frame_length=window_length_samples,
+          frame_step=hop_length_samples,
+          fft_length=fft_length)
+    else:
+      magnitude_spectrogram = tf.abs(tf.signal.stft(
+          signals=waveform,
+          frame_length=window_length_samples,
+          frame_step=hop_length_samples,
+          fft_length=fft_length))
+    # magnitude_spectrogram has shape [<# STFT frames>, num_spectrogram_bins]
+
+    # Convert spectrogram into log mel spectrogram.
+    linear_to_mel_weight_matrix = tf.signal.linear_to_mel_weight_matrix(
+        num_mel_bins=params.mel_bands,
+        num_spectrogram_bins=num_spectrogram_bins,
+        sample_rate=params.sample_rate,
+        lower_edge_hertz=params.mel_min_hz,
+        upper_edge_hertz=params.mel_max_hz)
+    mel_spectrogram = tf.matmul(
+      magnitude_spectrogram, linear_to_mel_weight_matrix)
+    log_mel_spectrogram = tf.math.log(mel_spectrogram + params.log_offset)
+    # log_mel_spectrogram has shape [<# STFT frames>, params.mel_bands]
+
+    # Frame spectrogram (shape [<# STFT frames>, params.mel_bands]) into patches
+    # (the input examples). Only complete frames are emitted, so if there is
+    # less than params.patch_window_seconds of waveform then nothing is emitted
+    # (to avoid this, zero-pad before processing).
+    spectrogram_hop_length_samples = int(
+      round(params.sample_rate * params.stft_hop_seconds))
+    spectrogram_sample_rate = params.sample_rate / spectrogram_hop_length_samples
+    patch_window_length_samples = int(
+      round(spectrogram_sample_rate * params.patch_window_seconds))
+    patch_hop_length_samples = int(
+      round(spectrogram_sample_rate * params.patch_hop_seconds))
+    features = tf.signal.frame(
+        signal=log_mel_spectrogram,
+        frame_length=patch_window_length_samples,
+        frame_step=patch_hop_length_samples,
+        axis=0)
+    # features has shape [<# patches>, <# STFT frames in an patch>, params.mel_bands]
+
+    return log_mel_spectrogram, features
+
+
+def pad_waveform(waveform, params):
+  """Pads waveform with silence if needed to get an integral number of patches."""
+  # In order to produce one patch of log mel spectrogram input to YAMNet, we
+  # need at least one patch window length of waveform plus enough extra samples
+  # to complete the final STFT analysis window.
+  min_waveform_seconds = (
+      params.patch_window_seconds +
+      params.stft_window_seconds - params.stft_hop_seconds)
+  min_num_samples = tf.cast(min_waveform_seconds * params.sample_rate, tf.int32)
+  num_samples = tf.shape(waveform)[0]
+  num_padding_samples = tf.maximum(0, min_num_samples - num_samples)
+
+  # In addition, there might be enough waveform for one or more additional
+  # patches formed by hopping forward. If there are more samples than one patch,
+  # round up to an integral number of hops.
+  num_samples = tf.maximum(num_samples, min_num_samples)
+  num_samples_after_first_patch = num_samples - min_num_samples
+  hop_samples = tf.cast(params.patch_hop_seconds * params.sample_rate, tf.int32)
+  num_hops_after_first_patch = tf.cast(tf.math.ceil(
+          tf.cast(num_samples_after_first_patch, tf.float32) /
+          tf.cast(hop_samples, tf.float32)), tf.int32)
+  num_padding_samples += (
+      hop_samples * num_hops_after_first_patch - num_samples_after_first_patch)
+
+  padded_waveform = tf.pad(waveform, [[0, num_padding_samples]],
+                           mode='CONSTANT', constant_values=0.0)
+  return padded_waveform
+
+
+def _tflite_stft_magnitude(signal, frame_length, frame_step, fft_length):
+  """TF-Lite-compatible version of tf.abs(tf.signal.stft())."""
+  def _hann_window():
+    return tf.reshape(
+      tf.constant(
+          (0.5 - 0.5 * np.cos(2 * np.pi * np.arange(0, 1.0, 1.0 / frame_length))
+          ).astype(np.float32),
+          name='hann_window'), [1, frame_length])
+
+  def _dft_matrix(dft_length):
+    """Calculate the full DFT matrix in NumPy."""
+    # See https://en.wikipedia.org/wiki/DFT_matrix
+    omega = (0 + 1j) * 2.0 * np.pi / float(dft_length)
+    # Don't include 1/sqrt(N) scaling, tf.signal.rfft doesn't apply it.
+    return np.exp(omega * np.outer(np.arange(dft_length), np.arange(dft_length)))
+
+  def _rdft(framed_signal, fft_length):
+    """Implement real-input Discrete Fourier Transform by matmul."""
+    # We are right-multiplying by the DFT matrix, and we are keeping only the
+    # first half ("positive frequencies").  So discard the second half of rows,
+    # but transpose the array for right-multiplication.  The DFT matrix is
+    # symmetric, so we could have done it more directly, but this reflects our
+    # intention better.
+    complex_dft_matrix_kept_values = _dft_matrix(fft_length)[:(
+        fft_length // 2 + 1), :].transpose()
+    real_dft_matrix = tf.constant(
+        np.real(complex_dft_matrix_kept_values).astype(np.float32),
+        name='real_dft_matrix')
+    imag_dft_matrix = tf.constant(
+        np.imag(complex_dft_matrix_kept_values).astype(np.float32),
+        name='imaginary_dft_matrix')
+    signal_frame_length = tf.shape(framed_signal)[-1]
+    half_pad = (fft_length - signal_frame_length) // 2
+    padded_frames = tf.pad(
+        framed_signal,
+        [
+            # Don't add any padding in the frame dimension.
+            [0, 0],
+            # Pad before and after the signal within each frame.
+            [half_pad, fft_length - signal_frame_length - half_pad]
+        ],
+        mode='CONSTANT',
+        constant_values=0.0)
+    real_stft = tf.matmul(padded_frames, real_dft_matrix)
+    imag_stft = tf.matmul(padded_frames, imag_dft_matrix)
+    return real_stft, imag_stft
+
+  def _complex_abs(real, imag):
+    return tf.sqrt(tf.add(real * real, imag * imag))
+
+  framed_signal = tf.signal.frame(signal, frame_length, frame_step)
+  windowed_signal = framed_signal * _hann_window()
+  real_stft, imag_stft = _rdft(windowed_signal, fft_length)
+  stft_magnitude = _complex_abs(real_stft, imag_stft)
+  return stft_magnitude

inference.py

@@ -0,0 +1,515 @@
+"""
+Audio Classification using YAMNet and Custom Models
+A streamlined tool for classifying audio using pre-trained and custom models.
+"""
+
+import os
+import argparse
+import logging
+from pathlib import Path
+from typing import Dict, List, Tuple, Optional, Union, Any
+from dataclasses import dataclass, field
+
+import numpy as np
+import pandas as pd
+import librosa
+import resampy
+import soundfile as sf
+import tensorflow as tf
+from tensorflow.keras.models import load_model
+
+
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+# Suppress TensorFlow warnings
+tf.get_logger().setLevel(logging.ERROR)
+os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
+
+
+@dataclass(frozen=True)
+class YAMNetParams:
+    """Parameters for YAMNet model."""
+    sample_rate: float = 16000.0
+    stft_window_seconds: float = 0.025
+    stft_hop_seconds: float = 0.010
+    mel_bands: int = 64
+    mel_min_hz: float = 125.0
+    mel_max_hz: float = 7500.0
+    log_offset: float = 0.001
+    patch_window_seconds: float = 0.96
+    patch_hop_seconds: float = 0.48
+    num_classes: int = 521
+    conv_padding: str = 'same'
+    batchnorm_center: bool = True
+    batchnorm_scale: bool = False
+    batchnorm_epsilon: float = 1e-4
+    classifier_activation: str = 'sigmoid'
+    tflite_compatible: bool = True
+
+    @property
+    def patch_frames(self) -> int:
+        """Calculate number of frames per patch."""
+        return int(round(self.patch_window_seconds / self.stft_hop_seconds))
+
+    @property
+    def patch_bands(self) -> int:
+        """Get number of mel bands."""
+        return self.mel_bands
+
+
+@dataclass
+class Config:
+    """Configuration for models and processing parameters."""
+
+    yamnet_model_path: str
+    yamnet_classes_path: str
+    model_path: Optional[str] = None
+    custom_classes_path: Optional[str] = None
+    output_dir: str = "results"
+    output_file: str = "classification.txt"
+    
+    # Processing parameters
+    window_length: int = 10  # seconds
+    hop_length: int = 1  # seconds
+    custom_weight_factor: float = 5.0
+    top_k: int = 10  # Number of top predictions to keep
+    
+    # Exclude certain classes
+    excluded_classes: List[str] = field(default_factory=lambda: ["Vehicle"])
+    
+    def __post_init__(self):
+        """Convert paths to absolute paths and ensure output directory exists."""
+        self.yamnet_model_path = os.path.abspath(self.yamnet_model_path)
+        self.yamnet_classes_path = os.path.abspath(self.yamnet_classes_path)
+        
+        if self.model_path:
+            self.model_path = os.path.abspath(self.model_path)
+        
+        if self.custom_classes_path:
+            self.custom_classes_path = os.path.abspath(self.custom_classes_path)
+        
+        # Create output directory
+        os.makedirs(Path(self.output_dir), exist_ok=True)
+    
+    @property
+    def output_path(self) -> str:
+        """Get full path to output file."""
+        return os.path.join(self.output_dir, self.output_file)
+
+    @classmethod
+    def from_args(cls, args: argparse.Namespace) -> 'Config':
+        """Create config from command line arguments."""
+        output_dir = os.path.dirname(args.output) or "results"
+        output_file = os.path.basename(args.output) or "classification.txt"
+        
+        return cls(
+            yamnet_model_path=args.yamnet_model,
+            yamnet_classes_path=args.yamnet_classes,
+            model_path=args.model if os.path.exists(args.model) else None,
+            custom_classes_path=args.custom_classes if os.path.exists(args.custom_classes) else None,
+            output_dir=output_dir,
+            output_file=output_file,
+            window_length=args.window,
+            hop_length=args.hop,
+            custom_weight_factor=args.weight
+        )
+
+
+class AudioClassifier:
+    """Audio classification using YAMNet and custom models."""
+    
+    def __init__(self, config: Config):
+        """Initialize classifier with configuration."""
+        self.config = config
+        self.params = YAMNetParams()
+        
+        # Initialize models
+        self.yamnet_model = None
+        self.model = None
+        self.yamnet_classes = []
+        self.custom_classes = []
+        
+        # Load models
+        self._load_models()
+    
+    def _load_models(self) -> None:
+        """Load YAMNet and custom models."""
+        # Load YAMNet model
+        try:
+            from yamnet import yamnet_frames_model, class_names
+            
+            logger.info(f"Loading YAMNet model from {self.config.yamnet_model_path}")
+            self.yamnet_model = yamnet_frames_model(self.params)
+            self.yamnet_model.load_weights(self.config.yamnet_model_path)
+            
+            logger.info(f"Loading YAMNet classes from {self.config.yamnet_classes_path}")
+            self.yamnet_classes = class_names(self.config.yamnet_classes_path)
+            
+        except ImportError:
+            logger.error("YAMNet module not found. Please install it or provide correct path.")
+            raise
+        except Exception as e:
+            logger.error(f"Failed to load YAMNet model: {e}")
+            raise
+        
+        # Load custom model if available
+        if self.config.model_path:
+            try:
+                logger.info(f"Loading custom model from {self.config.model_path}")
+                self.model = load_model(self.config.model_path)
+                
+                if self.config.custom_classes_path:
+                    logger.info(f"Loading custom classes from {self.config.custom_classes_path}")
+                    self.custom_classes = np.load(self.config.custom_classes_path, allow_pickle=True)
+                
+            except Exception as e:
+                logger.warning(f"Failed to load custom model: {e}")
+                logger.warning("Continuing with YAMNet model only.")
+                self.model = None
+                self.custom_classes = []
+    
+    def classify_file(self, audio_path: str) -> Dict[str, Any]:
+        """Classify audio file and return results."""
+        logger.info(f"Processing audio file: {audio_path}")
+        
+        # Load audio
+        waveform, sr = self._load_audio(audio_path)
+        
+        # Process audio segments
+        logger.info("Processing audio segments...")
+        segments_results = self._process_audio_segments(waveform, sr)
+        
+        # Aggregate results
+        logger.info("Aggregating results...")
+        final_results = self._aggregate_results(segments_results)
+        
+        # Save results
+        if self.config.output_path:
+            self._save_results(final_results)
+        
+        return final_results
+    
+    def _load_audio(self, file_path: str) -> Tuple[np.ndarray, int]:
+        """Load and preprocess audio file."""
+      
+        if not os.path.exists(file_path):
+            raise FileNotFoundError(f"Audio file not found: {file_path}")
+            
+        # Load audio data
+        logger.info(f"Loading audio from {file_path}")
+        wav_data, sr = sf.read(file_path, dtype=np.int16)
+        
+        # Convert to float32 in range [-1.0, 1.0]
+        waveform = wav_data / 32768.0
+        waveform = waveform.astype('float32')
+        
+        # Convert stereo to mono if needed
+        if len(waveform.shape) > 1:
+            logger.info("Converting stereo audio to mono")
+            waveform = np.mean(waveform, axis=1)
+        
+        # Resample if needed
+        if sr != self.params.sample_rate:
+            logger.info(f"Resampling audio from {sr}Hz to {self.params.sample_rate}Hz")
+            waveform = resampy.resample(waveform, sr, self.params.sample_rate)
+            sr = int(self.params.sample_rate)
+        
+        return waveform, sr
+    
+    def _process_audio_segments(self, waveform: np.ndarray, sr: int) -> List[Dict[str, Any]]:
+        """Process audio in segments."""
+        segment_length_samples = int(sr * self.config.window_length)
+        hop_length_samples = int(sr * self.config.hop_length)
+        
+        if segment_length_samples <= 0:
+            raise ValueError(f"Invalid segment length: {self.config.window_length} seconds")
+        
+        segments_results = []
+        
+        # Process each segment
+        total_segments = max(1, (len(waveform) - segment_length_samples + hop_length_samples) // hop_length_samples)
+        for i in range(0, len(waveform) - segment_length_samples + 1, hop_length_samples):
+            segment_idx = i // hop_length_samples + 1
+            logger.debug(f"Processing segment {segment_idx}/{total_segments}")
+            
+            end_idx = min(i + segment_length_samples, len(waveform))
+            window = waveform[i:end_idx]
+            
+            # Get YAMNet predictions
+            yamnet_predictions = self._get_yamnet_predictions(window)
+            
+            # Get custom model predictions if available
+            custom_predictions = None
+            if self.model is not None:
+                custom_predictions = self._get_custom_predictions(window)
+            
+            # Combine predictions
+            combined_results = self._combine_predictions(yamnet_predictions, custom_predictions)
+            
+            # Store results
+            segment_result = {
+                'yamnet_predictions': yamnet_predictions,
+                'custom_predictions': custom_predictions,
+                'combined_predictions': combined_results
+            }
+            
+            segments_results.append(segment_result)
+        
+        return segments_results
+    
+    def _get_yamnet_predictions(self, audio_segment: np.ndarray) -> Dict[str, float]:
+        """Get YAMNet predictions for an audio segment."""
+        try:
+            scores, embeddings, spectrogram = self.yamnet_model(audio_segment)
+            prediction = np.mean(scores, axis=0)
+            
+            # Get top predictions
+            top_indices = np.argsort(prediction)[::-1][:self.config.top_k]
+            top_labels = [self.yamnet_classes[i] for i in top_indices]
+            top_scores = prediction[top_indices]
+            
+            return {label: float(score) for label, score in zip(top_labels, top_scores)}
+            
+        except Exception as e:
+            logger.error(f"Error in YAMNet prediction: {e}")
+            return {}
+    
+    def _get_custom_predictions(self, audio_segment: np.ndarray) -> Dict[str, float]:
+        """Get custom model predictions for an audio segment."""
+        try:
+            # Get YAMNet embeddings first
+            embeddings = self.yamnet_model(audio_segment)[1]
+            
+            # Reshape embeddings for custom model
+            embeddings_reshaped = np.reshape(embeddings, (embeddings.shape[0], -1))
+            
+            # Get predictions from custom model
+            predictions = self.model.predict(embeddings_reshaped, verbose=0)
+            
+            # Calculate mean prediction over time
+            mean_predictions = np.mean(predictions, axis=0)
+            
+            # Get top predictions
+            top_indices = np.argsort(mean_predictions)[::-1][:self.config.top_k]
+            
+            # Check if custom classes are available
+            if len(self.custom_classes) > 0:
+                top_labels = [self.custom_classes[i] for i in top_indices]
+            else:
+                # Use numeric indices as labels if no class names are available
+                top_labels = [f"Class_{i}" for i in top_indices]
+                
+            top_scores = mean_predictions[top_indices]
+            
+            # Normalize scores
+            total_score = np.sum(top_scores)
+            if total_score > 0:
+                top_scores = top_scores / total_score
+            
+            return {label: float(score) for label, score in zip(top_labels, top_scores)}
+            
+        except Exception as e:
+            logger.error(f"Error in custom model prediction: {e}")
+            return {}
+    
+    def _combine_predictions(
+        self, 
+        yamnet_predictions: Dict[str, float], 
+        custom_predictions: Optional[Dict[str, float]]
+    ) -> Dict[str, float]:
+        """Combine predictions from different models."""
+        combined = {}
+        
+        # Add custom predictions with weighting 
+        if custom_predictions:
+            for label, score in custom_predictions.items():
+                combined[label] = score * self.config.custom_weight_factor
+        
+        # Add YAMNet predictions if not already present or if higher score
+        for label, score in yamnet_predictions.items():
+            if label not in self.config.excluded_classes:
+                if label not in combined or score > combined[label]:
+                    combined[label] = score
+        
+        return combined
+    
+    def _aggregate_results(self, segments_results: List[Dict[str, Any]]) -> Dict[str, Any]:
+        """Aggregate results across all segments."""
+        # Initialize aggregated predictions
+        aggregated_predictions = {}
+        
+        # Collect all combined predictions, keeping maximum score per label
+        for segment in segments_results:
+            for label, score in segment['combined_predictions'].items():
+                if label in aggregated_predictions:
+                    aggregated_predictions[label] = max(aggregated_predictions[label], score)
+                else:
+                    aggregated_predictions[label] = score
+        
+        # Process results
+        if aggregated_predictions:
+            # Get top predictions
+            sorted_predictions = sorted(
+                aggregated_predictions.items(),
+                key=lambda x: x[1],
+                reverse=True
+            )[:self.config.top_k]
+            
+            # Create a new dictionary with only the top predictions
+            top_predictions = {label: score for label, score in sorted_predictions}
+            
+            # Normalize scores
+            total_score = sum(top_predictions.values())
+            if total_score > 0:
+                normalized_predictions = {
+                    label: score / total_score 
+                    for label, score in top_predictions.items()
+                }
+            else:
+                # Default to equal probabilities if all scores are 0
+                normalized_predictions = {
+                    label: 1.0 / len(top_predictions) if len(top_predictions) > 0 else 0.0
+                    for label in top_predictions
+                }
+            
+            # Find dominant label
+            dominant_label, dominant_score = max(normalized_predictions.items(), key=lambda x: x[1])
+            dominant_score_percentage = round(dominant_score * 100)
+            
+            # Replace original predictions with normalized ones
+            aggregated_predictions = normalized_predictions
+        else:
+            dominant_label = "Unknown"
+            dominant_score = 0
+            dominant_score_percentage = 0
+        
+        return {
+            'aggregated_predictions': aggregated_predictions,
+            'dominant_label': dominant_label,
+            'dominant_score': dominant_score,
+            'dominant_score_percentage': dominant_score_percentage
+        }
+    
+    def _save_results(self, results: Dict[str, Any]) -> None:
+        """Save classification results to file."""
+        try:
+            with open(self.config.output_path, 'w') as file:
+                file.write("Audio Classification Results\n")
+                file.write("=========================\n\n")
+                file.write(f"Primary Classification: {results['dominant_label']} ({results['dominant_score_percentage']}%)\n\n")
+                
+                # Add detailed breakdown
+                file.write("Classification Details:\n")
+                file.write("-----------------\n")
+                sorted_predictions = sorted(
+                    results['aggregated_predictions'].items(), 
+                    key=lambda x: x[1], 
+                    reverse=True
+                )
+                
+                for label, score in sorted_predictions:
+                    percentage = round(score * 100)
+                    file.write(f"{label}: {percentage}%\n")
+                
+            logger.info(f"Results saved to {self.config.output_path}")
+            
+        except Exception as e:
+            logger.error(f"Error saving results: {e}")
+            raise
+
+
+def main():
+    """Main function to run audio classification."""
+    parser = argparse.ArgumentParser(
+        description='Audio classification using YAMNet and custom models',
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+    
+    # Required arguments
+    parser.add_argument('audio_file', type=str, help='Path to audio file for classification')
+    
+    # Model paths
+    parser.add_argument('--yamnet_model', type=str, default='yamnet/yamnet.h5', 
+                        help='Path to YAMNet model weights')
+    parser.add_argument('--yamnet_classes', type=str, default='yamnet/yamnet_class_map.csv', 
+                        help='Path to YAMNet class names')
+    parser.add_argument('--model', type=str, default='model/model.h5', 
+                        help='Path to custom model (optional)')
+    parser.add_argument('--custom_classes', type=str, default='model/model.npy', 
+                        help='Path to custom class names (optional)')
+    
+    # Processing parameters
+    parser.add_argument('--window', type=int, default=10, 
+                        help='Window length in seconds')
+    parser.add_argument('--hop', type=int, default=1, 
+                        help='Hop length in seconds')
+    parser.add_argument('--weight', type=float, default=5.0,
+                        help='Weighting factor for custom model predictions')
+    
+    # Output options
+    parser.add_argument('--output', type=str, default='results/classification.txt', 
+                        help='Path to output file')
+    
+    # Logging options
+    parser.add_argument('--verbose', action='store_true', 
+                        help='Enable verbose output')
+    parser.add_argument('--debug', action='store_true', 
+                        help='Enable debug logging')
+    
+    args = parser.parse_args()
+    
+    # Configure logging
+    if args.debug:
+        logger.setLevel(logging.DEBUG)
+    elif args.verbose:
+        logger.setLevel(logging.INFO)
+    else:
+        logger.setLevel(logging.WARNING)
+    
+    try:
+        # Create configuration
+        config = Config.from_args(args)
+        
+        # Create classifier
+        classifier = AudioClassifier(config)
+        
+        # Process audio file
+        results = classifier.classify_file(args.audio_file)
+        
+        # Print results
+        print("\nAudio Classification Results")
+        print("=========================")
+        print(f"\nPrimary Classification: {results['dominant_label']} ({results['dominant_score_percentage']}%)")
+        
+        if args.verbose:
+            print("\nTop 10 Predictions:")
+            print("-----------------")
+            sorted_predictions = sorted(
+                results['aggregated_predictions'].items(),
+                key=lambda x: x[1],
+                reverse=True
+            )
+            
+            for label, score in sorted_predictions:
+                percentage = round(score * 100)
+                print(f"{label}: {percentage}%")
+                
+        print(f"\nFull results saved to: {config.output_path}")
+        
+    except Exception as e:
+        logger.error(f"Error: {e}")
+        if args.debug:
+            import traceback
+            traceback.print_exc()
+        return 1
+    
+    return 0
+
+
+if __name__ == '__main__':
+    import sys
+    sys.exit(main())

model/model.h5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fad065a5b7abb72dfe41e9fe1d5bb238c43b174559a9f94d62da217d41f44b6b
+size 12671536

model/model.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0139d6fdecc7e44160ddbe67127e9b24760687be87ca6a030b9ffbccc2dbe126
+size 640

params.py

@@ -0,0 +1,36 @@
+"""Hyperparameters for YAMNet."""
+
+from dataclasses import dataclass
+
+# The following hyperparameters (except patch_hop_seconds) were used to train YAMNet,
+# so expect some variability in performance if you change these. The patch hop can
+# be changed arbitrarily: a smaller hop should give you more patches from the same
+# clip and possibly better performance at a larger computational cost.
+@dataclass(frozen=True)  # Instances of this class are immutable.
+class Params:
+  sample_rate: float = 16000.0
+  stft_window_seconds: float = 0.025
+  stft_hop_seconds: float = 0.010
+  mel_bands: int = 64
+  mel_min_hz: float = 125.0
+  mel_max_hz: float = 7500.0
+  log_offset: float = 0.001
+  patch_window_seconds: float = 0.96
+  patch_hop_seconds: float = 0.48
+
+  @property
+  def patch_frames(self):
+    return int(round(self.patch_window_seconds / self.stft_hop_seconds))
+
+  @property
+  def patch_bands(self):
+    return self.mel_bands
+
+  num_classes: int = 521
+  conv_padding: str = 'same'
+  batchnorm_center: bool = True
+  batchnorm_scale: bool = False
+  batchnorm_epsilon: float = 1e-4
+  classifier_activation: str = 'sigmoid'
+
+  tflite_compatible: bool = True

requirements.txt

@@ -0,0 +1,6 @@
+librosa==0.11.0
+numpy==1.26.4
+pandas==2.2.3
+tensorflow==2.9.0
+resampy==0.4.3
+ffmpeg

test/audio.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ebe25cec891496f3cdee7a8160dffe92f29fee75b2c1ac4424d7922350abfe98
+size 1920044

train.py

@@ -0,0 +1,514 @@
+"""
+Audio Classification System
+
+This module trains a neural network model on audio data using YAMNet embeddings.
+It extracts features from audio files and trains a classifier to recognize audio classes.
+
+Usage: 
+    python main.py --data_path <path_to_data> --model_name <model_name>
+    
+"""
+
+import os
+import sys
+import argparse
+import logging
+from pathlib import Path
+from typing import Tuple, List, Dict, Optional, Any
+
+import numpy as np
+import pandas as pd
+import tensorflow as tf
+import librosa
+from tqdm import tqdm
+from sklearn.preprocessing import LabelBinarizer
+from sklearn.utils import shuffle
+
+
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.StreamHandler() 
+    ]
+)
+logger = logging.getLogger(__name__)
+
+# Default configuration
+DEFAULT_CONFIG = {
+    'yamnet_path': 'yamnet/yamnet.h5',
+    'classes_path': 'yamnet/yamnet_class_map.csv',
+    'sample_rate': 16000,
+    'epochs': 100,
+    'batch_size': 32,
+    'learning_rate': 0.001,
+    'num_hidden': 1024,
+    'hidden_layer_size': 512,
+    'num_extra_layers': 1,
+    'dropout_rate': 0.3,
+    'regularization': 0.01,
+    'patience': 10,
+    'validation_split': 0.2,
+    'model_folder': 'model' 
+}
+
+
+class Configuration:
+    """Handles configuration for the audio classification system."""
+    
+    def __init__(self, custom_config: Optional[Dict[str, Any]] = None):
+        """
+        Initialize configuration handler.
+        
+        Args:
+            custom_config: Custom configuration to override defaults
+        """
+        self.config = DEFAULT_CONFIG.copy()
+        if custom_config:
+            self.config.update(custom_config)
+
+    def get(self, key: str, default: Any = None) -> Any:
+        return self.config.get(key, default)
+    
+    def set(self, key: str, value: Any) -> None:
+        self.config[key] = value
+    
+    def __getitem__(self, key: str) -> Any:
+        return self.config[key]
+
+
+class ClassMap:
+    """Handles audio class mapping and persistence."""
+    
+    def __init__(self, config: Configuration):
+        """
+        Initialize class map.
+        
+        Args:
+            config: Configuration handler
+        """
+        self.config = config
+        self.classes_path = config['classes_path']
+        self._ensure_classes_file_exists()
+    
+    def _ensure_classes_file_exists(self) -> None:
+        """Ensure the classes mapping file exists."""
+        if not os.path.exists(self.classes_path):
+            logger.info(f"Class map file not found: {self.classes_path}. Creating a new one.")
+ 
+            pd.DataFrame({"display_name": [], "index": [], "mid": []}).to_csv(
+                self.classes_path, index=False
+            )
+    
+    def load_yamnet_classes(self) -> np.ndarray:
+        """Load classes from YAMNet class map CSV file."""
+        try:
+            df = pd.read_csv(self.classes_path)
+            return df["display_name"].values
+        except Exception as e:
+            logger.error(f"Error loading classes: {str(e)}")
+            return np.array([])
+    
+    def update_classes(self, data_path: str) -> List[str]:
+        """
+        Update classes based on directory structure.
+        
+        Args:
+            data_path: Path to data directory
+            
+        Returns:
+            List of all class names
+        """
+        try:
+            # Load existing classes mapping
+            existing_classes_df = pd.read_csv(self.classes_path)
+            existing_classes_set = set(existing_classes_df['display_name'])
+            
+            # Find new classes
+            new_classes = []
+            for cls in sorted(os.listdir(data_path)):
+                class_path = os.path.join(data_path, cls)
+                if os.path.isdir(class_path) and cls not in existing_classes_set:
+                    new_classes.append(cls)
+            
+            # Append new classes to the existing classes dataframe
+            if new_classes:
+                logger.info(f"Adding {len(new_classes)} new classes: {', '.join(new_classes)}")
+                new_classes_df = pd.DataFrame({
+                    'display_name': new_classes, 
+                    'index': [''] * len(new_classes), 
+                    'mid': [''] * len(new_classes)
+                })
+                updated_classes_df = pd.concat([existing_classes_df, new_classes_df], ignore_index=True)
+                updated_classes_df.to_csv(self.classes_path, index=False)
+                
+            # Return all classes from data directory
+            return [cls for cls in sorted(os.listdir(data_path)) 
+                    if os.path.isdir(os.path.join(data_path, cls))]
+            
+        except Exception as e:
+            logger.error(f"Error updating classes: {str(e)}")
+            raise
+
+
+class FeatureExtractor:
+    """Extracts features from audio files using YAMNet."""
+    
+    def __init__(self, config: Configuration):
+        """
+        Initialize feature extractor.
+        
+        Args:
+            config: Configuration handler
+        """
+        self.config = config
+        self.yamnet_model = self._load_yamnet_model()
+    
+    def _load_yamnet_model(self):
+        """Load YAMNet model for feature extraction."""
+        try:
+            logger.info("Loading YAMNet model...")
+            # Import here to avoid circular imports
+            from yamnet import yamnet_frames_model
+            from params import Params
+            
+            model = yamnet_frames_model(Params())
+            model.load_weights(self.config['yamnet_path'])
+            return model
+        except Exception as e:
+            logger.error(f"Error loading YAMNet model: {str(e)}")
+            raise
+    
+    def extract_features(self, audio_path: str) -> np.ndarray:
+        """
+        Extract features from an audio file using YAMNet.
+        
+        Args:
+            audio_path: Path to audio file
+            
+        Returns:
+            Numpy array of extracted features
+        """
+        try:
+            # Load audio file
+            wav, _ = librosa.load(
+                audio_path, 
+                sr=self.config['sample_rate'], 
+                mono=True
+            )
+            wav = wav.astype(np.float32)
+            
+            if len(wav) == 0:
+                logger.warning(f"Warning: Empty audio file: {audio_path}")
+                return np.array([])
+            
+            # Extract embeddings using YAMNet
+            _, embeddings, _ = self.yamnet_model(wav)
+            return embeddings.numpy()
+            
+        except Exception as e:
+            logger.error(f"Error extracting features from {audio_path}: {str(e)}")
+            return np.array([])
+
+
+class DatasetLoader:
+    """Creates a dataset from audio files."""
+    
+    def __init__(self, config: Configuration, feature_extractor: FeatureExtractor):
+        """
+        Initialize dataset creator.
+        
+        Args:
+            config: Configuration handler
+            feature_extractor: Feature extractor
+        """
+        self.config = config
+        self.feature_extractor = feature_extractor
+    
+    def create_dataset(self, data_path: str, classes: List[str]) -> Tuple[np.ndarray, np.ndarray]:
+        """
+        Create a dataset from audio files in the specified path.
+        
+        Args:
+            data_path: Path to the directory containing audio files organized in class folders
+            classes: List of class names
+            
+        Returns:
+            samples: Numpy array of audio features
+            labels: Numpy array of corresponding labels
+        """
+        samples, labels = [], []
+        
+        for cls in classes:
+            class_path = os.path.join(data_path, cls)
+            if not os.path.isdir(class_path):
+                continue
+                
+            logger.info(f"Processing class: {cls}")
+            audio_files = os.listdir(class_path)
+            
+            for sound in tqdm(audio_files, desc=f"Processing {cls}"):
+                audio_path = os.path.join(class_path, sound)
+                embeddings = self.feature_extractor.extract_features(audio_path)
+                
+                if len(embeddings) == 0:
+                    continue
+                
+                # Store each embedding frame with its label
+                for embedding in embeddings:
+                    samples.append(embedding)
+                    labels.append(cls)
+        
+        # Convert to numpy arrays
+        if not samples:
+            error_msg = "No valid audio samples were processed!"
+            logger.error(error_msg)
+            raise ValueError(error_msg)
+            
+        samples = np.asarray(samples)
+        labels = np.asarray(labels)
+        
+        logger.info(f"Created dataset with {len(samples)} samples across {len(set(labels))} classes")
+        return samples, labels
+
+
+class ModelBuilder:
+    """Builds and trains neural network models for audio classification."""
+    
+    def __init__(self, config: Configuration):
+        """
+        Initialize model builder.
+        
+        Args:
+            config: Configuration handler
+        """
+        self.config = config
+    
+    def build_model(self, num_classes: int) -> tf.keras.Model:
+        """
+        Build a neural network model for audio classification.
+        
+        Args:
+            num_classes: Number of output classes
+            
+        Returns:
+            Keras Model object
+        """
+        # Input layer (YAMNet embeddings are 1024-dimensional)
+        inputs = tf.keras.layers.Input(shape=(1024,))
+        
+        # First hidden layer with L2 regularization
+        x = tf.keras.layers.Dense(
+            self.config['num_hidden'], 
+            activation='relu',
+            kernel_regularizer=tf.keras.regularizers.l2(self.config['regularization'])
+        )(inputs)
+        x = tf.keras.layers.BatchNormalization()(x)
+        x = tf.keras.layers.Dropout(self.config['dropout_rate'])(x)
+        
+        # Additional hidden layers
+        for i in range(self.config['num_extra_layers']):
+            layer_size = self.config['hidden_layer_size'] // (i+1)
+            x = tf.keras.layers.Dense(
+                layer_size, 
+                activation='relu',
+                kernel_regularizer=tf.keras.regularizers.l2(self.config['regularization'])
+            )(x)
+            x = tf.keras.layers.BatchNormalization()(x)
+            x = tf.keras.layers.Dropout(self.config['dropout_rate'])(x)
+        
+        # Output layer
+        outputs = tf.keras.layers.Dense(num_classes, activation='softmax')(x)
+        
+        # Create and return model
+        model = tf.keras.Model(inputs=inputs, outputs=outputs)
+        return model
+    
+    def _create_callbacks(self, model_path: str) -> List[tf.keras.callbacks.Callback]:
+        """
+        Create callbacks for model training.
+        
+        Args:
+            model_path: Path to save the model
+            
+        Returns:
+            List of callbacks
+        """
+        # Create tensorboard callback
+        log_dir = Path(f"logs/{os.path.basename(model_path)}")
+        log_dir.mkdir(parents=True, exist_ok=True)
+        
+        tensorboard = tf.keras.callbacks.TensorBoard(
+            log_dir=log_dir,
+            histogram_freq=1
+        )
+        
+        # Early stopping callback
+        early_stopping = tf.keras.callbacks.EarlyStopping(
+            monitor='val_accuracy',
+            patience=self.config['patience'],
+            restore_best_weights=True,
+            verbose=1
+        )
+        
+        # Learning rate reduction callback
+        reduce_lr = tf.keras.callbacks.ReduceLROnPlateau(
+            monitor='val_loss',
+            factor=0.5,
+            patience=5,
+            min_lr=0.00001,
+            verbose=1
+        )
+        
+        return [early_stopping, reduce_lr, tensorboard]
+    
+    def train_model(self, X: np.ndarray, y: np.ndarray, model_name: str) -> Tuple[tf.keras.Model, LabelBinarizer]:
+        """
+        Train a model on the provided data.
+        
+        Args:
+            X: Input features
+            y: Target labels
+            model_name: Name of the model
+            
+        Returns:
+            Tuple of (trained model, label encoder)
+        """
+        # Encode the labels (one-hot encoding)
+        encoder = LabelBinarizer()
+        encoded_labels = encoder.fit_transform(y)
+        num_classes = len(encoder.classes_)
+        
+        logger.info(f"Training model with {num_classes} classes: {', '.join(encoder.classes_)}")
+        
+        # Create model
+        model = self.build_model(num_classes=num_classes)
+        
+        # Print model summary
+        model.summary()
+        
+        # Compile model
+        optimizer = tf.keras.optimizers.Adam(learning_rate=self.config['learning_rate'])
+        model.compile(
+            optimizer=optimizer,
+            loss=tf.keras.losses.CategoricalCrossentropy(),
+            metrics=['accuracy']
+        )
+        
+        model_folder = os.path.join(self.config['model_folder'])
+        os.makedirs(model_folder, exist_ok=True)
+       
+        model_path = os.path.join(model_folder, model_name)
+        
+        callbacks = self._create_callbacks(model_path)
+        
+        # Train the model
+        history = model.fit(
+            X, encoded_labels,
+            epochs=self.config['epochs'],
+            batch_size=self.config['batch_size'],
+            validation_split=self.config['validation_split'],
+            callbacks=callbacks,
+            verbose=1
+        )
+        
+        # Save the model and class names
+        model.save(f"{model_path}.h5")
+        np.save(f"{model_path}_classes.npy", encoder.classes_)
+        
+        # Save training history
+        hist_df = pd.DataFrame(history.history)
+        hist_df.to_csv(f"{model_path}_history.csv", index=False)
+        
+        logger.info(f"Model saved as {model_path}.h5")
+        logger.info(f"Class names saved as {model_path}_classes.npy")
+        
+        return model, encoder
+
+
+def parse_arguments() -> argparse.Namespace:
+    """Parse command-line arguments."""
+    parser = argparse.ArgumentParser(description="Train an audio classification model")
+    parser.add_argument('--data_path', type=str, required=True, 
+                        help='Path to the directory containing audio files')
+    parser.add_argument('--model_name', type=str, required=True, 
+                        help='Name for the saved model')
+    parser.add_argument('--config', type=str, 
+                        help='Path to config JSON file (optional)')
+    parser.add_argument('--epochs', type=int, default=DEFAULT_CONFIG['epochs'],
+                        help='Number of training epochs')
+    parser.add_argument('--batch_size', type=int, default=DEFAULT_CONFIG['batch_size'],
+                        help='Batch size for training')
+    parser.add_argument('--learning_rate', type=float, default=DEFAULT_CONFIG['learning_rate'],
+                        help='Initial learning rate')
+    parser.add_argument('--model_folder', type=str, default=DEFAULT_CONFIG['model_folder'],
+                        help='Folder to save the model')
+    
+    return parser.parse_args()
+
+
+def load_custom_config(config_path: Optional[str]) -> Dict[str, Any]:
+    """Load custom configuration from a JSON file."""
+    if not config_path:
+        return {}
+    
+    try:
+        import json
+        with open(config_path, 'r') as f:
+            return json.load(f)
+    except Exception as e:
+        logger.error(f"Error loading config file: {str(e)}")
+        return {}
+
+
+def main():
+    """Main function to run the script."""
+    try:
+     
+        args = parse_arguments()
+        
+        # Load custom configuration
+        custom_config = load_custom_config(args.config)
+        
+
+        custom_config.update({
+            'epochs': args.epochs,
+            'batch_size': args.batch_size,
+            'learning_rate': args.learning_rate,
+            'model_folder': args.model_folder
+        })
+        
+        # Create configuration handler
+        config = Configuration(custom_config)
+        
+        logger.info(f"Data path: {args.data_path}")
+        logger.info(f"Model name: {args.model_name}")
+        logger.info(f"Model folder: {config['model_folder']}")
+        
+        # Initialize components
+        class_map = ClassMap(config)
+        feature_extractor = FeatureExtractor(config)
+        dataset_creator = DatasetLoader(config, feature_extractor)
+        model_builder = ModelBuilder(config)
+        
+        # Update classes and get class list
+        classes = class_map.update_classes(args.data_path)
+        
+        # Create dataset
+        samples, labels = dataset_creator.create_dataset(args.data_path, classes)
+        
+        # Shuffle the data for better training
+        samples, labels = shuffle(samples, labels, random_state=42)
+        
+        # Train model
+        model, encoder = model_builder.train_model(samples, labels, args.model_name)
+        
+        logger.info("Training completed successfully!")
+        
+    except Exception as e:
+        logger.error(f"Error during execution: {str(e)}", exc_info=True)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

yamnet.py

@@ -0,0 +1,164 @@
+"""Core model definition of YAMNet."""
+from tensorflow.keras.models import load_model
+
+import csv
+
+import numpy as np
+import tensorflow as tf
+from tensorflow.keras import Model, layers
+
+import features as features_lib
+
+
+def _batch_norm(name, params):
+    def _bn_layer(layer_input):
+        return layers.BatchNormalization(
+            name=name,
+            center=params.batchnorm_center,
+            scale=params.batchnorm_scale,
+            epsilon=params.batchnorm_epsilon)(layer_input)
+
+    return _bn_layer
+
+
+def _conv(name, kernel, stride, filters, params):
+    def _conv_layer(layer_input):
+        output = layers.Conv2D(name='{}/conv'.format(name),
+                               filters=filters,
+                               kernel_size=kernel,
+                               strides=stride,
+                               padding=params.conv_padding,
+                               use_bias=False,
+                               activation=None)(layer_input)
+        output = _batch_norm('{}/conv/bn'.format(name), params)(output)
+        output = layers.ReLU(name='{}/relu'.format(name))(output)
+        return output
+
+    return _conv_layer
+
+
+def _separable_conv(name, kernel, stride, filters, params):
+    def _separable_conv_layer(layer_input):
+        output = layers.DepthwiseConv2D(name='{}/depthwise_conv'.format(name),
+                                        kernel_size=kernel,
+                                        strides=stride,
+                                        depth_multiplier=1,
+                                        padding=params.conv_padding,
+                                        use_bias=False,
+                                        activation=None)(layer_input)
+        output = _batch_norm('{}/depthwise_conv/bn'.format(name), params)(output)
+        output = layers.ReLU(name='{}/depthwise_conv/relu'.format(name))(output)
+        output = layers.Conv2D(name='{}/pointwise_conv'.format(name),
+                               filters=filters,
+                               kernel_size=(1, 1),
+                               strides=1,
+                               padding=params.conv_padding,
+                               use_bias=False,
+                               activation=None)(output)
+        output = _batch_norm('{}/pointwise_conv/bn'.format(name), params)(output)
+        output = layers.ReLU(name='{}/pointwise_conv/relu'.format(name))(output)
+        return output
+
+    return _separable_conv_layer
+
+
+_YAMNET_LAYER_DEFS = [
+    # (layer_function, kernel, stride, num_filters)
+    (_conv, [3, 3], 2, 32),
+    (_separable_conv, [3, 3], 1, 64),
+    (_separable_conv, [3, 3], 2, 128),
+    (_separable_conv, [3, 3], 1, 128),
+    (_separable_conv, [3, 3], 2, 256),
+    (_separable_conv, [3, 3], 1, 256),
+    (_separable_conv, [3, 3], 2, 512),
+    (_separable_conv, [3, 3], 1, 512),
+    (_separable_conv, [3, 3], 1, 512),
+    (_separable_conv, [3, 3], 1, 512),
+    (_separable_conv, [3, 3], 1, 512),
+    (_separable_conv, [3, 3], 1, 512),
+    (_separable_conv, [3, 3], 2, 1024),
+    (_separable_conv, [3, 3], 1, 1024)
+]
+
+
+def yamnet(features, params):
+    """Define the core YAMNet mode in Keras."""
+    net = layers.Reshape(
+        (params.patch_frames, params.patch_bands, 1),
+        input_shape=(params.patch_frames, params.patch_bands))(features)
+    for (i, (layer_fun, kernel, stride, filters)) in enumerate(_YAMNET_LAYER_DEFS):
+        net = layer_fun('layer{}'.format(i + 1), kernel, stride, filters, params)(net)
+    embeddings = layers.GlobalAveragePooling2D()(net)
+    logits = layers.Dense(units=params.num_classes, use_bias=True)(embeddings)
+    predictions = layers.Activation(activation=params.classifier_activation)(logits)
+    return predictions, embeddings
+
+
+def yamnet_frames_model(params):
+    """Defines the YAMNet waveform-to-class-scores model.
+
+  Args:
+    params: An instance of Params containing hyperparameters.
+
+  Returns:
+    A model accepting (num_samples,) waveform input and emitting:
+    - predictions: (num_patches, num_classes) matrix of class scores per time frame
+    - embeddings: (num_patches, embedding size) matrix of embeddings per time frame
+    - log_mel_spectrogram: (num_spectrogram_frames, num_mel_bins) spectrogram feature matrix
+  """
+    waveform = layers.Input(batch_shape=(None,), dtype=tf.float32)
+    waveform_padded = features_lib.pad_waveform(waveform, params)
+    log_mel_spectrogram, features = features_lib.waveform_to_log_mel_spectrogram_patches(
+        waveform_padded, params)
+    predictions, embeddings = yamnet(features, params)
+    frames_model = Model(
+        name='yamnet_frames', inputs=waveform,
+        outputs=[predictions, embeddings, log_mel_spectrogram])
+    return frames_model
+
+
+def yamnet_transfer(features, params):
+    net = layers.Reshape(
+        (params.patch_frames, params.patch_bands, 1),
+        input_shape=(params.patch_frames, params.patch_bands))(features)
+    for (i, (layer_fun, kernel, stride, filters)) in enumerate(_YAMNET_LAYER_DEFS):
+        net = layer_fun('layer{}'.format(i + 1), kernel, stride, filters, params)(net)
+    embeddings = layers.GlobalAveragePooling2D()(net)
+    return embeddings
+
+
+def yamnet_frames_model_transfer(params, last_layers):
+    """Defines the YAMNet waveform-to-class-scores model.
+
+    Args:
+      params: An instance of Params containing hyperparameters;
+      last_layers: Path to the classifier model.
+    Returns:
+    A model accepting (num_samples,) waveform input and emitting:
+    - predictions: (num_patches, num_classes) matrix of class scores per time frame
+    - embeddings: (num_patches, embedding size) matrix of embeddings per time frame
+    """
+
+    waveform = layers.Input(batch_shape=(None,), dtype=tf.float32)
+    waveform_padded = features_lib.pad_waveform(waveform, params)
+    _, features = features_lib.waveform_to_log_mel_spectrogram_patches(
+        waveform_padded, params)
+    embeddings = yamnet_transfer(features, params)
+    prediction = embeddings
+    last_layers = load_model(last_layers)
+    for layer in last_layers.layers[1:]:
+        prediction = layer(prediction)
+    frames_model = Model(
+        name='yamnet_frames', inputs=waveform,
+        outputs=[prediction, embeddings])
+    return frames_model
+
+
+def class_names(class_map_csv):
+    """Read the class name definition file and return a list of strings."""
+    if tf.is_tensor(class_map_csv):
+        class_map_csv = class_map_csv.numpy()
+    with open(class_map_csv) as csv_file:
+        reader = csv.reader(csv_file)
+        next(reader)  # Skip header
+        return np.array([display_name for (_, _, display_name) in reader])

yamnet_test.py

@@ -0,0 +1,56 @@
+"""Installation test for YAMNet."""
+
+import numpy as np
+import tensorflow as tf
+
+import params
+import yamnet
+
+
+class YAMNetTest(tf.test.TestCase):
+  _params = None
+  _yamnet = None
+  _yamnet_classes = None
+
+  @classmethod
+  def setUpClass(cls):
+    super().setUpClass()
+    cls._params = params.Params()
+    cls._yamnet = yamnet.yamnet_frames_model(cls._params)
+    cls._yamnet.load_weights('yamnet/yamnet.h5')
+    cls._yamnet_classes = yamnet.class_names('yamnet/yamnet_class_map.csv')
+
+  def clip_test(self, waveform, expected_class_name, top_n=10):
+    """Run the model on the waveform, check that expected class is in top-n."""
+    predictions, _, _ = YAMNetTest._yamnet(waveform)
+    clip_predictions = np.mean(predictions, axis=0)
+    top_n_indices = np.argsort(clip_predictions)[-top_n:]
+    top_n_scores = clip_predictions[top_n_indices]
+    top_n_class_names = YAMNetTest._yamnet_classes[top_n_indices]
+    top_n_predictions = list(zip(top_n_class_names, top_n_scores))
+    self.assertIn(expected_class_name, top_n_class_names,
+                  'Did not find expected class {} in top {} predictions: {}'.format(
+                      expected_class_name, top_n, top_n_predictions))
+
+  def testZeros(self):
+    self.clip_test(
+        waveform=np.zeros((int(3 * YAMNetTest._params.sample_rate),)),
+        expected_class_name='Silence')
+
+  def testRandom(self):
+    # Create a numpy random Generator with a fixed seed for repeatability
+    rng = np.random.default_rng(51773)
+    self.clip_test(
+        waveform=rng.uniform(-1.0, +1.0,
+                              (int(3 * YAMNetTest._params.sample_rate),)),
+        expected_class_name='White noise')
+
+  def testSine(self):
+    self.clip_test(
+        waveform=np.sin(2 * np.pi * 440 *
+                        np.arange(0, 3, 1 / YAMNetTest._params.sample_rate)),
+        expected_class_name='Sine wave')
+
+
+if __name__ == '__main__':
+  tf.test.main()