test_composition_dedup.py
6.51 KB
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
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
"""作曲去重模块测试。
测试覆盖:
- Chromagram 提取
- 时间归一化输出维度
- Cosine 相似度计算
- 向量展开维度为 1536
"""
import os
import tempfile
import wave
import numpy as np
import pytest
from scipy.signal import resample
from composition_dedup.extractor import extract_chroma_feature, _normalize_audio_ffmpeg
from composition_dedup.similarity import (
CompositionSimilarity,
SimilarityDecision,
DUPLICATE_THRESHOLD,
SUSPECTED_THRESHOLD,
)
def _generate_test_wav(duration_sec: float = 1.0, sample_rate: int = 22050, frequency: float = 440.0) -> str:
"""生成测试用的 WAV 文件(正弦波)。
Args:
duration_sec: 持续时间(秒)。
sample_rate: 采样率。
frequency: 频率(Hz)。
Returns:
临时 WAV 文件路径。
"""
t = np.linspace(0, duration_sec, int(sample_rate * duration_sec), endpoint=False)
audio_data = (0.5 * np.sin(2 * np.pi * frequency * t)).astype(np.float32)
tmp_path = tempfile.mktemp(suffix=".wav")
with wave.open(tmp_path, "wb") as wf:
wf.setnchannels(1)
wf.setsampwidth(2) # 16-bit
wf.setframerate(sample_rate)
wf.writeframes((audio_data * 32767).astype(np.int16).tobytes())
return tmp_path
class TestChromaExtraction:
"""Chromagram 提取测试。"""
def test_extract_chroma_returns_1536_dim(self):
"""测试 Chromagram 提取返回 1536 维向量。"""
wav_path = _generate_test_wav(duration_sec=2.0, frequency=440.0)
try:
feature = extract_chroma_feature(wav_path)
assert isinstance(feature, np.ndarray)
assert feature.shape == (1536,), f"期望 (1536,), 实际 {feature.shape}"
assert feature.dtype == np.float32
finally:
if os.path.exists(wav_path):
os.remove(wav_path)
def test_extract_chroma_file_not_found(self):
"""测试不存在的音频文件抛出 FileNotFoundError。"""
with pytest.raises(FileNotFoundError):
extract_chroma_feature("/nonexistent/path/audio.mp3")
def test_extract_chroma_different_frequencies(self):
"""测试不同频率的音频产生不同特征。"""
wav_a = _generate_test_wav(duration_sec=2.0, frequency=440.0)
wav_b = _generate_test_wav(duration_sec=2.0, frequency=880.0)
try:
feature_a = extract_chroma_feature(wav_a)
feature_b = extract_chroma_feature(wav_b)
# 不同频率的音频特征不应完全相同
assert not np.allclose(feature_a, feature_b)
finally:
for path in [wav_a, wav_b]:
if os.path.exists(path):
os.remove(path)
def test_extract_chroma_same_audio_consistent(self):
"""测试同一音频多次提取结果一致。"""
wav_path = _generate_test_wav(duration_sec=1.0, frequency=440.0)
try:
feature_1 = extract_chroma_feature(wav_path)
feature_2 = extract_chroma_feature(wav_path)
np.testing.assert_array_almost_equal(feature_1, feature_2, decimal=5)
finally:
if os.path.exists(wav_path):
os.remove(wav_path)
class TestTimeNormalization:
"""时间归一化测试。"""
def test_resample_chroma_to_128_frames(self):
"""测试 Chromagram 时间归一化到 128 帧。"""
# 模拟不同长度的 Chromagram
for num_frames in [100, 256, 512, 1000, 2000]:
chroma = np.random.rand(12, num_frames).astype(np.float32)
if chroma.shape[1] != 128:
chroma = resample(chroma, 128, axis=1)
assert chroma.shape == (12, 128), f"帧数归一化失败: {chroma.shape}"
def test_flatten_to_1536(self):
"""测试展平后维度为 1536。"""
chroma = np.random.rand(12, 128).astype(np.float32)
feature = chroma.flatten()
assert feature.shape[0] == 12 * 128 == 1536
class TestCosineSimilarity:
"""Cosine 相似度计算测试。"""
def test_identical_vectors(self):
"""测试相同向量相似度为 1。"""
vec = np.random.rand(1536).astype(np.float32)
sim = CompositionSimilarity.cosine_similarity(vec, vec)
assert abs(sim - 1.0) < 1e-6
def test_orthogonal_vectors(self):
"""测试正交向量相似度接近 0。"""
vec_a = np.zeros(1536)
vec_a[0] = 1.0
vec_b = np.zeros(1536)
vec_b[1] = 1.0
sim = CompositionSimilarity.cosine_similarity(vec_a, vec_b)
assert abs(sim) < 1e-6
def test_zero_vector(self):
"""测试零向量返回 0 相似度。"""
vec_a = np.random.rand(1536).astype(np.float32)
vec_b = np.zeros(1536)
sim = CompositionSimilarity.cosine_similarity(vec_a, vec_b)
assert sim == 0.0
def test_similarity_range(self):
"""测试相似度值在 [0, 1] 范围内。"""
vec_a = np.random.rand(1536).astype(np.float32)
vec_b = np.random.rand(1536).astype(np.float32)
sim = CompositionSimilarity.cosine_similarity(vec_a, vec_b)
assert 0.0 <= sim <= 1.0
def test_classify_duplicate(self):
"""测试重复判定。"""
assert CompositionSimilarity.classify_similarity(0.96) == SimilarityDecision.DUPLICATE
assert CompositionSimilarity.classify_similarity(0.95) == SimilarityDecision.DUPLICATE
def test_classify_suspected(self):
"""测试疑似判定。"""
assert CompositionSimilarity.classify_similarity(0.94) == SimilarityDecision.SUSPECTED
assert CompositionSimilarity.classify_similarity(0.85) == SimilarityDecision.SUSPECTED
def test_classify_new(self):
"""测试非重复判定。"""
assert CompositionSimilarity.classify_similarity(0.84) == SimilarityDecision.NEW
assert CompositionSimilarity.classify_similarity(0.5) == SimilarityDecision.NEW
def test_compare_method(self):
"""测试 compare 方法同时返回相似度和判定。"""
vec = np.random.rand(1536).astype(np.float32)
sim, decision = CompositionSimilarity.compare(vec, vec)
assert abs(sim - 1.0) < 1e-6
assert decision == SimilarityDecision.DUPLICATE
class TestThresholds:
"""阈值常量测试。"""
def test_threshold_order(self):
"""测试阈值顺序正确。"""
assert DUPLICATE_THRESHOLD > SUSPECTED_THRESHOLD
def test_threshold_values(self):
"""测试阈值符合设计值。"""
assert DUPLICATE_THRESHOLD == 0.95
assert SUSPECTED_THRESHOLD == 0.85