File size: 5,509 Bytes
4d3248c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
import os
import glob
import json
import logging
import subprocess
from concurrent.futures import ProcessPoolExecutor, as_completed
import pandas as pd
from tqdm import tqdm

logging.basicConfig(level=logging.INFO)

BAK_JSONL = r"D:/dataset/audio/New folder/thaiserextract/metadata.jsonl.bak"
PARQUET_DIR = r"D:/dataset/audio/New folder/thaiser/data"
OUTPUT_AUDIO_DIR = r"D:/dataset/audio/New folder/thaiserextract/audio"
OUTPUT_JSONL = r"D:/dataset/audio/New folder/thaiserextract/metadata_clean.jsonl"

os.makedirs(OUTPUT_AUDIO_DIR, exist_ok=True)

def extract_base_id(audio_id):
    return audio_id.replace("_con_", "_").replace("_clip_", "_").replace("_middle_", "_")

def safe_json_loads(line):
    # Hack to avoid cp1252 weirdness since some lines might be broken in Windows console
    # but we are in python so utf-8 should just work if we ignore errors.
    try:
        return json.loads(line)
    except:
        return None

def load_bak_mapping():
    mapping = {}
    if os.path.exists(BAK_JSONL):
        with open(BAK_JSONL, "r", encoding="utf-8", errors="ignore") as f:
            for line in f:
                data = safe_json_loads(line)
                if data and "audio_id" in data and "text" in data:
                    base_id = extract_base_id(data["audio_id"])
                    mapping[base_id] = {
                        "text": data["text"],
                        "speaker_id": data.get("speaker_id", "Unknown")
                    }
    return mapping

def convert_bytes_to_mp3(audio_bytes, out_mp3):
    if os.path.exists(out_mp3):
        return True
    try:
        cmd = ['ffmpeg', '-y', '-f', 'flac', '-i', 'pipe:0', '-c:a', 'libmp3lame', '-q:a', '2', out_mp3]
        subprocess.run(cmd, input=audio_bytes, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        return True
    except Exception as e:
        return False

def process_parquet(parquet_file, mapping):
    results = []
    try:
        df = pd.read_parquet(parquet_file)
        # Filter matching rows
        df = df[df['assigned_emo'].notna() & (df['assigned_emo'].astype(str).str.lower() != 'none')]
        
        for _, row in df.iterrows():
            audio_id_con = row['audio_id']
            base_id = extract_base_id(audio_id_con)
            
            text = None
            speaker_id = str(row.get('actor_id', 'Unknown'))
            emo = row['assigned_emo']

            if base_id in mapping:
                text = mapping[base_id]["text"]
                speaker_id = mapping[base_id]["speaker_id"]
            elif pd.notna(row.get('script_sent')):
                text = str(row['script_sent'])
            
            if not text or len(text.strip()) == 0:
                continue

            mics = [("mic_con", "_con_"), ("mic_clip", "_clip_"), ("mic_middle", "_middle_")]
            # Construct the base name parts. Usually "s001_actor001_..."
            # We want to insert the mic infix properly.
            # Example base_id: s001_actor001_impro1_1.flac
            # So replace the first '_' with '_con_' -> s001_con_actor001_impro1_1.flac
            for col, infix in mics:
                if col in row and pd.notna(row[col]) and isinstance(row[col], dict) and 'bytes' in row[col]:
                    first_underscore = base_id.find('_')
                    if first_underscore != -1:
                        new_audio_id = base_id[:first_underscore] + infix + base_id[first_underscore+1:]
                    else:
                        new_audio_id = infix.strip('_') + "_" + base_id
                        
                    out_mp3_name = new_audio_id + ".mp3"
                    out_mp3_path = os.path.join(OUTPUT_AUDIO_DIR, out_mp3_name)
                    
                    success = convert_bytes_to_mp3(row[col]['bytes'], out_mp3_path)
                    if success:
                        results.append({
                            "audio_id": new_audio_id,
                            "file_path": out_mp3_path.replace("\\", "/"),
                            "text": text,
                            "speaker_id": speaker_id,
                            "emotion": emo,
                            "path": out_mp3_path.replace("\\", "/")
                        })
    except Exception as e:
        print(f"Error processing {parquet_file}: {e}")
    return results

def main():
    print("Loading mapping from metadata.jsonl.bak...")
    mapping = load_bak_mapping()
    print(f"Loaded {len(mapping)} verified texts.")
    
    parquet_files = glob.glob(os.path.join(PARQUET_DIR, "*.parquet"))
    print(f"Found {len(parquet_files)} parquet files to process.")
    
    all_results = []
    
    # Process sequentially for safety (ffmpeg multi processing can hit cpu limits if not careful, 
    # but let's use a small process pool 4 workers).
    with ProcessPoolExecutor(max_workers=4) as executor:
        futures = {executor.submit(process_parquet, pf, mapping): pf for pf in parquet_files}
        
        for future in tqdm(as_completed(futures), total=len(futures), desc="Extracting Audio"):
            res = future.result()
            all_results.extend(res)
            
    print(f"Extraction complete! Total multi-mic segments: {len(all_results)}")
    
    print("Writing metadata_clean.jsonl...")
    with open(OUTPUT_JSONL, "w", encoding="utf-8") as f:
        for r in all_results:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")
    print("Done!")

if __name__ == "__main__":
    main()