257 lines
9.0 KiB
Python
257 lines
9.0 KiB
Python
#!/usr/bin/env python3
|
||
"""Audit an audiobook/EPUB pair before building a single-speaker TTS corpus."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import collections
|
||
import difflib
|
||
import json
|
||
import re
|
||
import subprocess
|
||
import sys
|
||
import xml.etree.ElementTree as ET
|
||
import zipfile
|
||
from pathlib import Path, PurePosixPath
|
||
|
||
|
||
WORD_RE = re.compile(r"[а-яёa-z0-9]+", re.IGNORECASE)
|
||
WHITESPACE_RE = re.compile(r"\s+")
|
||
|
||
|
||
def normalized_words(text: str) -> list[str]:
|
||
return [word.replace("ё", "е") for word in WORD_RE.findall(text.lower())]
|
||
|
||
|
||
def element_text(root: ET.Element) -> str:
|
||
ignored = {"script", "style", "svg", "math"}
|
||
parts: list[str] = []
|
||
|
||
def visit(node: ET.Element) -> None:
|
||
tag = node.tag.rsplit("}", 1)[-1].lower()
|
||
if tag in ignored:
|
||
return
|
||
if node.text:
|
||
parts.append(node.text)
|
||
for child in node:
|
||
visit(child)
|
||
if child.tail:
|
||
parts.append(child.tail)
|
||
if tag in {"p", "div", "section", "h1", "h2", "h3", "h4", "h5", "h6", "li", "br"}:
|
||
parts.append("\n")
|
||
|
||
visit(root)
|
||
lines = [WHITESPACE_RE.sub(" ", line).strip() for line in "".join(parts).splitlines()]
|
||
return "\n".join(line for line in lines if line)
|
||
|
||
|
||
def extract_epub(epub_path: Path, output_dir: Path) -> tuple[str, list[dict[str, object]]]:
|
||
with zipfile.ZipFile(epub_path) as archive:
|
||
container = ET.fromstring(archive.read("META-INF/container.xml"))
|
||
rootfile = next(
|
||
node.attrib["full-path"]
|
||
for node in container.iter()
|
||
if node.tag.rsplit("}", 1)[-1] == "rootfile"
|
||
)
|
||
opf = ET.fromstring(archive.read(rootfile))
|
||
opf_dir = PurePosixPath(rootfile).parent
|
||
manifest = {
|
||
node.attrib["id"]: node.attrib["href"]
|
||
for node in opf.iter()
|
||
if node.tag.rsplit("}", 1)[-1] == "item" and "id" in node.attrib and "href" in node.attrib
|
||
}
|
||
spine_ids = [
|
||
node.attrib["idref"]
|
||
for node in opf.iter()
|
||
if node.tag.rsplit("}", 1)[-1] == "itemref" and "idref" in node.attrib
|
||
]
|
||
sections: list[dict[str, object]] = []
|
||
texts: list[str] = []
|
||
for index, item_id in enumerate(spine_ids):
|
||
href = manifest.get(item_id)
|
||
if not href:
|
||
continue
|
||
member = str(opf_dir / PurePosixPath(href))
|
||
try:
|
||
root = ET.fromstring(archive.read(member))
|
||
except (KeyError, ET.ParseError):
|
||
continue
|
||
text = element_text(root)
|
||
if not text:
|
||
continue
|
||
headings = [
|
||
WHITESPACE_RE.sub(" ", "".join(node.itertext())).strip()
|
||
for node in root.iter()
|
||
if node.tag.rsplit("}", 1)[-1].lower() in {"h1", "h2", "h3"}
|
||
]
|
||
sections.append(
|
||
{
|
||
"index": index,
|
||
"href": member,
|
||
"characters": len(text),
|
||
"words": len(normalized_words(text)),
|
||
"headings": [heading for heading in headings if heading],
|
||
}
|
||
)
|
||
texts.append(text)
|
||
|
||
book_text = "\n\n".join(texts)
|
||
output_dir.mkdir(parents=True, exist_ok=True)
|
||
(output_dir / "book.txt").write_text(book_text, encoding="utf-8")
|
||
(output_dir / "spine.json").write_text(
|
||
json.dumps(sections, ensure_ascii=False, indent=2), encoding="utf-8"
|
||
)
|
||
return book_text, sections
|
||
|
||
|
||
def best_text_match(transcript: str, book_words: list[str]) -> dict[str, object]:
|
||
spoken = normalized_words(transcript)
|
||
if not spoken or not book_words:
|
||
return {"ratio": 0.0, "book_excerpt": "", "word_offset": None}
|
||
positions: dict[str, list[int]] = collections.defaultdict(list)
|
||
for index, word in enumerate(book_words):
|
||
if len(word) >= 4:
|
||
positions[word].append(index)
|
||
votes: collections.Counter[int] = collections.Counter()
|
||
for spoken_index, word in enumerate(spoken):
|
||
candidates = positions.get(word, ())
|
||
if len(candidates) <= 200:
|
||
votes.update(book_index - spoken_index for book_index in candidates)
|
||
offsets = [offset for offset, _ in votes.most_common(30)] or [0]
|
||
best_ratio = 0.0
|
||
best_offset = 0
|
||
best_window: list[str] = []
|
||
for offset in offsets:
|
||
start = max(0, offset - 8)
|
||
window = book_words[start : start + len(spoken) + 16]
|
||
ratio = difflib.SequenceMatcher(None, spoken, window, autojunk=False).ratio()
|
||
if ratio > best_ratio:
|
||
best_ratio, best_offset, best_window = ratio, start, window
|
||
return {
|
||
"ratio": round(best_ratio, 4),
|
||
"book_excerpt": " ".join(best_window),
|
||
"word_offset": best_offset,
|
||
}
|
||
|
||
|
||
def transcribe_samples(
|
||
audio_files: list[Path],
|
||
book_text: str,
|
||
output_dir: Path,
|
||
model_name: str,
|
||
sample_seconds: int,
|
||
) -> list[dict[str, object]]:
|
||
try:
|
||
import av
|
||
import imageio_ffmpeg
|
||
from faster_whisper import WhisperModel
|
||
except ImportError as error:
|
||
raise SystemExit(
|
||
"Install audit dependencies into PYTHONPATH: faster-whisper imageio-ffmpeg"
|
||
) from error
|
||
|
||
chosen = [audio_files[0], audio_files[len(audio_files) // 2], audio_files[-1]]
|
||
samples_dir = output_dir / "samples"
|
||
samples_dir.mkdir(parents=True, exist_ok=True)
|
||
ffmpeg = imageio_ffmpeg.get_ffmpeg_exe()
|
||
model = WhisperModel(
|
||
model_name,
|
||
device="cpu",
|
||
compute_type="int8",
|
||
download_root=str(output_dir / "models"),
|
||
)
|
||
book_words = normalized_words(book_text)
|
||
results: list[dict[str, object]] = []
|
||
for source in chosen:
|
||
with av.open(str(source)) as media:
|
||
stream = media.streams.audio[0]
|
||
duration = float(stream.duration * stream.time_base) if stream.duration is not None else 0.0
|
||
start = min(max(45.0, duration * 0.35), max(0.0, duration - sample_seconds - 5.0))
|
||
wav_path = samples_dir / f"{source.stem[:4]}-{int(start):05d}.wav"
|
||
subprocess.run(
|
||
[
|
||
ffmpeg,
|
||
"-hide_banner",
|
||
"-loglevel",
|
||
"error",
|
||
"-y",
|
||
"-ss",
|
||
f"{start:.3f}",
|
||
"-i",
|
||
str(source),
|
||
"-t",
|
||
str(sample_seconds),
|
||
"-ac",
|
||
"1",
|
||
"-ar",
|
||
"16000",
|
||
str(wav_path),
|
||
],
|
||
check=True,
|
||
)
|
||
segments, info = model.transcribe(
|
||
str(wav_path),
|
||
language="ru",
|
||
beam_size=5,
|
||
vad_filter=True,
|
||
condition_on_previous_text=True,
|
||
)
|
||
transcript = " ".join(segment.text.strip() for segment in segments).strip()
|
||
result = {
|
||
"source": source.name,
|
||
"source_duration_seconds": round(duration, 3),
|
||
"sample_start_seconds": round(start, 3),
|
||
"sample_duration_seconds": sample_seconds,
|
||
"detected_language": info.language,
|
||
"language_probability": round(info.language_probability, 4),
|
||
"transcript": transcript,
|
||
}
|
||
result.update(best_text_match(transcript, book_words))
|
||
results.append(result)
|
||
(output_dir / "alignment_samples.json").write_text(
|
||
json.dumps(results, ensure_ascii=False, indent=2), encoding="utf-8"
|
||
)
|
||
return results
|
||
|
||
|
||
def main() -> int:
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument("--source-dir", type=Path, required=True)
|
||
parser.add_argument("--epub", type=Path)
|
||
parser.add_argument("--output-dir", type=Path, required=True)
|
||
parser.add_argument("--model", default="base")
|
||
parser.add_argument("--sample-seconds", type=int, default=75)
|
||
parser.add_argument("--skip-asr", action="store_true")
|
||
args = parser.parse_args()
|
||
|
||
audio_files = sorted(args.source_dir.glob("*.mp3"))
|
||
epub_path = args.epub or next(args.source_dir.glob("*.epub"), None)
|
||
if not audio_files:
|
||
parser.error("No MP3 files found")
|
||
if epub_path is None or not epub_path.is_file():
|
||
parser.error("EPUB file not found")
|
||
|
||
book_text, sections = extract_epub(epub_path, args.output_dir)
|
||
report: dict[str, object] = {
|
||
"source_dir": str(args.source_dir.resolve()),
|
||
"epub": str(epub_path.resolve()),
|
||
"audio_files": [str(path.resolve()) for path in audio_files],
|
||
"audio_file_count": len(audio_files),
|
||
"epub_section_count": len(sections),
|
||
"book_characters": len(book_text),
|
||
"book_words": len(normalized_words(book_text)),
|
||
}
|
||
if not args.skip_asr:
|
||
report["alignment_samples"] = transcribe_samples(
|
||
audio_files, book_text, args.output_dir, args.model, args.sample_seconds
|
||
)
|
||
(args.output_dir / "audit.json").write_text(
|
||
json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8"
|
||
)
|
||
print(json.dumps(report, ensure_ascii=False, indent=2))
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main())
|