LyndonCatan/audio-forensic-app

1import os2import torch3import soundfile as sf4import json5import io6from pydub import AudioSegment7from pyannote.audio import Pipeline8from huggingface_hub import login9from dotenv import load_dotenv10 11# --- 1. LOCAL FFMEPG SETUP ---12BASE_DIR = os.path.dirname(os.path.abspath(__file__))13FFMPEG_DIR = os.path.join(BASE_DIR, "ffmpeg")14os.environ["PATH"] += os.pathsep + FFMPEG_DIR15 16AudioSegment.converter = os.path.join(FFMPEG_DIR, "ffmpeg.exe")17AudioSegment.ffprobe = os.path.join(FFMPEG_DIR, "ffprobe.exe")18 19# --- 2. AUTHENTICATION (The Secure Way) ---20# This tells Python to look for the .env file in the same folder as this script21load_dotenv(os.path.join(BASE_DIR, ".env"))22HF_TOKEN = os.getenv("HF_TOKEN")23 24def run_forensic_analysis(audio_path):25    print(f"\n[INFO] Initializing Offline Forensic Analysis...")26    27    if not HF_TOKEN:28        print("[ERROR] HF_TOKEN not found! Ensure the .env file is in the scripts folder.")29        return30 31    try:32        # Authenticate with Hugging Face using the token from .env33        login(token=HF_TOKEN)34        35        print("[INFO] Building AI Brain from local cache...")36        pipeline = Pipeline.from_pretrained(37            "pyannote/speaker-diarization-3.1", 38            token=HF_TOKEN39        )40        41        # Use CPU for processing42        pipeline.to(torch.device("cpu"))43 44        # Step B: Manual Audio Decoding45        print(f"[INFO] Decoding: {os.path.basename(audio_path)}")46        audio = AudioSegment.from_file(audio_path).set_frame_rate(16000).set_channels(1)47        48        buffer = io.BytesIO()49        audio.export(buffer, format="wav")50        buffer.seek(0)51        data, samplerate = sf.read(buffer)52        waveform = torch.tensor(data).float().unsqueeze(0)53 54        # Step C: Run Analysis55        print("[INFO] Analyzing voices... (Processing locally)")56        diarization = pipeline({"waveform": waveform, "sample_rate": samplerate})57 58        # --- 3. ORGANIZE DATA FOR JSON ---59        json_output = {60            "fileName": os.path.basename(audio_path),61            "totalDuration": round(len(audio) / 1000.0, 2),62            "segments": []63        }64 65        for turn, _, speaker in diarization.itertracks(yield_label=True):66            json_output["segments"].append({67                "start": round(turn.start, 2),68                "end": round(turn.end, 2),69                "speaker": speaker70            })71 72        # Save results to JSON73        output_file = os.path.join(BASE_DIR, "analysis_results.json")74        with open(output_file, 'w') as f:75            json.dump(json_output, f, indent=4)76 77        print("\n" + "="*45)78        print(f" SUCCESS: Results saved to analysis_results.json")79        print("="*45)80        81    except Exception as e:82        print(f"\n[ERROR] Analysis failed: {e}")83 84if __name__ == "__main__":85    # Assumes test_audio.wav is in the scripts folder86    target_audio = os.path.join(BASE_DIR, "test_audio.wav")87    if os.path.exists(target_audio):88        run_forensic_analysis(target_audio)89    else:90        print(f"[!] File not found: {target_audio}")