Commit 63589d3c 63589d3c84acde4f51e11b0857525c12c41cbca2 by 沈秋雨

更新支持pg下载多版本歌曲

1 parent 5ae39abf
...@@ -13,10 +13,11 @@ ...@@ -13,10 +13,11 @@
13 python scripts/acrcloud/generate_acrcloud_testset.py \ 13 python scripts/acrcloud/generate_acrcloud_testset.py \
14 --audio-dir /Volumes/移动硬盘/composition_test \ 14 --audio-dir /Volumes/移动硬盘/composition_test \
15 --negative-audio-dir /Volumes/移动硬盘/composition_drop \ 15 --negative-audio-dir /Volumes/移动硬盘/composition_drop \
16 --out-dir acrcloud_testset_cloud \ 16 --out-dir /Volumes/移动硬盘/acrcloud_testset_cloud \
17 --num-songs 20 \ 17 --num-songs 600 \
18 --num-negative-songs 80 \ 18 --num-negative-songs 399 \
19 --seed 123 19 --seed 123 \
20 --negative-variants
20 21
21 输出: 22 输出:
22 reference.csv — 参照曲(原始文件),需提前入库 23 reference.csv — 参照曲(原始文件),需提前入库
...@@ -79,6 +80,8 @@ POSITIVE_VARIANTS: list[tuple[str, str | None]] = [ ...@@ -79,6 +80,8 @@ POSITIVE_VARIANTS: list[tuple[str, str | None]] = [
79 ("codec_320k", "acodec=libmp3lame,b:a=320k"), 80 ("codec_320k", "acodec=libmp3lame,b:a=320k"),
80 # 小范围音效:轻微混响 81 # 小范围音效:轻微混响
81 ("reverb_small", "aecho=0.8:0.88:60:0.4"), 82 ("reverb_small", "aecho=0.8:0.88:60:0.4"),
83 # 10 秒短片段:应保留足够指纹特征,应被识别为去重
84 ("short_clip", None),
82 ] 85 ]
83 86
84 # -------------------------------------------------------------------------- 87 # --------------------------------------------------------------------------
...@@ -91,8 +94,6 @@ NEGATIVE_VARIANTS: list[tuple[str, str | None]] = [ ...@@ -91,8 +94,6 @@ NEGATIVE_VARIANTS: list[tuple[str, str | None]] = [
91 # 升调 / 降调(改变音高,破坏指纹频域特征) 94 # 升调 / 降调(改变音高,破坏指纹频域特征)
92 ("pitch_up2", "aresample=22050,asetrate=22050*1.1225,aresample=22050"), 95 ("pitch_up2", "aresample=22050,asetrate=22050*1.1225,aresample=22050"),
93 ("pitch_down2", "aresample=22050,asetrate=22050*0.8909,aresample=22050"), 96 ("pitch_down2", "aresample=22050,asetrate=22050*0.8909,aresample=22050"),
94 # 极端片段(过短的片段不足以提取有效指纹)
95 ("short_clip", None),
96 ] 97 ]
97 98
98 # 片段拼接负样本:每个样本从几首不同歌曲各取一段拼接 99 # 片段拼接负样本:每个样本从几首不同歌曲各取一段拼接
...@@ -295,38 +296,84 @@ def main() -> None: ...@@ -295,38 +296,84 @@ def main() -> None:
295 ref_rows = [] 296 ref_rows = []
296 query_rows = [] 297 query_rows = []
297 298
299 # 断点续跑:加载已有 CSV,跳过已生成的条目
300 ref_path = out_dir / "reference.csv"
301 query_path = out_dir / "queries.csv"
302
303 done_ref_ids: set[str] = set()
304 done_query_paths: set[str] = set()
305
306 if ref_path.exists():
307 try:
308 with ref_path.open(newline="", encoding="utf-8") as f:
309 for row in csv.DictReader(f):
310 done_ref_ids.add(row["song_id"])
311 logger.info("续跑:已加载 reference.csv,跳过 %d 首参照歌", len(done_ref_ids))
312 except (UnicodeDecodeError, Exception) as e:
313 logger.warning("reference.csv 读取失败(%s),将重新生成", e)
314
315 if query_path.exists():
316 try:
317 with query_path.open(newline="", encoding="utf-8") as f:
318 for row in csv.DictReader(f):
319 done_query_paths.add(row["audio_path"])
320 logger.info("续跑:已加载 queries.csv,跳过 %d 条查询", len(done_query_paths))
321 except (UnicodeDecodeError, Exception) as e:
322 logger.warning("queries.csv 读取失败(%s),将重新生成", e)
323
298 # ---- 应去重(expected=duplicate):参照歌 + 轻微变换 ---- 324 # ---- 应去重(expected=duplicate):参照歌 + 轻微变换 ----
299 for wav in _tqdm(selected, desc="生成正样本变体", total=len(selected)): 325 for wav in _tqdm(selected, desc="生成正样本变体", total=len(selected)):
300 song_id = _song_id(wav) 326 song_id = _song_id(wav)
301 327
302 ref_rows.append({ 328 if song_id not in done_ref_ids:
303 "song_id": song_id, 329 ref_rows.append({
304 "audio_path": str(wav.resolve()), 330 "song_id": song_id,
305 "variant": "original", 331 "audio_path": str(wav.resolve()),
306 }) 332 "variant": "original",
333 })
307 334
308 # 参照歌原始(自身查询,验证入库和识别链路) 335 audio_path_str = str(wav.resolve())
309 query_rows.append({ 336 if audio_path_str not in done_query_paths:
310 "song_id": song_id, 337 query_rows.append({
311 "audio_path": str(wav.resolve()), 338 "song_id": song_id,
312 "variant": "acr_original", 339 "audio_path": audio_path_str,
313 "sample_class": "positive", 340 "variant": "acr_original",
314 "expected_song_id": song_id, 341 "sample_class": "positive",
315 "expected": "duplicate", 342 "expected_song_id": song_id,
316 }) 343 "expected": "duplicate",
344 })
317 345
318 for variant_name, af in POSITIVE_VARIANTS: 346 for variant_name, af in POSITIVE_VARIANTS:
319 dst = variants_dir / f"{song_id}_{variant_name}.wav" 347 dst = variants_dir / f"{song_id}_{variant_name}.wav"
320 if variant_name == "slight_trim": 348 dst_str = str(dst.resolve())
321 ok = _ffmpeg_trim(wav, dst, start_ratio=0.05, duration_ratio=0.90) 349 if dst_str in done_query_paths:
322 else:
323 ok = _ffmpeg_variant(wav, dst, af)
324 if not ok:
325 logger.warning("正样本变换失败,跳过: %s %s", wav.name, variant_name)
326 continue 350 continue
351 if not dst.exists():
352 if variant_name == "slight_trim":
353 ok = _ffmpeg_trim(wav, dst, start_ratio=0.05, duration_ratio=0.90)
354 elif variant_name == "short_clip":
355 duration = _probe_duration(wav)
356 if duration is not None:
357 usable_end = max(0.0, duration - 10.0)
358 start_min = min(duration * 0.30, usable_end)
359 start_max = min(duration * 0.50, usable_end)
360 ss = random.uniform(start_min, start_max) if start_max > start_min else start_min
361 ok = _run_ffmpeg([
362 "ffmpeg", "-y", "-i", str(wav),
363 "-ss", f"{ss:.3f}", "-t", "10.0",
364 "-ar", "22050", "-ac", "1",
365 str(dst),
366 ])
367 else:
368 ok = False
369 else:
370 ok = _ffmpeg_variant(wav, dst, af)
371 if not ok:
372 logger.warning("正样本变换失败,跳过: %s %s", wav.name, variant_name)
373 continue
327 query_rows.append({ 374 query_rows.append({
328 "song_id": song_id, 375 "song_id": song_id,
329 "audio_path": str(dst.resolve()), 376 "audio_path": dst_str,
330 "variant": variant_name, 377 "variant": variant_name,
331 "sample_class": "positive", 378 "sample_class": "positive",
332 "expected_song_id": song_id, 379 "expected_song_id": song_id,
...@@ -336,9 +383,12 @@ def main() -> None: ...@@ -336,9 +383,12 @@ def main() -> None:
336 # ---- 不应去重(expected=not_duplicate):composition_drop 原始文件 ---- 383 # ---- 不应去重(expected=not_duplicate):composition_drop 原始文件 ----
337 for wav in _tqdm(negative_selected, desc="生成负样本原始", total=len(negative_selected)): 384 for wav in _tqdm(negative_selected, desc="生成负样本原始", total=len(negative_selected)):
338 song_id = _song_id(wav) 385 song_id = _song_id(wav)
386 audio_path_str = str(wav.resolve())
387 if audio_path_str in done_query_paths:
388 continue
339 query_rows.append({ 389 query_rows.append({
340 "song_id": song_id, 390 "song_id": song_id,
341 "audio_path": str(wav.resolve()), 391 "audio_path": audio_path_str,
342 "variant": "negative_original", 392 "variant": "negative_original",
343 "sample_class": "negative", 393 "sample_class": "negative",
344 "expected_song_id": "", 394 "expected_song_id": "",
...@@ -353,30 +403,17 @@ def main() -> None: ...@@ -353,30 +403,17 @@ def main() -> None:
353 403
354 for variant_name, af in NEGATIVE_VARIANTS: 404 for variant_name, af in NEGATIVE_VARIANTS:
355 dst = variants_dir / f"{song_id}_{variant_name}.wav" 405 dst = variants_dir / f"{song_id}_{variant_name}.wav"
356 if variant_name == "short_clip": 406 dst_str = str(dst.resolve())
357 # 截取 5 秒极短片段,不足以提取有效指纹 407 if dst_str in done_query_paths:
358 duration = _probe_duration(wav)
359 if duration is not None:
360 usable_end = max(0.0, duration - 5.0)
361 start_min = min(duration * 0.30, usable_end)
362 start_max = min(duration * 0.50, usable_end)
363 ss = random.uniform(start_min, start_max) if start_max > start_min else start_min
364 ok = _run_ffmpeg([
365 "ffmpeg", "-y", "-i", str(wav),
366 "-ss", f"{ss:.3f}", "-t", "5.0",
367 "-ar", "22050", "-ac", "1",
368 str(dst),
369 ])
370 else:
371 ok = False
372 else:
373 ok = _ffmpeg_variant(wav, dst, af)
374 if not ok:
375 logger.warning("破坏性变换失败,跳过: %s %s", wav.name, variant_name)
376 continue 408 continue
409 if not dst.exists():
410 ok = _ffmpeg_variant(wav, dst, af)
411 if not ok:
412 logger.warning("破坏性变换失败,跳过: %s %s", wav.name, variant_name)
413 continue
377 query_rows.append({ 414 query_rows.append({
378 "song_id": song_id, 415 "song_id": song_id,
379 "audio_path": str(dst.resolve()), 416 "audio_path": dst_str,
380 "variant": variant_name, 417 "variant": variant_name,
381 "sample_class": "negative", 418 "sample_class": "negative",
382 "expected_song_id": song_id, 419 "expected_song_id": song_id,
...@@ -389,48 +426,58 @@ def main() -> None: ...@@ -389,48 +426,58 @@ def main() -> None:
389 srcs = random.sample(selected, SPLICE_SONGS_PER_SAMPLE) 426 srcs = random.sample(selected, SPLICE_SONGS_PER_SAMPLE)
390 splice_id = f"splice_{i:04d}" 427 splice_id = f"splice_{i:04d}"
391 dst = variants_dir / f"{splice_id}_negative_splice.wav" 428 dst = variants_dir / f"{splice_id}_negative_splice.wav"
392 ok = _ffmpeg_splice(srcs, dst) 429 dst_str = str(dst.resolve())
393 if not ok: 430 if dst_str in done_query_paths:
394 logger.warning("片段拼接生成失败,跳过: splice %d", i)
395 continue 431 continue
432 if not dst.exists():
433 ok = _ffmpeg_splice(srcs, dst)
434 if not ok:
435 logger.warning("片段拼接生成失败,跳过: splice %d", i)
436 continue
396 query_rows.append({ 437 query_rows.append({
397 "song_id": splice_id, 438 "song_id": splice_id,
398 "audio_path": str(dst.resolve()), 439 "audio_path": dst_str,
399 "variant": "negative_splice", 440 "variant": "negative_splice",
400 "sample_class": "negative", 441 "sample_class": "negative",
401 "expected_song_id": "", 442 "expected_song_id": "",
402 "expected": "not_duplicate", 443 "expected": "not_duplicate",
403 }) 444 })
404 445
405 # ---- 输出 CSV ---- 446 # ---- 输出 CSV(追加模式,只写本次新增行)----
406 ref_path = out_dir / "reference.csv"
407 query_path = out_dir / "queries.csv"
408
409 fieldnames = ["song_id", "audio_path", "variant", "sample_class", "expected_song_id", "expected"] 447 fieldnames = ["song_id", "audio_path", "variant", "sample_class", "expected_song_id", "expected"]
410 with ref_path.open("w", newline="", encoding="utf-8") as f: 448
449 ref_is_new = not ref_path.exists() or done_ref_ids == set()
450 with ref_path.open("a" if ref_path.exists() else "w", newline="", encoding="utf-8") as f:
411 writer = csv.DictWriter(f, fieldnames=["song_id", "audio_path", "variant"]) 451 writer = csv.DictWriter(f, fieldnames=["song_id", "audio_path", "variant"])
412 writer.writeheader() 452 if ref_is_new:
453 writer.writeheader()
413 writer.writerows(ref_rows) 454 writer.writerows(ref_rows)
414 455
415 with query_path.open("w", newline="", encoding="utf-8") as f: 456 query_is_new = not query_path.exists() or done_query_paths == set()
457 with query_path.open("a" if query_path.exists() else "w", newline="", encoding="utf-8") as f:
416 writer = csv.DictWriter(f, fieldnames=fieldnames) 458 writer = csv.DictWriter(f, fieldnames=fieldnames)
417 writer.writeheader() 459 if query_is_new:
460 writer.writeheader()
418 writer.writerows(query_rows) 461 writer.writerows(query_rows)
419 462
463 total_ref = len(done_ref_ids) + len(ref_rows)
464 total_query = len(done_query_paths) + len(query_rows)
420 pos = sum(1 for r in query_rows if r["expected"] == "duplicate") 465 pos = sum(1 for r in query_rows if r["expected"] == "duplicate")
421 neg = sum(1 for r in query_rows if r["expected"] == "not_duplicate") 466 neg = sum(1 for r in query_rows if r["expected"] == "not_duplicate")
422 logger.info("参照集: %s (%d 条)", ref_path, len(ref_rows)) 467 logger.info("参照集: %s(累计 %d 条,本次新增 %d 条)", ref_path, total_ref, len(ref_rows))
423 logger.info("查询集: %s (%d 条,正样本 %d,负样本 %d)", query_path, len(query_rows), pos, neg) 468 logger.info("查询集: %s(累计 %d 条,本次新增 %d 条,正样本 +%d,负样本 +%d)",
469 query_path, total_query, len(query_rows), pos, neg)
424 470
425 # 按 sample_class + variant 统计 471 # 按 sample_class + variant 统计(全量)
426 from collections import Counter 472 from collections import Counter
427 by_class = Counter(r["sample_class"] for r in query_rows) 473 with query_path.open(newline="", encoding="utf-8") as f:
474 all_query_rows = list(csv.DictReader(f))
475 by_class = Counter(r["sample_class"] for r in all_query_rows)
428 for cls, cnt in sorted(by_class.items()): 476 for cls, cnt in sorted(by_class.items()):
429 logger.info(" %-20s %d 条", cls, cnt) 477 logger.info(" %-20s %d 条", cls, cnt)
430 by_variant = Counter(r["variant"] for r in query_rows) 478 by_variant = Counter(r["variant"] for r in all_query_rows)
431 for variant, cnt in sorted(by_variant.items()): 479 for variant, cnt in sorted(by_variant.items()):
432 logger.info(" %-25s %d 条", variant, cnt) 480 logger.info(" %-25s %d 条", variant, cnt)
433 481
434
435 if __name__ == "__main__": 482 if __name__ == "__main__":
436 main() 483 main()
......
...@@ -12,7 +12,8 @@ ...@@ -12,7 +12,8 @@
12 用法: 12 用法:
13 conda activate hikoon-data-spider 13 conda activate hikoon-data-spider
14 python scripts/aliyun_dna/evaluate_aliyun_dna.py \ 14 python scripts/aliyun_dna/evaluate_aliyun_dna.py \
15 --queries acrcloud_testset_cloud/queries.csv \ 15 --queries /Volumes/移动硬盘/acrcloud_testset_cloud/queries.csv \
16 --concurrency 8 \
16 --out results/aliyun_dna_eval.csv 17 --out results/aliyun_dna_eval.csv
17 18
18 # 只评测特定 variant 19 # 只评测特定 variant
...@@ -35,8 +36,10 @@ import logging ...@@ -35,8 +36,10 @@ import logging
35 import os 36 import os
36 import sys 37 import sys
37 import tempfile 38 import tempfile
39 import threading
38 import time 40 import time
39 import urllib.request 41 import urllib.request
42 from concurrent.futures import ThreadPoolExecutor, as_completed
40 from pathlib import Path 43 from pathlib import Path
41 44
42 sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent)) 45 sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent))
...@@ -229,6 +232,10 @@ def main() -> None: ...@@ -229,6 +232,10 @@ def main() -> None:
229 help=f"降采样目标采样率(默认 {DEFAULT_SAMPLE_RATE}Hz)") 232 help=f"降采样目标采样率(默认 {DEFAULT_SAMPLE_RATE}Hz)")
230 parser.add_argument("--no-resample", action="store_true", 233 parser.add_argument("--no-resample", action="store_true",
231 help="不降采样,使用原始文件") 234 help="不降采样,使用原始文件")
235 parser.add_argument("--save-interval", type=int, default=20,
236 help="每处理 N 条就追加保存一次结果(默认 10,0 表示只在结束时保存)")
237 parser.add_argument("--concurrency", type=int, default=1,
238 help="并发查询数(默认 1,建议不超过 8 以免触发限流)")
232 args = parser.parse_args() 239 args = parser.parse_args()
233 240
234 # 验证配置 241 # 验证配置
...@@ -256,16 +263,46 @@ def main() -> None: ...@@ -256,16 +263,46 @@ def main() -> None:
256 263
257 logger.info("评测样本过滤: 原始 %d 条,保留 %d 条", original_count, len(rows)) 264 logger.info("评测样本过滤: 原始 %d 条,保留 %d 条", original_count, len(rows))
258 265
259 oss_bucket = _get_oss_client()
260 ice_client = _get_ice_client()
261
262 out_path = Path(args.out) 266 out_path = Path(args.out)
263 out_path.parent.mkdir(parents=True, exist_ok=True) 267 out_path.parent.mkdir(parents=True, exist_ok=True)
264 268
265 tmp_files = [] # 跟踪临时降采样文件 269 fieldnames = ["query_song_id", "audio_song_id", "audio_path", "variant", "sample_class",
270 "expected_song_id", "expected", "top1_song_id", "top1_similarity",
271 "top1_hit", "topk_hit", "expected_rank", "expected_similarity",
272 "expected_duplicate", "predicted_duplicate", "correct",
273 "upload_ms", "poll_ms", "total_ms", "error"]
266 274
267 result_rows = [] 275 # 断点续跑:加载已处理的结果
268 for i, row in enumerate(rows, 1): 276 done_paths: set[str] = set()
277 result_rows: list[dict] = []
278 if out_path.exists():
279 with out_path.open(newline="", encoding="utf-8") as f:
280 for r in csv.DictReader(f):
281 result_rows.append(r)
282 done_paths.add(r["audio_path"])
283 logger.info("断点续跑:已加载 %d 条历史结果,跳过已处理条目", len(result_rows))
284
285 rows = [r for r in rows if r["audio_path"] not in done_paths]
286 logger.info("本次待处理: %d 条", len(rows))
287
288 if not rows:
289 logger.info("所有样本均已处理完毕")
290 else:
291 oss_bucket = _get_oss_client()
292 ice_client = _get_ice_client()
293
294 def _flush_results(rows_to_write: list[dict]) -> None:
295 with out_path.open("w", newline="", encoding="utf-8") as f:
296 writer = csv.DictWriter(f, fieldnames=fieldnames)
297 writer.writeheader()
298 writer.writerows(rows_to_write)
299
300 total_pending = len(rows)
301 lock = threading.Lock()
302 completed_count = 0
303
304 def _process_one(row: dict, idx: int) -> dict:
305 """处理单条查询,tmp 文件用完即删,返回结果 dict(含错误时也返回,不抛出)。"""
269 audio_path = row["audio_path"] 306 audio_path = row["audio_path"]
270 query_song_id = row.get("song_id") or _song_id_from_audio_path(audio_path) 307 query_song_id = row.get("song_id") or _song_id_from_audio_path(audio_path)
271 audio_song_id = _song_id_from_audio_path(audio_path) 308 audio_song_id = _song_id_from_audio_path(audio_path)
...@@ -273,6 +310,7 @@ def main() -> None: ...@@ -273,6 +310,7 @@ def main() -> None:
273 expected_dup = row.get("expected", "").strip().lower() == "duplicate" 310 expected_dup = row.get("expected", "").strip().lower() == "duplicate"
274 311
275 upload_path = audio_path 312 upload_path = audio_path
313 tmp_resampled = None
276 try: 314 try:
277 t0 = time.perf_counter() 315 t0 = time.perf_counter()
278 316
...@@ -281,10 +319,10 @@ def main() -> None: ...@@ -281,10 +319,10 @@ def main() -> None:
281 resampled = _resample_audio(audio_path, target_sr=args.sample_rate) 319 resampled = _resample_audio(audio_path, target_sr=args.sample_rate)
282 if resampled: 320 if resampled:
283 upload_path = resampled 321 upload_path = resampled
284 tmp_files.append(resampled) 322 tmp_resampled = resampled
285 323
286 # 2. 上传到 OSS 324 # 2. 上传到 OSS
287 query_id = f"query_{i}_{query_song_id}" 325 query_id = f"query_{idx}_{query_song_id}"
288 if not args.skip_upload: 326 if not args.skip_upload:
289 oss_url = upload_to_oss(oss_bucket, upload_path, query_id) 327 oss_url = upload_to_oss(oss_bucket, upload_path, query_id)
290 else: 328 else:
...@@ -294,12 +332,8 @@ def main() -> None: ...@@ -294,12 +332,8 @@ def main() -> None:
294 332
295 # 3. 提交 DNA 查询 333 # 3. 提交 DNA 查询
296 t1 = time.perf_counter() 334 t1 = time.perf_counter()
297 if oss_url: 335 job_id = submit_dna_query(ice_client, oss_url or "", query_id)
298 job_id = submit_dna_query(ice_client, oss_url, query_id) 336 logger.info("[%d/%d] 提交查询: job_id=%s", idx, total_pending, job_id)
299 else:
300 job_id = submit_dna_query(ice_client, "", query_id)
301
302 logger.info("[%d/%d] 提交查询: job_id=%s", i, len(rows), job_id)
303 337
304 # 4. 轮询结果 338 # 4. 轮询结果
305 result = poll_job_result(ice_client, job_id) 339 result = poll_job_result(ice_client, job_id)
...@@ -320,9 +354,7 @@ def main() -> None: ...@@ -320,9 +354,7 @@ def main() -> None:
320 pk = match.get("PrimaryKey", "") 354 pk = match.get("PrimaryKey", "")
321 sim = match.get("GlobalSimilarity", 0.0) 355 sim = match.get("GlobalSimilarity", 0.0)
322 topk_song_ids.append((pk, sim)) 356 topk_song_ids.append((pk, sim))
323
324 topk_song_ids.sort(key=lambda x: x[1], reverse=True) 357 topk_song_ids.sort(key=lambda x: x[1], reverse=True)
325
326 if topk_song_ids: 358 if topk_song_ids:
327 top1_song_id = topk_song_ids[0][0] 359 top1_song_id = topk_song_ids[0][0]
328 top1_sim = round(topk_song_ids[0][1], 4) 360 top1_sim = round(topk_song_ids[0][1], 4)
...@@ -342,7 +374,15 @@ def main() -> None: ...@@ -342,7 +374,15 @@ def main() -> None:
342 374
343 correct = expected_dup == predicted_dup 375 correct = expected_dup == predicted_dup
344 376
345 result_rows.append({ 377 logger.info(
378 "[%d/%d] variant=%s expected=%s predicted_dup=%s top1=%s sim=%s"
379 " top1_hit=%s topk_hit=%s correct=%s time=%dms",
380 idx, total_pending, row.get("variant", ""), row.get("expected", ""),
381 predicted_dup, top1_song_id or "-", top1_sim if top1_sim != "" else "-",
382 top1_hit, topk_hit, correct, total_ms,
383 )
384
385 return {
346 "query_song_id": query_song_id, 386 "query_song_id": query_song_id,
347 "audio_song_id": audio_song_id, 387 "audio_song_id": audio_song_id,
348 "audio_path": audio_path, 388 "audio_path": audio_path,
...@@ -363,19 +403,12 @@ def main() -> None: ...@@ -363,19 +403,12 @@ def main() -> None:
363 "poll_ms": poll_ms, 403 "poll_ms": poll_ms,
364 "total_ms": total_ms, 404 "total_ms": total_ms,
365 "error": "", 405 "error": "",
366 }) 406 }
367
368 logger.info(
369 "[%d/%d] variant=%s expected=%s predicted_dup=%s top1=%s sim=%s top1_hit=%s topk_hit=%s correct=%s time=%dms",
370 i, len(rows), row.get("variant", ""), row.get("expected", ""),
371 predicted_dup, top1_song_id or "-", top1_sim if top1_sim != "" else "-",
372 top1_hit, topk_hit, correct, total_ms,
373 )
374 407
375 except Exception as e: 408 except Exception as e:
376 total_ms = round((time.perf_counter() - t0) * 1000, 1) 409 total_ms = round((time.perf_counter() - t0) * 1000, 1)
377 logger.error("[%d/%d] 查询失败: %s, %s", i, len(rows), audio_path, e) 410 logger.error("[%d/%d] 查询失败: %s, %s", idx, total_pending, audio_path, e)
378 result_rows.append({ 411 return {
379 "query_song_id": query_song_id, 412 "query_song_id": query_song_id,
380 "audio_song_id": audio_song_id, 413 "audio_song_id": audio_song_id,
381 "audio_path": audio_path, 414 "audio_path": audio_path,
...@@ -396,33 +429,22 @@ def main() -> None: ...@@ -396,33 +429,22 @@ def main() -> None:
396 "poll_ms": "", 429 "poll_ms": "",
397 "total_ms": total_ms, 430 "total_ms": total_ms,
398 "error": str(e), 431 "error": str(e),
399 }) 432 }
400 433 finally:
401 # 清理临时文件 434 if tmp_resampled:
402 for tmp in tmp_files: 435 Path(tmp_resampled).unlink(missing_ok=True)
403 try:
404 Path(tmp).unlink(missing_ok=True)
405 except Exception:
406 pass
407
408 # 写逐条结果
409 fieldnames = ["query_song_id", "audio_song_id", "audio_path", "variant", "sample_class",
410 "expected_song_id", "expected", "top1_song_id", "top1_similarity",
411 "top1_hit", "topk_hit", "expected_rank", "expected_similarity",
412 "expected_duplicate", "predicted_duplicate", "correct",
413 "upload_ms", "poll_ms", "total_ms", "error"]
414 436
415 with out_path.open("w", newline="", encoding="utf-8") as f: 437 # 汇总辅助函数(需在 _collect/_flush_summary 调用前定义)
416 writer = csv.DictWriter(f, fieldnames=fieldnames) 438 def _to_bool(v) -> bool:
417 writer.writeheader() 439 if isinstance(v, bool):
418 writer.writerows(result_rows) 440 return v
441 return str(v).strip().lower() in ("true", "1", "yes")
419 442
420 # 汇总指标
421 def _metrics(rows: list[dict]) -> dict: 443 def _metrics(rows: list[dict]) -> dict:
422 tp = sum(1 for r in rows if r["expected_duplicate"] and r["predicted_duplicate"]) 444 tp = sum(1 for r in rows if _to_bool(r["expected_duplicate"]) and _to_bool(r["predicted_duplicate"]))
423 fp = sum(1 for r in rows if not r["expected_duplicate"] and r["predicted_duplicate"]) 445 fp = sum(1 for r in rows if not _to_bool(r["expected_duplicate"]) and _to_bool(r["predicted_duplicate"]))
424 tn = sum(1 for r in rows if not r["expected_duplicate"] and not r["predicted_duplicate"]) 446 tn = sum(1 for r in rows if not _to_bool(r["expected_duplicate"]) and not _to_bool(r["predicted_duplicate"]))
425 fn = sum(1 for r in rows if r["expected_duplicate"] and not r["predicted_duplicate"]) 447 fn = sum(1 for r in rows if _to_bool(r["expected_duplicate"]) and not _to_bool(r["predicted_duplicate"]))
426 precision = tp / (tp + fp) if tp + fp else 0.0 448 precision = tp / (tp + fp) if tp + fp else 0.0
427 recall = tp / (tp + fn) if tp + fn else 0.0 449 recall = tp / (tp + fn) if tp + fn else 0.0
428 f1 = 2 * precision * recall / (precision + recall) if precision + recall else 0.0 450 f1 = 2 * precision * recall / (precision + recall) if precision + recall else 0.0
...@@ -436,21 +458,6 @@ def main() -> None: ...@@ -436,21 +458,6 @@ def main() -> None:
436 "tp": tp, "fp": fp, "tn": tn, "fn": fn, 458 "tp": tp, "fp": fp, "tn": tn, "fn": fn,
437 } 459 }
438 460
439 metrics = _metrics(result_rows)
440
441 from collections import defaultdict
442 by_variant: dict[str, dict] = defaultdict(lambda: {"correct": 0, "total": 0})
443 for r in result_rows:
444 v = r["variant"] or "unknown"
445 by_variant[v]["total"] += 1
446 if r["correct"]:
447 by_variant[v]["correct"] += 1
448
449 # 耗时统计
450 total_times = [r["total_ms"] for r in result_rows if r.get("total_ms", "") != ""]
451 upload_times = [r["upload_ms"] for r in result_rows if r.get("upload_ms", "") != ""]
452 poll_times = [r["poll_ms"] for r in result_rows if r.get("poll_ms", "") != ""]
453
454 def _time_stats(values: list) -> dict: 461 def _time_stats(values: list) -> dict:
455 s = sorted(values) if values else [] 462 s = sorted(values) if values else []
456 n = len(s) 463 n = len(s)
...@@ -463,33 +470,70 @@ def main() -> None: ...@@ -463,33 +470,70 @@ def main() -> None:
463 "max": round(max(values), 1) if n else 0, 470 "max": round(max(values), 1) if n else 0,
464 } 471 }
465 472
466 summary = { 473 from collections import defaultdict
467 "total": len(result_rows), 474
468 "filters": { 475 def _flush_summary(rows: list[dict]) -> None:
469 "variants": sorted(variant_filter) if variant_filter else None, 476 metrics = _metrics(rows)
470 "sample_classes": sorted(sample_class_filter) if sample_class_filter else None, 477 by_variant: dict[str, dict] = defaultdict(lambda: {"correct": 0, "total": 0})
471 "expected": args.expected, 478 for r in rows:
472 "original_total": original_count, 479 v = r["variant"] or "unknown"
473 }, 480 by_variant[v]["total"] += 1
474 "duplicate_threshold": args.duplicate_threshold, 481 if _to_bool(r["correct"]):
475 "accuracy": metrics["accuracy"], 482 by_variant[v]["correct"] += 1
476 "precision": metrics["precision"], 483 total_times = [float(r["total_ms"]) for r in rows if r.get("total_ms", "") != ""]
477 "recall": metrics["recall"], 484 upload_times = [float(r["upload_ms"]) for r in rows if r.get("upload_ms", "") != ""]
478 "f1": metrics["f1"], 485 poll_times = [float(r["poll_ms"]) for r in rows if r.get("poll_ms", "") != ""]
479 "tp": metrics["tp"], "fp": metrics["fp"], "tn": metrics["tn"], "fn": metrics["fn"], 486 summary = {
480 "query_time_ms": _time_stats(total_times), 487 "total": len(rows),
481 "upload_time_ms": _time_stats(upload_times), 488 "pending": total_pending - len(rows),
482 "poll_time_ms": _time_stats(poll_times), 489 "filters": {
483 "by_variant": { 490 "variants": sorted(variant_filter) if variant_filter else None,
484 v: {"accuracy": round(d["correct"] / d["total"], 4), "total": d["total"]} 491 "sample_classes": sorted(sample_class_filter) if sample_class_filter else None,
485 for v, d in sorted(by_variant.items()) 492 "expected": args.expected,
486 }, 493 "original_total": original_count,
487 "out": str(out_path), 494 },
488 } 495 "duplicate_threshold": args.duplicate_threshold,
496 "accuracy": metrics["accuracy"],
497 "precision": metrics["precision"],
498 "recall": metrics["recall"],
499 "f1": metrics["f1"],
500 "tp": metrics["tp"], "fp": metrics["fp"], "tn": metrics["tn"], "fn": metrics["fn"],
501 "query_time_ms": _time_stats(total_times),
502 "upload_time_ms": _time_stats(upload_times),
503 "poll_time_ms": _time_stats(poll_times),
504 "by_variant": {
505 v: {"accuracy": round(d["correct"] / d["total"], 4), "total": d["total"]}
506 for v, d in sorted(by_variant.items())
507 },
508 "out": str(out_path),
509 }
510 summary_path = out_path.with_suffix(".summary.json")
511 summary_path.write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
512
513 def _collect(result_dict: dict) -> None:
514 nonlocal completed_count
515 with lock:
516 result_rows.append(result_dict)
517 completed_count += 1
518 if args.save_interval > 0 and completed_count % args.save_interval == 0:
519 _flush_results(result_rows)
520 _flush_summary(result_rows)
521 logger.info(" [checkpoint] 已保存 %d 条结果到 %s", len(result_rows), out_path)
522
523 with ThreadPoolExecutor(max_workers=args.concurrency) as executor:
524 futures = {
525 executor.submit(_process_one, row, idx): idx
526 for idx, row in enumerate(rows, 1)
527 }
528 for future in as_completed(futures):
529 _collect(future.result())
530
531 # 写逐条结果(最终落盘)
532 _flush_results(result_rows)
489 533
490 summary_path = out_path.with_suffix(".summary.json") 534 summary_path = out_path.with_suffix(".summary.json")
491 summary_path.write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8") 535 _flush_summary(result_rows)
492 print(json.dumps(summary, ensure_ascii=False, indent=2)) 536 print(json.dumps(json.loads(summary_path.read_text(encoding="utf-8")), ensure_ascii=False, indent=2))
493 537
494 538
495 if __name__ == "__main__": 539 if __name__ == "__main__":
......
1 #!/usr/bin/env python3
2 import argparse
3 import csv
4 import os
5 import re
6 import sys
7 from pathlib import Path
8 from urllib.parse import urlparse
9
10 from dotenv import load_dotenv
11 from qcloud_cos import CosConfig, CosS3Client
12
13 load_dotenv(Path(__file__).resolve().parents[1] / '.env')
14
15 AUDIO_EXTS = {'.mp3', '.wav', '.flac', '.m4a', '.aac', '.ogg', '.wma', '.ape', '.alac'}
16
17
18 def normalize_key(raw_url: str) -> str:
19 if not raw_url:
20 return ''
21 raw_url = raw_url.strip()
22 if re.match(r'^https?://', raw_url, re.I):
23 return urlparse(raw_url).path.lstrip('/')
24 return raw_url.lstrip('/')
25
26
27 def ext_ok(path: str) -> bool:
28 return Path(path).suffix.lower() in AUDIO_EXTS
29
30
31 def main():
32 ap = argparse.ArgumentParser()
33 ap.add_argument('--manifest', default='output_selection_budgeted/selected_files.csv')
34 ap.add_argument('--output-dir', default='downloads')
35 ap.add_argument('--types', default='1,7,8,11,16')
36 ap.add_argument('--song-limit', type=int, default=3)
37 ap.add_argument('--start-song-offset', type=int, default=0)
38 ap.add_argument('--overwrite', action='store_true')
39 ap.add_argument('--fail-log', default='download_failures.csv')
40 args = ap.parse_args()
41
42 region = os.environ['COS_REGION']
43 secret_id = os.environ['COS_SECRET_ID']
44 secret_key = os.environ['COS_SECRET_KEY']
45 bucket = os.environ['COS_BUCKET']
46
47 config = CosConfig(Region=region, SecretId=secret_id, SecretKey=secret_key)
48 client = CosS3Client(config)
49 wanted_types = {int(x) for x in args.types.split(',') if x.strip()}
50
51 out_dir = Path(args.output_dir)
52 out_dir.mkdir(parents=True, exist_ok=True)
53 fail_log = Path(args.fail_log)
54 fail_exists = fail_log.exists()
55 fail_f = fail_log.open('a', encoding='utf-8', newline='')
56 fail_w = csv.writer(fail_f)
57 if not fail_exists:
58 fail_w.writerow(['song_id', 'type', 'url', 'error'])
59
60 downloaded = skipped = errors = 0
61 active_song_ids = []
62 current_song = None
63 started = False
64 song_counter = 0
65
66 with open(args.manifest, 'r', encoding='utf-8', newline='') as f:
67 reader = csv.DictReader(f)
68 for row in reader:
69 sid = row['song_id']
70 if sid != current_song:
71 current_song = sid
72 if song_counter < args.start_song_offset:
73 song_counter += 1
74 continue
75 if args.song_limit and started and len(active_song_ids) >= args.song_limit and sid not in set(active_song_ids):
76 break
77 if sid not in active_song_ids:
78 active_song_ids.append(sid)
79 started = True
80 if sid not in active_song_ids:
81 continue
82 try:
83 t = int(row['type'])
84 except Exception:
85 skipped += 1
86 continue
87 if t not in wanted_types:
88 skipped += 1
89 continue
90 key = normalize_key(row['url'])
91 if not key or not ext_ok(key):
92 skipped += 1
93 continue
94 target_dir = out_dir / sid / f'type_{t}'
95 target_dir.mkdir(parents=True, exist_ok=True)
96 target = target_dir / Path(key).name
97 if target.exists() and not args.overwrite:
98 skipped += 1
99 continue
100 try:
101 client.download_file(Bucket=bucket, Key=key, DestFilePath=str(target))
102 downloaded += 1
103 if downloaded % 50 == 0:
104 print({'downloaded': downloaded, 'songs': len(active_song_ids), 'last': str(target)})
105 except Exception as e:
106 errors += 1
107 fail_w.writerow([sid, t, row['url'], str(e)])
108 fail_f.flush()
109 print({'error_song': sid, 'type': t, 'key': key, 'error': str(e)}, file=sys.stderr)
110
111 fail_f.close()
112 print({
113 'downloaded_files': downloaded,
114 'skipped_rows': skipped,
115 'error_files': errors,
116 'song_count': len(active_song_ids),
117 'output_dir': str(out_dir.resolve()),
118 })
119
120
121 if __name__ == '__main__':
122 main()
1 #!/usr/bin/env python3
2 """从 embed_db 数据库下载歌曲音频,保留完整元数据。
3
4 用法:
5 # 只下载主版本,限 100 首歌
6 python scripts/download_from_db.py --song-limit 100
7
8 # 主版本 + 每首歌最多 3 个其他版本
9 python scripts/download_from_db.py --song-limit 100 --extra-versions 3
10
11 # 主版本 + 所有其他版本
12 python scripts/download_from_db.py --song-limit 100 --extra-versions -1
13
14 # 指定歌曲 ID
15 python scripts/download_from_db.py --song-ids 1,2,3 --extra-versions -1
16
17 # 本地:导出查询结果到 CSV(不下载)
18 python scripts/download_from_db.py --song-limit 500 --extra-versions 3 --export-csv records.csv
19
20 # 服务器:从导出的 CSV 下载(不需要连接数据库)
21 python scripts/download_from_db.py --from-csv records.csv --concurrency 8
22
23 输出目录结构:
24 {output_dir}/
25 metadata.csv — 所有下载记录的完整元数据
26 reference.csv — DNA 入库参照集(is_main_version=1 的主版本)
27 queries.csv — DNA 评测查询集(其余版本,expected=duplicate)
28 download_failures.csv — 下载失败记录
29 {song_id}_{safe_song_name}/
30 {record_id}_{safe_singer_name}{ext}
31 """
32
33 import argparse
34 import csv
35 import logging
36 import re
37 import time
38 import threading
39 import urllib.parse
40 import urllib.request
41 from concurrent.futures import ThreadPoolExecutor, as_completed
42 from pathlib import Path
43
44 from dotenv import load_dotenv
45
46 try:
47 from tqdm import tqdm
48 except ImportError:
49 tqdm = None
50
51 load_dotenv(Path(__file__).resolve().parent.parent / ".env")
52
53 logger = logging.getLogger(__name__)
54
55 DB_DSN = "postgresql://postgres:postgres@localhost:5432/embed_db"
56
57 # fetch_records 查询返回的列,也是 --export-csv / --from-csv 的 CSV 格式
58 RECORDS_FIELDS = [
59 "record_id", "song_id", "song_name", "original_singer",
60 "lyricist", "composer", "genre", "language_tag",
61 "singer_name", "version_name", "platform_name",
62 "duration", "is_main_version", "album", "pub_time", "audio_url",
63 ]
64
65 METADATA_FIELDS = [
66 "record_id", "song_id", "song_name", "original_singer",
67 "lyricist", "composer", "genre", "language_tag",
68 "singer_name", "version_name", "platform_name",
69 "duration", "is_main_version", "album", "pub_time",
70 "audio_url", "local_path", "status",
71 ]
72
73 REFERENCE_FIELDS = [
74 "song_id", "audio_path", "variant",
75 "song_name", "original_singer", "lyricist", "composer",
76 "genre", "language_tag",
77 ]
78
79 QUERY_FIELDS = [
80 "song_id", "audio_path", "variant", "sample_class",
81 "expected_song_id", "expected",
82 "song_name", "original_singer", "singer_name",
83 "platform_name", "duration", "is_main_version",
84 ]
85
86
87 def _safe_name(s: str, max_len: int = 40) -> str:
88 if not s:
89 return "unknown"
90 s = re.sub(r'[\\/:*?"<>|]', "_", s).strip()
91 return s[:max_len] or "unknown"
92
93
94 def _ext_from_url(url: str) -> str:
95 path = url.split("?")[0]
96 suffix = Path(path).suffix.lower()
97 return suffix if suffix in {".mp3", ".wav", ".flac", ".m4a", ".aac", ".ogg"} else ".mp3"
98
99
100 def _encode_url(url: str) -> str:
101 parsed = urllib.parse.urlsplit(url)
102 encoded_path = urllib.parse.quote(parsed.path, safe="/:@!$&'()*+,;=")
103 return urllib.parse.urlunsplit(parsed._replace(path=encoded_path))
104
105
106 def fetch_records(
107 conn: psycopg.Connection,
108 song_ids: list[int] | None,
109 extra_versions: int,
110 song_limit: int,
111 song_offset: int,
112 ) -> list[dict]:
113 """查询待下载记录,JOIN embed_song 获取完整元数据。
114
115 extra_versions:
116 0 — 只下载主版本(is_main_version=1)
117 -1 — 主版本 + 所有其他版本
118 N — 主版本 + 每首歌最多 N 个其他版本(按 id 排序取前 N)
119 """
120 base_where = "r.audio_url IS NOT NULL AND r.audio_url != '' AND r.status = 'ready'"
121
122 if song_ids:
123 ids_str = ",".join(str(i) for i in song_ids)
124 base_where += f" AND r.song_id IN ({ids_str})"
125
126 # 确定要下载的 song_id 范围
127 if song_limit > 0 and not song_ids:
128 song_range_sql = f"""
129 SELECT DISTINCT song_id FROM embed_record
130 WHERE {base_where}
131 ORDER BY song_id
132 LIMIT {song_limit} OFFSET {song_offset}
133 """
134 song_range_clause = f"r.song_id IN ({song_range_sql})"
135 else:
136 song_range_clause = "TRUE"
137
138 full_where = f"{base_where} AND {song_range_clause}"
139
140 # 主版本
141 main_sql = f"""
142 SELECT
143 r.id AS record_id, r.song_id,
144 s.song_name, s.original_singer, s.lyricist, s.composer,
145 s.genre, s.language_tag,
146 r.singer_name, r.version_name, r.platform_name,
147 r.duration, r.is_main_version, r.album, r.pub_time, r.audio_url
148 FROM embed_record r
149 LEFT JOIN embed_song s ON s.id = r.song_id
150 WHERE {full_where} AND r.is_main_version = 1
151 """
152
153 if extra_versions == 0:
154 union_sql = main_sql
155 elif extra_versions == -1:
156 extra_sql = f"""
157 SELECT
158 r.id AS record_id, r.song_id,
159 s.song_name, s.original_singer, s.lyricist, s.composer,
160 s.genre, s.language_tag,
161 r.singer_name, r.version_name, r.platform_name,
162 r.duration, r.is_main_version, r.album, r.pub_time, r.audio_url
163 FROM embed_record r
164 LEFT JOIN embed_song s ON s.id = r.song_id
165 WHERE {full_where} AND r.is_main_version = 0
166 """
167 union_sql = f"{main_sql} UNION ALL {extra_sql}"
168 else:
169 # 每首歌最多取 extra_versions 个非主版本,用窗口函数在 SQL 层截断
170 extra_sql = f"""
171 SELECT record_id, song_id, song_name, original_singer, lyricist, composer,
172 genre, language_tag, singer_name, version_name, platform_name,
173 duration, is_main_version, album, pub_time, audio_url
174 FROM (
175 SELECT
176 r.id AS record_id, r.song_id,
177 s.song_name, s.original_singer, s.lyricist, s.composer,
178 s.genre, s.language_tag,
179 r.singer_name, r.version_name, r.platform_name,
180 r.duration, r.is_main_version, r.album, r.pub_time, r.audio_url,
181 ROW_NUMBER() OVER (PARTITION BY r.song_id ORDER BY r.id) AS rn
182 FROM embed_record r
183 LEFT JOIN embed_song s ON s.id = r.song_id
184 WHERE {full_where} AND r.is_main_version = 0
185 ) ranked
186 WHERE rn <= {extra_versions}
187 """
188 union_sql = f"{main_sql} UNION ALL {extra_sql}"
189
190 final_sql = f"""
191 SELECT * FROM ({union_sql}) combined
192 ORDER BY song_id, is_main_version DESC, record_id
193 """
194
195 with conn.cursor() as cur:
196 cur.execute(final_sql)
197 cols = [desc[0] for desc in cur.description]
198 return [dict(zip(cols, row)) for row in cur.fetchall()]
199
200
201 def load_records_from_csv(path: str) -> list[dict]:
202 """从 --export-csv 导出的文件读取记录,替代数据库查询。"""
203 with open(path, newline="", encoding="utf-8") as f:
204 rows = list(csv.DictReader(f))
205 # 将字符串还原为适当类型
206 for r in rows:
207 r["record_id"] = int(r["record_id"])
208 r["song_id"] = int(r["song_id"])
209 r["is_main_version"] = int(r["is_main_version"]) if r.get("is_main_version") else 0
210 r["duration"] = int(r["duration"]) if r.get("duration") else None
211 return rows
212
213
214 def download_audio(url: str, dest: Path, timeout: int = 60, retries: int = 2) -> bool:
215 url = _encode_url(url)
216 for attempt in range(retries + 1):
217 try:
218 req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
219 with urllib.request.urlopen(req, timeout=timeout) as resp:
220 data = resp.read()
221 dest.write_bytes(data)
222 return True
223 except Exception as e:
224 if attempt < retries:
225 time.sleep(2 ** attempt)
226 else:
227 logger.warning("下载失败 [%d/%d]: %s — %s", attempt + 1, retries + 1, url, e)
228 return False
229
230
231 def _build_paths(record: dict, output_dir: Path) -> tuple[Path, Path]:
232 song_dir = output_dir / f"{record['song_id']}_{_safe_name(record['song_name'] or '')}"
233 ext = _ext_from_url(record["audio_url"])
234 singer = _safe_name(record["singer_name"] or record["original_singer"] or "unknown")
235 filename = f"{record['record_id']}_{singer}{ext}"
236 return song_dir, song_dir / filename
237
238
239 def main() -> None:
240 logging.basicConfig(
241 level=logging.INFO,
242 format="%(asctime)s [%(levelname)s] %(message)s",
243 )
244
245 ap = argparse.ArgumentParser(description="从 embed_db 下载歌曲音频并生成测试集 CSV")
246 ap.add_argument("--output-dir", default="downloads_db", help="输出根目录(默认 downloads_db)")
247 ap.add_argument("--song-limit", type=int, default=0, help="限制歌曲数量(0=不限)")
248 ap.add_argument("--song-offset", type=int, default=0, help="按 song_id 跳过前 N 首歌")
249 ap.add_argument("--song-ids", help="只下载指定 song_id,逗号分隔")
250 ap.add_argument(
251 "--extra-versions", type=int, default=0, metavar="N",
252 help="每首歌额外下载 N 个非主版本(0=只主版本,-1=全部,默认 0)",
253 )
254 ap.add_argument("--concurrency", type=int, default=4, help="并发下载线程数(默认 4)")
255 ap.add_argument("--timeout", type=int, default=60, help="单文件下载超时秒数(默认 60)")
256 ap.add_argument("--overwrite", action="store_true", help="覆盖已存在的文件")
257 ap.add_argument("--dsn", default=DB_DSN, help="PostgreSQL 连接串")
258 ap.add_argument(
259 "--export-csv", metavar="FILE",
260 help="只导出查询结果到 CSV 后退出,不执行下载(在本地机器上运行)",
261 )
262 ap.add_argument(
263 "--from-csv", metavar="FILE",
264 help="从已导出的 CSV 读取记录,跳过数据库连接(在服务器上运行)",
265 )
266 args = ap.parse_args()
267
268 song_ids = [int(x) for x in args.song_ids.split(",")] if args.song_ids else None
269
270 # ── 模式一:从数据库查询 ──────────────────────────────────────────────────
271 if args.from_csv:
272 logger.info("从 CSV 读取记录: %s", args.from_csv)
273 records = load_records_from_csv(args.from_csv)
274 else:
275 logger.info("连接数据库: %s", args.dsn)
276 try:
277 import psycopg
278 except ImportError:
279 logger.error("缺少 psycopg,请执行: pip install psycopg")
280 raise SystemExit(1)
281 with psycopg.connect(args.dsn) as conn:
282 records = fetch_records(
283 conn,
284 song_ids=song_ids,
285 extra_versions=args.extra_versions,
286 song_limit=args.song_limit,
287 song_offset=args.song_offset,
288 )
289
290 song_count_total = len({r["song_id"] for r in records})
291 logger.info("共 %d 条记录,涉及 %d 首歌", len(records), song_count_total)
292
293 # ── 模式二:只导出 CSV,不下载 ────────────────────────────────────────────
294 if args.export_csv:
295 out_csv = Path(args.export_csv)
296 out_csv.parent.mkdir(parents=True, exist_ok=True)
297 with out_csv.open("w", newline="", encoding="utf-8") as f:
298 writer = csv.DictWriter(f, fieldnames=RECORDS_FIELDS, extrasaction="ignore")
299 writer.writeheader()
300 writer.writerows(records)
301 logger.info("已导出 %d 条记录到 %s,可同步到服务器后用 --from-csv 下载", len(records), out_csv)
302 return
303
304 # ── 模式三:下载 ──────────────────────────────────────────────────────────
305 output_dir = Path(args.output_dir)
306 output_dir.mkdir(parents=True, exist_ok=True)
307
308 # 断点续跑:加载已有 metadata.csv
309 metadata_path = output_dir / "metadata.csv"
310 done_record_ids: set[int] = set()
311 existing_metadata: list[dict] = []
312 if metadata_path.exists():
313 with metadata_path.open(newline="", encoding="utf-8") as f:
314 for row in csv.DictReader(f):
315 if row.get("status") == "ok":
316 done_record_ids.add(int(row["record_id"]))
317 existing_metadata.append(row)
318 logger.info("断点续跑:跳过已下载 %d 条", len(done_record_ids))
319
320 pending = [r for r in records if r["record_id"] not in done_record_ids]
321 logger.info("本次待下载: %d 条", len(pending))
322
323 for rec in pending:
324 song_dir, dest = _build_paths(rec, output_dir)
325 rec["_song_dir"] = song_dir
326 rec["_dest"] = dest
327
328 results: list[dict] = list(existing_metadata)
329 results_lock = threading.Lock()
330
331 def _download_one(rec: dict) -> dict:
332 song_dir: Path = rec["_song_dir"]
333 dest: Path = rec["_dest"]
334 song_dir.mkdir(parents=True, exist_ok=True)
335
336 if dest.exists() and not args.overwrite:
337 status = "ok"
338 else:
339 ok = download_audio(rec["audio_url"], dest, timeout=args.timeout)
340 status = "ok" if ok else "failed"
341
342 return {
343 "record_id": rec["record_id"],
344 "song_id": rec["song_id"],
345 "song_name": rec["song_name"] or "",
346 "original_singer": rec["original_singer"] or "",
347 "lyricist": rec["lyricist"] or "",
348 "composer": rec["composer"] or "",
349 "genre": rec["genre"] or "",
350 "language_tag": rec["language_tag"] or "",
351 "singer_name": rec["singer_name"] or "",
352 "version_name": rec["version_name"] or "",
353 "platform_name": rec["platform_name"] or "",
354 "duration": rec["duration"] or "",
355 "is_main_version": rec["is_main_version"],
356 "album": rec["album"] or "",
357 "pub_time": rec["pub_time"] or "",
358 "audio_url": rec["audio_url"],
359 "local_path": str(dest) if status == "ok" else "",
360 "status": status,
361 }
362
363 progress = (
364 tqdm(total=len(pending), unit="文件", dynamic_ncols=True)
365 if tqdm is not None
366 else None
367 )
368
369 with ThreadPoolExecutor(max_workers=args.concurrency) as executor:
370 futures = {executor.submit(_download_one, rec): rec for rec in pending}
371 for future in as_completed(futures):
372 result = future.result()
373 with results_lock:
374 results.append(result)
375 if progress is not None:
376 status_str = "✓" if result["status"] == "ok" else "✗"
377 progress.set_postfix_str(
378 f"{status_str} {result['song_name']} — {result['singer_name'] or result['original_singer']}",
379 refresh=False,
380 )
381 progress.update(1)
382 else:
383 done = sum(1 for r in results if r not in existing_metadata)
384 if done % 20 == 0 or done == len(pending):
385 logger.info("[%d/%d] %s — %s (%s)",
386 done, len(pending),
387 result["song_name"], result["singer_name"], result["status"])
388
389 if progress is not None:
390 progress.close()
391
392 # 写 metadata.csv
393 with metadata_path.open("w", newline="", encoding="utf-8") as f:
394 writer = csv.DictWriter(f, fieldnames=METADATA_FIELDS)
395 writer.writeheader()
396 writer.writerows(results)
397
398 # 写 download_failures.csv
399 failures = [r for r in results if r.get("status") == "failed"]
400 if failures:
401 fail_path = output_dir / "download_failures.csv"
402 with fail_path.open("w", newline="", encoding="utf-8") as f:
403 writer = csv.DictWriter(f, fieldnames=METADATA_FIELDS, extrasaction="ignore")
404 writer.writeheader()
405 writer.writerows(failures)
406 logger.warning("失败 %d 条,见 %s", len(failures), fail_path)
407
408 # 生成 reference.csv 和 queries.csv
409 ok_results = [r for r in results if r.get("status") == "ok" and r.get("local_path")]
410
411 ref_rows, query_rows = [], []
412 for r in ok_results:
413 if int(r["is_main_version"]) == 1:
414 ref_rows.append({
415 "song_id": r["song_id"],
416 "audio_path": r["local_path"],
417 "variant": "original",
418 "song_name": r["song_name"],
419 "original_singer": r["original_singer"],
420 "lyricist": r["lyricist"],
421 "composer": r["composer"],
422 "genre": r["genre"],
423 "language_tag": r["language_tag"],
424 })
425 else:
426 platform = r["platform_name"] or "unknown"
427 query_rows.append({
428 "song_id": r["song_id"],
429 "audio_path": r["local_path"],
430 "variant": f"cover_{platform}",
431 "sample_class": "positive",
432 "expected_song_id": r["song_id"],
433 "expected": "duplicate",
434 "song_name": r["song_name"],
435 "original_singer": r["original_singer"],
436 "singer_name": r["singer_name"],
437 "platform_name": r["platform_name"],
438 "duration": r["duration"],
439 "is_main_version": r["is_main_version"],
440 })
441
442 ref_path = output_dir / "reference.csv"
443 with ref_path.open("w", newline="", encoding="utf-8") as f:
444 writer = csv.DictWriter(f, fieldnames=REFERENCE_FIELDS)
445 writer.writeheader()
446 writer.writerows(ref_rows)
447
448 query_path = output_dir / "queries.csv"
449 with query_path.open("w", newline="", encoding="utf-8") as f:
450 writer = csv.DictWriter(f, fieldnames=QUERY_FIELDS)
451 writer.writeheader()
452 writer.writerows(query_rows)
453
454 ok_count = sum(1 for r in results if r.get("status") == "ok")
455 song_count_ok = len({r["song_id"] for r in results if r.get("status") == "ok"})
456 logger.info("完成: 成功 %d 条,失败 %d 条,涉及 %d 首歌",
457 ok_count, len(failures), song_count_ok)
458 logger.info("参照集: %s(%d 条)", ref_path, len(ref_rows))
459 logger.info("查询集: %s(%d 条)", query_path, len(query_rows))
460 logger.info("元数据: %s", metadata_path)
461
462
463 if __name__ == "__main__":
464 main()