Spaces:
Running
Running
| import os | |
| import sys | |
| import numpy as np | |
| import librosa | |
| import soundfile as sf | |
| import pandas as pd | |
| from tqdm import tqdm | |
| import gc | |
| import tensorflow as tf | |
| import time | |
| # Add project root to path | |
| sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) | |
| from utils.hear_extractor import HeARExtractor | |
| # --- Config --- | |
| OUTPUT_DIR = r"c:\Users\ASUS\lung_ai_project\data\hear_embeddings_optimized" | |
| TEMP_AUDIO_DIR = r"c:\Users\ASUS\lung_ai_project\data\temp_aug_audio" | |
| ORIG_HEAR_DIR = r"c:\Users\ASUS\lung_ai_project\data\hear_embeddings" | |
| RESP_BASE = r"c:\Users\ASUS\lung_ai_project\data\extracted_cough\Respiratory_Sound_Dataset-main" | |
| COS_BASE = r"c:\Users\ASUS\lung_ai_project\data\coswara" | |
| CHECKPOINT_INTERVAL = 50 | |
| # --- Augmentations --- | |
| def add_noise(data, noise_factor=0.005): | |
| noise = np.random.randn(len(data)) | |
| augmented_data = data + noise_factor * noise | |
| return augmented_data | |
| def speed_change(data, speed_factor=0.9): | |
| # Resample is much faster than time_stretch/pitch_shift (FFT based) | |
| # This changes both pitch and speed, which is a valid augmentation | |
| new_len = int(len(data) / speed_factor) | |
| return librosa.resample(data, orig_sr=16000, target_sr=int(16000*speed_factor)) | |
| # --- Data Collection --- | |
| def get_sick_files(): | |
| files = [] | |
| # Coswara | |
| csv_dir = os.path.join(COS_BASE, "csvs") | |
| data_dir = os.path.join(COS_BASE, "coswara_data", "kaggle_data") | |
| status_map = {} | |
| if os.path.exists(csv_dir): | |
| for csv_file in os.listdir(csv_dir): | |
| if csv_file.endswith(".csv"): | |
| df = pd.read_csv(os.path.join(csv_dir, csv_file)) | |
| if 'id' in df.columns and 'covid_status' in df.columns: | |
| for _, row in df.iterrows(): | |
| status_map[row['id']] = row['covid_status'] | |
| if os.path.exists(data_dir): | |
| for pid in os.listdir(data_dir): | |
| status = status_map.get(pid) | |
| if status and status.lower() != "healthy": | |
| pid_dir = os.path.join(data_dir, pid) | |
| for af in ["cough.wav", "cough-heavy.wav", "cough-shallow.wav"]: | |
| path = os.path.join(pid_dir, af) | |
| if os.path.exists(path): | |
| files.append(path) | |
| break | |
| # Respiratory | |
| resp_audio = os.path.join(RESP_BASE, "audio_and_txt_files") | |
| resp_csv = os.path.join(RESP_BASE, "patient_diagnosis.csv") | |
| if os.path.exists(resp_csv): | |
| df = pd.read_csv(resp_csv) | |
| diag_map = dict(zip(df['Patient_ID'], df['DIAGNOSIS'])) | |
| for f in os.listdir(resp_audio): | |
| if f.endswith(".wav"): | |
| try: | |
| pid = int(f.split('_')[0]) | |
| diag = diag_map.get(pid) | |
| if diag and diag.lower() != "healthy": | |
| files.append(os.path.join(resp_audio, f)) | |
| except: continue | |
| return files | |
| def main(): | |
| # Setup directories | |
| if not os.path.exists(OUTPUT_DIR): | |
| os.makedirs(OUTPUT_DIR) | |
| if not os.path.exists(TEMP_AUDIO_DIR): | |
| os.makedirs(TEMP_AUDIO_DIR) | |
| print("Identifying Sick Files...") | |
| sick_files = get_sick_files() | |
| print(f"Found {len(sick_files)} sick files.") | |
| # Load feature lists | |
| features = [] | |
| labels = [] | |
| # Check for existing checkpoint | |
| checkpoint_path = os.path.join(OUTPUT_DIR, "checkpoint_indices.npy") | |
| start_idx = 0 | |
| if os.path.exists(checkpoint_path): | |
| start_idx = np.load(checkpoint_path).item() | |
| print(f"Resuming from index {start_idx}") | |
| features = list(np.load(os.path.join(OUTPUT_DIR, "X_checkpoint.npy"))) | |
| labels = list(np.load(os.path.join(OUTPUT_DIR, "y_checkpoint.npy"))) | |
| print("Loading HeAR Extractor...") | |
| extractor = HeARExtractor() | |
| # Processing Loop | |
| processed_count = 0 | |
| for i in tqdm(range(start_idx, len(sick_files))): | |
| file_path = sick_files[i] | |
| try: | |
| # Memory Cleanup | |
| if processed_count % CHECKPOINT_INTERVAL == 0 and processed_count > 0: | |
| gc.collect() | |
| tf.keras.backend.clear_session() | |
| # Save Checkpoint | |
| np.save(os.path.join(OUTPUT_DIR, "X_checkpoint.npy"), np.array(features)) | |
| np.save(os.path.join(OUTPUT_DIR, "y_checkpoint.npy"), np.array(labels)) | |
| np.save(checkpoint_path, i) | |
| # Load Audio (limit duration to 5s to save speed) | |
| y, sr = librosa.load(file_path, sr=16000, duration=5.0) | |
| if len(y) < 2000: # Skip empty/too short | |
| continue | |
| # Aug 1: Noise (Fast) | |
| y_noise = add_noise(y) | |
| temp_path_1 = os.path.join(TEMP_AUDIO_DIR, "temp_noise.wav") | |
| sf.write(temp_path_1, y_noise, 16000) | |
| emb1 = extractor.extract(temp_path_1) | |
| if emb1 is not None: | |
| features.append(emb1) | |
| labels.append("sick") | |
| # Aug 2: Speed/Pitch Change (Resampling - Fast) | |
| y_speed = speed_change(y, speed_factor=0.9) # Slightly slower/deeper | |
| temp_path_2 = os.path.join(TEMP_AUDIO_DIR, "temp_speed.wav") | |
| sf.write(temp_path_2, y_speed, 16000) | |
| emb2 = extractor.extract(temp_path_2) | |
| if emb2 is not None: | |
| features.append(emb2) | |
| labels.append("sick") | |
| except Exception as e: | |
| print(f"Error on {file_path}: {e}") | |
| continue | |
| processed_count += 1 | |
| # Merge | |
| print("Merging with Original Data...") | |
| if os.path.exists(os.path.join(ORIG_HEAR_DIR, "X_hear.npy")): | |
| X_orig = np.load(os.path.join(ORIG_HEAR_DIR, "X_hear.npy")) | |
| y_orig = np.load(os.path.join(ORIG_HEAR_DIR, "y_hear.npy")) | |
| X_final = np.concatenate([X_orig, np.array(features)]) | |
| y_final = np.concatenate([y_orig, np.array(labels)]) | |
| else: | |
| X_final = np.array(features) | |
| y_final = np.array(labels) | |
| np.save(os.path.join(OUTPUT_DIR, "X_hear_opt_merged.npy"), X_final) | |
| np.save(os.path.join(OUTPUT_DIR, "y_hear_opt_merged.npy"), y_final) | |
| print(f"DONE. Saved {len(X_final)} total samples to {OUTPUT_DIR}") | |
| # Cleanup Temp | |
| try: | |
| import shutil | |
| shutil.rmtree(TEMP_AUDIO_DIR) | |
| except: pass | |
| if __name__ == "__main__": | |
| main() | |