LyndonCatan/audio-forensic-app

1import sys2import os3import json4import subprocess5import shutil6import warnings7import numpy as np8from scipy.io import wavfile9 10import tempfile11 12# FORCE SILENCE13warnings.filterwarnings("ignore")14os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'15 16def convert_to_wav_if_needed(input_path, log_func):17    """18    Tries to read the file. If it fails, converts to WAV using FFmpeg.19    Returns (path_to_read, is_temp)20    """21    try:22        # Check if readable23        try:24            wavfile.read(input_path)25            return input_path, False26        except Exception:27            log_func(f"Direct read failed, attempting conversion for {input_path}")28            29            tmp = tempfile.NamedTemporaryFile(delete=False, suffix='.wav')30            tmp.close()31            output_path = tmp.name32            33            cmd = [34                "ffmpeg", "-y", 35                "-i", input_path, 36                "-ar", "44100", 37                 output_path38            ]39            # Suppress output40            subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True)41            log_func(f"Converted to {output_path}")42            return output_path, True43    except Exception as e:44        log_func(f"Conversion failed: {str(e)}")45        return input_path, False46 47def separate_audio(input_path, output_dir, job_id, classification_path=None):48    debug_log = []49    50    def log(msg):51        debug_log.append(str(msg))52 53    converted_audio_path = None54    is_temp_file = False55 56    try:57        log(f"Start separation. Input: {input_path}, Job: {job_id}")58        input_path = os.path.abspath(input_path.strip('"'))59        output_dir = os.path.abspath(output_dir.strip('"'))60        61        # 0. Ensure Input is Valid WAV62        # Demucs might handle MP3, but since we had ID3 issues, let's normalize first.63        read_path, is_temp = convert_to_wav_if_needed(input_path, log)64        converted_audio_path = read_path65        is_temp_file = is_temp66        67        if classification_path:68            classification_path = os.path.abspath(classification_path.strip('"'))69        70        input_filename = os.path.basename(converted_audio_path)71        input_no_ext = os.path.splitext(input_filename)[0]72        # Demucs output folder is based on the input filename. 73        # If we converted to a temp file 'tmp123.wav', Demucs will output to 'htdemucs/tmp123'.74        # We need to map this back or rename.75        76        # actually, to preserve the job_id or original name context, we might want to check77        # but let's see what Demucs does.78 79        # 1. Run Demucs (In-process to bypass torchaudio.save issues)80        print(f"[Demucs] Loading model htdemucs...", file=sys.stderr)81        82        # Imports inside function to avoid heavy load if not needed83        import torch84        from demucs.pretrained import get_model85        from demucs.apply import apply_model86        import torchaudio.transforms as T87        88        # Load Model89        model = get_model("htdemucs")90        model.cpu()91        model.eval()92        93        # Load Audio via Scipy (safe)94        sr, audio_data = wavfile.read(read_path)95        96        # Convert to float32 and normalize to [-1, 1]97        if audio_data.dtype == np.int16:98             audio_data = audio_data.astype(np.float32) / 32768.099        elif audio_data.dtype == np.int32:100             audio_data = audio_data.astype(np.float32) / 2147483648.0101        elif audio_data.dtype == np.uint8:102             audio_data = (audio_data.astype(np.float32) - 128) / 128.0103             104        # Ensure shape (Channels, Samples)105        if len(audio_data.shape) == 1:106            audio_data = np.expand_dims(audio_data, axis=0) # (1, Samples)107        else:108            audio_data = audio_data.T # (Channels, Samples)109            110        # Demucs expects Stereo (2 channels). Mix/Duplicate if Mono.111        if audio_data.shape[0] == 1:112            audio_data = np.concatenate([audio_data, audio_data], axis=0)113            114        # Convert to Tensor115        wav = torch.tensor(audio_data)116        117        # Resample if needed (Demucs htdemucs is 44100Hz)118        if sr != model.samplerate:119            print(f"[Demucs] Resampling {sr} -> {model.samplerate}Hz", file=sys.stderr)120            resampler = T.Resample(sr, model.samplerate)121            wav = resampler(wav)122            123        # Normalization (Standard Demucs procedure)124        ref = wav.mean(0)125        wav = (wav - ref.mean()) / ref.std()126        127        # Separate128        print(f"[Demucs] Separating...", file=sys.stderr)129        # sources shape: (Sources, Channels, Samples)130        sources = apply_model(model, wav[None], device="cpu", shifts=1, split=True, overlap=0.25, progress=True)[0]131        132        # De-normalize133        sources = sources * ref.std() + ref.mean()134        135        print("[Demucs] Separation finished. Saving stems...", file=sys.stderr)136        137        # Save Stems manually using scipy138        stem_names = model.sources # ['drums', 'bass', 'other', 'vocals'] for htdemucs139        140        # Create output structure matching standard Demucs141        # htdemucs/filename_no_ext/142        demucs_folder_name = os.path.splitext(os.path.basename(read_path))[0]143        separated_folder = os.path.join(output_dir, "htdemucs", demucs_folder_name)144        os.makedirs(separated_folder, exist_ok=True)145        146        sources_np = sources.numpy()147        148        for i, name in enumerate(stem_names):149            stem_audio = sources_np[i] # (Channels, Samples)150            stem_audio = stem_audio.T # (Samples, Channels)151            152            out_file = os.path.join(separated_folder, f"{name}.wav")153            wavfile.write(out_file, model.samplerate, stem_audio)154            155        log(f"Demucs output saved to: {separated_folder}")156 157        # Populate final_stems158        final_stems = {}159        if os.path.exists(os.path.join(separated_folder, "vocals.wav")):160            final_stems["vocals"] = f"/separated_audio/htdemucs/{demucs_folder_name}/vocals.wav"161        if os.path.exists(os.path.join(separated_folder, "other.wav")):162            final_stems["background"] = f"/separated_audio/htdemucs/{demucs_folder_name}/other.wav"163 164 165        # 2. Forensic Event Masking (if classification provided)166        if classification_path and os.path.exists(classification_path):167            try:168                log("Starting forensic masking...")169                with open(classification_path, 'r') as f:170                    classification_data = json.load(f)171                172                log(f"Loaded classification data. Keys: {list(classification_data.keys())}")173                if "status" in classification_data and classification_data["status"] == "error":174                     log(f"Classification ERROR: {classification_data.get('message', 'No message')}")175 176                # Load original Audio (already converted/validated)177                sr, audio_data = wavfile.read(read_path)178                log(f"Loaded audio. Sample rate: {sr}, Shape: {audio_data.shape}")179                180                # Prepare empty containers (silence)181                stems_to_generate = {182                    "vocals": "Human Voice",183                    "background": "Musical Content",184                    "vehicles": "Vehicle Sound",185                    "footsteps": "Footsteps",186                    "animals": "Animal Signal",187                    "wind": "Atmospheric Wind",188                    "gunshots": "Gunshot / Explosion",189                    "screams": "Scream / Aggression",190                    "sirens": "Siren / Alarm",191                    "impact": "Impact / Breach"192                }193                194                # Check what we already have from Demucs195                has_demucs_vocals = "vocals" in final_stems196                has_demucs_background = "background" in final_stems197                198                if len(audio_data.shape) > 1:199                     generated_audio = { key: np.zeros_like(audio_data) for key in stems_to_generate }200                else:201                     generated_audio = { key: np.zeros_like(audio_data) for key in stems_to_generate }202 203                # Iterate events and fill segments204                events = classification_data.get("soundEvents", [])205                log(f"Found {len(events)} sound events.")206                207                CLIP_DURATION = 0.975208                209                count_generated = 0210                for event in events:211                    etype = event.get("type", "")212                    target_stem = None213                    for stem_key, trigger_word in stems_to_generate.items():214                        if trigger_word.lower() == etype.lower():215                            target_stem = stem_key216                            break217                    218                    if target_stem:219                        # Skip forensic generation if Demucs already provided it (higher quality)220                        if target_stem == "vocals" and has_demucs_vocals:221                            continue222                        if target_stem == "background" and has_demucs_background:223                            continue224 225                        start_time = float(event.get("time", 0))226                        end_time = start_time + CLIP_DURATION227                        start_idx = int(start_time * sr)228                        end_idx = int(end_time * sr)229                        start_idx = max(0, start_idx)230                        end_idx = min(len(audio_data), end_idx)231                        232                        if start_idx < end_idx:233                            generated_audio[target_stem][start_idx:end_idx] = audio_data[start_idx:end_idx]234                            count_generated += 1235 236                log(f"Processed {count_generated} event segments matches.")237 238                # Save generated stems239                gen_dir = os.path.join(output_dir, "generated", job_id)240                os.makedirs(gen_dir, exist_ok=True)241                242                for stem_key, audio_arr in generated_audio.items():243                    peak = np.max(np.abs(audio_arr))244                    log(f"Stem {stem_key} peak amplitude: {peak}")245                    246                    if peak > 0:247                        out_file = os.path.join(gen_dir, f"{stem_key}.wav")248                        wavfile.write(out_file, sr, audio_arr)249                        final_stems[stem_key] = f"/separated_audio/generated/{job_id}/{stem_key}.wav"250            251            except Exception as e:252                log(f"Masking Exception: {str(e)}")253 254        if not final_stems:255             log("No stems were generated.")256             return {"status": "error", "message": "Separation failed, no stems found.", "debug": debug_log}257 258        return {"status": "success", "stems": final_stems, "debug": debug_log}259    except Exception as e:260        return {"status": "error", "message": str(e), "debug": debug_log}261    finally:262        if is_temp_file and converted_audio_path and os.path.exists(converted_audio_path):263            try:264                os.unlink(converted_audio_path)265            except:266                pass267 268if __name__ == "__main__":269    # Ensure no other prints exist in this file!270    if len(sys.argv) > 3:271        # Check for optional 4th arg272        cls_path = sys.argv[4] if len(sys.argv) > 4 else None273        result = separate_audio(sys.argv[1], sys.argv[2], sys.argv[3], cls_path)274        sys.stdout.write(json.dumps(result))275    else:276        sys.stdout.write(json.dumps({"status": "error", "message": "Insufficient arguments"}))277    sys.stdout.flush()