test.py

"""
Simple polling-based folder monitor without watchdog dependency.
"""

import os
import time
import shutil
import numpy as np
import tensorflow as tf
import librosa
from pathlib import Path
import json
from datetime import datetime
import csv

# STFT params (tune later)
FRAME_LENGTH = 256
FRAME_STEP = 128
FFT_LENGTH = FRAME_LENGTH  # keep same as frame length for simplicity
CLIP_LEN_S = 4.0
TARGET_SR = 4000  # standardize to 4kHz
EXPECTED_SAMPLES = int(TARGET_SR * CLIP_LEN_S)
NYQUIST_HZ = TARGET_SR // 2  # 2000 for 4kHz


class SimpleGunshotMonitor:
    """Simple polling-based gunshot detection monitor."""
    
    def __init__(self, input_folder, model_path, params_path=None, 
                 threshold=0.5, poll_interval=2):
        self.input_folder = Path(input_folder)
        self.poll_interval = poll_interval
        self.threshold = threshold
        
        # Create input folder if needed
        self.input_folder.mkdir(parents=True, exist_ok=True)
        
        # Setup output folders
        self.output_folder = self.input_folder.parent / "processed"
        self.gunshot_folder = self.output_folder / "gunshots"
        self.non_gunshot_folder = self.output_folder / "non_gunshots"
        
        for folder in [self.output_folder, self.gunshot_folder, self.non_gunshot_folder]:
            folder.mkdir(exist_ok=True)
        
        # Load model
        self.model = tf.keras.models.load_model(model_path)
        
        # Load params or use defaults
        if params_path and os.path.exists(params_path):
            with open(params_path, 'r') as f:
                params = json.load(f)
                self.target_sr = params.get('target_sr', 4000)
                self.clip_len_s = params.get('clip_len_s', 4.0)
                self.frame_length = params.get('frame_length', 256)
                self.frame_step = params.get('frame_step', 128)
                self.fft_length = params.get('fft_length', 256)
                self.spec_shape = params.get('spec_shape', [124, 129,1])
        else:
            # Default values matching your original code
            self.target_sr = 4000
            self.clip_len_s = 4.0
            self.frame_length = 256
            self.frame_step = 128
            self.fft_length = 256
            self.spec_shape = [124, 129, 1]
        
        self.expected_samples = int(self.target_sr * self.clip_len_s)
        
        # Track processed files
        self.processed_files = set()
        
        # Print model input shape for debugging
        print(f"Model input shape: {self.model.input_shape}")
        print(f"Monitoring: {self.input_folder}")
        print(f"Check interval: {poll_interval} seconds")
        print(f"Threshold: {threshold}")
        print(f"Target SR: {self.target_sr} Hz")
        print(f"Clip length: {self.clip_len_s} seconds")
        print(f"Expected samples: {self.expected_samples}")
        print("Press Ctrl+C to stop\n")

    def parse_tfrecord(self,spec: tf.Tensor):
        """
        Process spectrogram output from wav_to_logspec() to match 
        the format expected by the trained model.
        
        Args:
            spec: Output from wav_to_logspec(), shape [T, F, 1]
        
        Returns:
            Processed spectrogram ready for model inference, shape [1, 124, 129, 1]
            (with batch dimension added)
        """
        # Ensure it has the channel dimension
        spec = (spec - tf.reduce_min(spec)) / (tf.reduce_max(spec) - tf.reduce_min(spec) + 1e-8)


        if tf.rank(spec) == 2:
            spec = tf.expand_dims(spec, axis=-1)
        
        # Set expected shape
        spec.set_shape(self.spec_shape)  # [124, 129, 1]
        
        # Add batch dimension for model input
        spec = tf.expand_dims(spec, axis=0)  # [1, 124, 129, 1]
        
        return spec

    def wav_to_logspec(self, wav_path: str):
        """
        Reads a WAV clip, forces SR to 4kHz (if 8kHz, downsample by 2),
        pads/trims to 4s, returns log-magnitude STFT spectrogram [T, F, 1].
        T is time frames, F is frequency bins.
        """
        wav_bytes = tf.io.read_file(wav_path)

        # Decode WAV to mono float32 samples in range [-1, 1];
        # returns audio with shape [N, 1] and sr (sample rate) as a scalar tensor.
        audio, sr = tf.audio.decode_wav(wav_bytes, desired_channels=1)
        
        # Remove single channel dimension ([N, 1] → [N])
        audio = tf.squeeze(audio, axis=-1)

        sr = tf.cast(sr, tf.int32)

        # Standardize sample rate to 4kHz:
        # - if 8kHz: take every other sample (2x downsample)
        # - if 4kHz: leave as-is
        if sr == 8000:
            audio = audio[::2]
            sr = 4000
            print("-------->>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>")
        elif sr != 4000:
            raise ValueError(f"Unsupported sample rate {int(sr.numpy())} for {wav_path}")

        # # Force fixed length (4s @ 4kHz)
        # n = tf.shape(audio)[0]
        # if n < EXPECTED_SAMPLES:
        #     audio = tf.pad(audio, [[0, EXPECTED_SAMPLES - n]])
        # else:
        #     audio = audio[:EXPECTED_SAMPLES]

        # Compute short-time Fourier transform to obtain a time–frequency representation
        # of the audio signal (outputs complex numbers)
        stft = tf.signal.stft(
            audio,
            frame_length=FRAME_LENGTH,
            frame_step=FRAME_STEP,
            fft_length=FFT_LENGTH,
            pad_end=False,
        )

        # Convert complex STFT to magnitude spectrogram shape [T, F],
        # then apply log scaling to compress dynamic range for CNN stability
        mag = tf.abs(stft)
        mag = tf.math.log1p(mag) # log(1 + mag)

        # # Optional frequency mask for spectrograms
        # if cfg.get("mask", False):
        #     # Frequency bins for rfft: 0 to sr/2 with (FFT_LENGTH//2 + 1) bins
        #     freqs = tf.linspace(0.0, float(TARGET_SR) / 2.0, FFT_LENGTH // 2 + 1) # [F]
            
        #     # Create mask for frequencies to keep
        #     keep = tf.logical_and(freqs >= cfg["low_hz"], freqs <= cfg["high_hz"]) # [F]
        #     keep = tf.cast(keep, mag.dtype)

        #     # Apply the frequency mask across all time frames
        #     mag = mag * keep[tf.newaxis, :]

        mag = tf.expand_dims(mag, axis=-1)  # [T, F, 1]
        return self.parse_tfrecord(mag)
    
    def process_audio(self, audio_path):
        """Process a single audio file."""
        try:
            # Convert audio to log spectrogram
            spectrogram = self.wav_to_logspec(str(audio_path))
            
            # Predict
            prediction = self.model.predict(spectrogram, verbose=0)
            probability = float(prediction[0][0])
            is_gunshot = probability > self.threshold
            
            return probability, is_gunshot
            
        except Exception as e:
            print(f"Error processing {audio_path}: {e}")
            import traceback
            traceback.print_exc()
            return None, None
    
    def monitor_loop(self):
        """Main monitoring loop."""
        try:
            while True:
                # Get all WAV files
                wav_files = list(self.input_folder.glob("*.wav")) + \
                           list(self.input_folder.glob("*.WAV"))
                
                for wav_file in wav_files:
                    if str(wav_file) not in self.processed_files:
                        print(f"\nProcessing: {wav_file.name}")
                        
                        # Process the file
                        probability, is_gunshot = self.process_audio(str(wav_file))
                        
                        if probability is not None:
                            # Determine destination
                            if is_gunshot:
                                dest_folder = self.gunshot_folder
                                status = "GUNSHOT"
                            else:
                                dest_folder = self.non_gunshot_folder
                                status = "Non-gunshot"
                            
                            # Move file
                            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                            new_name = f"{timestamp}_{wav_file.stem}{wav_file.suffix}"
                            dest_path = dest_folder / new_name
                            
                            shutil.move(str(wav_file), str(dest_path))
                            
                            # Print result
                            print(f"  Probability: {probability:.4f}")
                            print(f"  Threshold: {self.threshold}")
                            print(f"  Result: {status}")
                            print(f"  Moved to: {dest_folder.name}/")
                            
                            # Mark as processed
                            self.processed_files.add(str(wav_file))
                
                # Wait before next check
                time.sleep(self.poll_interval)
                
        except KeyboardInterrupt:
            print("\nMonitoring stopped by user")
        except Exception as e:
            print(f"Error: {e}")
            import traceback
            traceback.print_exc()


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", "-i", required=True, help="Folder to monitor")
    parser.add_argument("--model", "-m", required=True, help="Model path")
    parser.add_argument("--params", "-p", help="Params JSON path")
    parser.add_argument("--threshold", "-t", type=float, default=0.5, help="Threshold")
    parser.add_argument("--interval", type=float, default=2, help="Poll interval")
    
    args = parser.parse_args()
    
    monitor = SimpleGunshotMonitor(
        input_folder=args.input,
        model_path=args.model,
        params_path=args.params,
        threshold=args.threshold,
        poll_interval=args.interval
    )
    
    monitor.monitor_loop()