File size: 4,639 Bytes
1425afc | 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 182 183 184 185 186 187 188 | from typing import List, Dict, Any, Union
# =====================================================
# PUBLIC API (USED BY MAIN.PY)
# =====================================================
def generate_srt(data: List[Dict[str, Any]]) -> str:
"""
Universal SRT generator for:
- Whisper word output (V1–V7)
- Highlight segments (start/end grouped words)
- Mixed/partial structures
Expected input formats:
1. Word-level:
{"text": "...", "start": float, "end": float}
2. Segment-level:
[{"start": float, "end": float, "text": "..."}]
Returns:
SRT formatted string
"""
if not data:
return ""
normalized = _normalize_input(data)
return _build_srt(normalized)
# =====================================================
# NORMALIZATION LAYER (CRITICAL FOR V1–V7 COMPATIBILITY)
# =====================================================
def _normalize_input(data: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
"""
Converts any supported structure into unified subtitle blocks
"""
normalized = []
# CASE 1: Already segment-based
if isinstance(data[0], dict) and "start" in data[0] and "end" in data[0] and "text" in data[0]:
for item in data:
normalized.append({
"start": float(item.get("start", 0)),
"end": float(item.get("end", 0)),
"text": str(item.get("text", "")).strip()
})
return normalized
# CASE 2: Whisper word-level output
buffer = []
current_start = None
for w in data:
if not isinstance(w, dict):
continue
text = str(w.get("text", "")).strip()
start = w.get("start", None)
end = w.get("end", None)
if start is None or end is None:
continue
if current_start is None:
current_start = start
buffer.append(text)
# Chunking strategy: group every ~8–12 words
if len(buffer) >= 10:
normalized.append({
"start": current_start,
"end": end,
"text": " ".join(buffer)
})
buffer = []
current_start = None
# flush remaining buffer
if buffer:
normalized.append({
"start": current_start or 0,
"end": data[-1].get("end", 0),
"text": " ".join(buffer)
})
return normalized
# =====================================================
# SRT BUILDER
# =====================================================
def _build_srt(items: List[Dict[str, Any]]) -> str:
"""
Converts normalized subtitle blocks → SRT format
"""
output = []
index = 1
for item in items:
start = _format_time(item["start"])
end = _format_time(item["end"])
text = _clean_text(item["text"])
if not text:
continue
output.append(f"{index}")
output.append(f"{start} --> {end}")
output.append(f"{text}")
output.append("") # blank line separator
index += 1
return "\n".join(output).strip()
# =====================================================
# TIME FORMATTER
# =====================================================
def _format_time(seconds: Union[int, float]) -> str:
"""
Converts seconds → SRT timestamp format
HH:MM:SS,mmm
"""
try:
seconds = float(seconds)
except:
seconds = 0.0
hrs = int(seconds // 3600)
mins = int((seconds % 3600) // 60)
secs = int(seconds % 60)
ms = int((seconds - int(seconds)) * 1000)
return f"{hrs:02}:{mins:02}:{secs:02},{ms:03}"
# =====================================================
# TEXT CLEANER (IMPORTANT FOR VIDEO RENDERING STABILITY)
# =====================================================
def _clean_text(text: str) -> str:
"""
Sanitizes subtitle text for rendering engines
"""
if not text:
return ""
text = text.replace("\n", " ")
text = text.replace("\r", " ")
# remove excessive spacing
text = " ".join(text.split())
return text.strip()
# =====================================================
# OPTIONAL DEBUG HELPER (SAFE IN PRODUCTION)
# =====================================================
def debug_srt(data: List[Dict[str, Any]]) -> dict:
"""
Returns structured preview for debugging pipelines
"""
normalized = _normalize_input(data)
return {
"blocks": len(normalized),
"sample": normalized[:3],
"duration": normalized[-1]["end"] if normalized else 0
} |