|
| 1 | +"""Parse real DJI .SRT telemetry dialects. |
| 2 | +
|
| 3 | +Supported layouts (documented public dialects): |
| 4 | +- Classic / Format 3 (+3b): HTML <font>-wrapped lines with SrtCnt/FrameCnt, |
| 5 | + datetime, and bracket fields [latitude:] [longitude:] [rel_alt: abs_alt:]. |
| 6 | +- Bracket-style / Format 1: Mini/Air/Mavic key-value brackets (same GPS keys; |
| 7 | + may include [iso], gb_yaw, and optional speed tags when present). |
| 8 | +- GPS function / Format 2 (+2b): GPS(lat,lon,alt[M]) / GPS (lat, lon, alt). |
| 9 | +- Compact RTK / Format 2c: single-line with GPS (...) and H.S Nm/s. |
| 10 | +
|
| 11 | +Speed: prefer recorded horizontal speed (H.S, [h_spd], [speed]) when present; |
| 12 | +otherwise estimate ground speed from successive GPS points (haversine). |
| 13 | +""" |
| 14 | + |
| 15 | +from __future__ import annotations |
| 16 | + |
| 17 | +import re |
| 18 | +from datetime import datetime, timezone |
| 19 | +from pathlib import Path |
| 20 | +from typing import List, Optional, Sequence, Tuple, Union |
| 21 | + |
| 22 | +from .models import TelemetryFrame |
| 23 | +from .speed import estimate_speeds, haversine_m |
| 24 | + |
| 25 | +PathLike = Union[str, Path] |
| 26 | + |
| 27 | +# Cue timing: 00:00:00,000 --> 00:00:00,033 (comma or period milliseconds) |
| 28 | +_TIME_RE = re.compile( |
| 29 | + r"(?P<sh>\d{1,2}):(?P<sm>\d{2}):(?P<ss>\d{2})[,.](?P<sms>\d{1,3})" |
| 30 | + r"\s*-->\s*" |
| 31 | + r"(?P<eh>\d{1,2}):(?P<em>\d{2}):(?P<es>\d{2})[,.](?P<ems>\d{1,3})" |
| 32 | +) |
| 33 | + |
| 34 | +_LAT_RE = re.compile(r"\[latitude\s*:\s*([+-]?\d+(?:\.\d+)?)\]", re.I) |
| 35 | +_LON_RE = re.compile(r"\[longitude\s*:\s*([+-]?\d+(?:\.\d+)?)\]", re.I) |
| 36 | +# rel_alt and abs_alt may share one bracket: [rel_alt: 1.300 abs_alt: 132.860] |
| 37 | +_REL_ALT_RE = re.compile(r"rel_alt\s*:\s*([+-]?\d+(?:\.\d+)?)", re.I) |
| 38 | +_ABS_ALT_RE = re.compile(r"abs_alt\s*:\s*([+-]?\d+(?:\.\d+)?)", re.I) |
| 39 | + |
| 40 | +# Format 2 / 2b: GPS(lat,lon,alt) or GPS(lat,lon,altM) |
| 41 | +_GPS_FUNC_RE = re.compile( |
| 42 | + r"GPS\s*\(\s*([+-]?\d+(?:\.\d+)?)\s*,\s*([+-]?\d+(?:\.\d+)?)\s*,\s*" |
| 43 | + r"([+-]?\d+(?:\.\d+)?)[A-Za-z]*\s*\)", |
| 44 | + re.I, |
| 45 | +) |
| 46 | +# Format 2c: GPS (lat, lon, alt) with spaces after commas |
| 47 | +_GPS_SPACE_RE = re.compile( |
| 48 | + r"GPS\s+\(\s*([+-]?\d+(?:\.\d+)?)\s*,\s*([+-]?\d+(?:\.\d+)?)\s*,\s*" |
| 49 | + r"([+-]?\d+(?:\.\d+)?)[A-Za-z]*\s*\)", |
| 50 | + re.I, |
| 51 | +) |
| 52 | + |
| 53 | +# Recorded horizontal speed variants |
| 54 | +_HS_RE = re.compile(r"H\.S\s*([+-]?\d+(?:\.\d+)?)\s*m/s", re.I) |
| 55 | +_H_SPD_RE = re.compile(r"\[h_spd\s*:\s*([+-]?\d+(?:\.\d+)?)\]", re.I) |
| 56 | +_SPEED_BRACKET_RE = re.compile(r"\[speed\s*:\s*([+-]?\d+(?:\.\d+)?)\]", re.I) |
| 57 | + |
| 58 | +# Datetime on its own line: 2024-01-15 14:30:22,123 or with . |
| 59 | +_DT_RE = re.compile( |
| 60 | + r"(?P<y>\d{4})-(?P<mo>\d{2})-(?P<d>\d{2})\s+" |
| 61 | + r"(?P<h>\d{2}):(?P<mi>\d{2}):(?P<s>\d{2})[,.](?P<ms>\d{1,3})" |
| 62 | +) |
| 63 | + |
| 64 | +_HTML_TAG_RE = re.compile(r"</?font[^>]*>", re.I) |
| 65 | + |
| 66 | + |
| 67 | +def _parse_ms(h: str, m: str, s: str, ms: str) -> int: |
| 68 | + ms_pad = (ms + "000")[:3] |
| 69 | + return ( |
| 70 | + int(h) * 3_600_000 |
| 71 | + + int(m) * 60_000 |
| 72 | + + int(s) * 1_000 |
| 73 | + + int(ms_pad) |
| 74 | + ) |
| 75 | + |
| 76 | + |
| 77 | +def _parse_datetime(text: str) -> Optional[datetime]: |
| 78 | + m = _DT_RE.search(text) |
| 79 | + if not m: |
| 80 | + return None |
| 81 | + ms = (m.group("ms") + "000")[:3] |
| 82 | + try: |
| 83 | + return datetime( |
| 84 | + int(m.group("y")), |
| 85 | + int(m.group("mo")), |
| 86 | + int(m.group("d")), |
| 87 | + int(m.group("h")), |
| 88 | + int(m.group("mi")), |
| 89 | + int(m.group("s")), |
| 90 | + int(ms) * 1000, |
| 91 | + tzinfo=timezone.utc, |
| 92 | + ) |
| 93 | + except ValueError: |
| 94 | + return None |
| 95 | + |
| 96 | + |
| 97 | +def _extract_speed(text: str) -> Optional[float]: |
| 98 | + for rx in (_HS_RE, _H_SPD_RE, _SPEED_BRACKET_RE): |
| 99 | + m = rx.search(text) |
| 100 | + if m: |
| 101 | + try: |
| 102 | + return float(m.group(1)) |
| 103 | + except ValueError: |
| 104 | + continue |
| 105 | + return None |
| 106 | + |
| 107 | + |
| 108 | +def _extract_gps(text: str) -> Optional[Tuple[float, float, Optional[float], Optional[float], Optional[float]]]: |
| 109 | + """Return (lat, lon, abs_alt, rel_alt, gps_alt_from_tuple) or None.""" |
| 110 | + lat_m = _LAT_RE.search(text) |
| 111 | + lon_m = _LON_RE.search(text) |
| 112 | + if lat_m and lon_m: |
| 113 | + lat = float(lat_m.group(1)) |
| 114 | + lon = float(lon_m.group(1)) |
| 115 | + abs_m = _ABS_ALT_RE.search(text) |
| 116 | + rel_m = _REL_ALT_RE.search(text) |
| 117 | + abs_alt = float(abs_m.group(1)) if abs_m else None |
| 118 | + rel_alt = float(rel_m.group(1)) if rel_m else None |
| 119 | + return lat, lon, abs_alt, rel_alt, None |
| 120 | + |
| 121 | + for rx in (_GPS_FUNC_RE, _GPS_SPACE_RE): |
| 122 | + m = rx.search(text) |
| 123 | + if m: |
| 124 | + lat = float(m.group(1)) |
| 125 | + lon = float(m.group(2)) |
| 126 | + gps_alt = float(m.group(3)) |
| 127 | + return lat, lon, None, None, gps_alt |
| 128 | + |
| 129 | + return None |
| 130 | + |
| 131 | + |
| 132 | +def _split_cues(content: str) -> List[Tuple[int, int, int, str]]: |
| 133 | + """Split SRT into (index, start_ms, end_ms, body_text) tuples.""" |
| 134 | + # Normalize newlines; strip BOM |
| 135 | + text = content.replace("\r\n", "\n").replace("\r", "\n") |
| 136 | + if text.startswith("\ufeff"): |
| 137 | + text = text[1:] |
| 138 | + |
| 139 | + blocks = re.split(r"\n\s*\n+", text.strip()) |
| 140 | + cues: List[Tuple[int, int, int, str]] = [] |
| 141 | + for block in blocks: |
| 142 | + lines = [ln for ln in block.split("\n") if ln.strip() != "" or True] |
| 143 | + # Keep structure: first non-empty may be index, then timing, then body |
| 144 | + nonempty = [ln for ln in block.split("\n")] |
| 145 | + if not any(ln.strip() for ln in nonempty): |
| 146 | + continue |
| 147 | + |
| 148 | + timing_idx = None |
| 149 | + start_ms = end_ms = 0 |
| 150 | + index = len(cues) + 1 |
| 151 | + for i, ln in enumerate(nonempty): |
| 152 | + tm = _TIME_RE.search(ln) |
| 153 | + if tm: |
| 154 | + timing_idx = i |
| 155 | + start_ms = _parse_ms(tm.group("sh"), tm.group("sm"), tm.group("ss"), tm.group("sms")) |
| 156 | + end_ms = _parse_ms(tm.group("eh"), tm.group("em"), tm.group("es"), tm.group("ems")) |
| 157 | + # Index line above timing if present |
| 158 | + if i > 0 and nonempty[i - 1].strip().isdigit(): |
| 159 | + index = int(nonempty[i - 1].strip()) |
| 160 | + break |
| 161 | + if timing_idx is None: |
| 162 | + continue |
| 163 | + body_lines = nonempty[timing_idx + 1 :] |
| 164 | + body = "\n".join(body_lines) |
| 165 | + body = _HTML_TAG_RE.sub("", body) |
| 166 | + cues.append((index, start_ms, end_ms, body)) |
| 167 | + return cues |
| 168 | + |
| 169 | + |
| 170 | +def parse_srt_text(content: str) -> List[TelemetryFrame]: |
| 171 | + """Parse SRT text into frames that have lat/lon. Skip GPS-less cues.""" |
| 172 | + frames: List[TelemetryFrame] = [] |
| 173 | + for index, start_ms, end_ms, body in _split_cues(content): |
| 174 | + gps = _extract_gps(body) |
| 175 | + if gps is None: |
| 176 | + continue |
| 177 | + lat, lon, abs_alt, rel_alt, gps_alt = gps |
| 178 | + # Prefer absolute altitude when present, else relative, else GPS tuple alt |
| 179 | + if abs_alt is not None: |
| 180 | + altitude_m = abs_alt |
| 181 | + altitude_source = "absolute" |
| 182 | + elif rel_alt is not None: |
| 183 | + altitude_m = rel_alt |
| 184 | + altitude_source = "relative" |
| 185 | + elif gps_alt is not None: |
| 186 | + # Format 2 GPS() altitude is typically MSL / absolute-ish |
| 187 | + altitude_m = gps_alt |
| 188 | + altitude_source = "absolute" |
| 189 | + abs_alt = gps_alt |
| 190 | + else: |
| 191 | + altitude_m = 0.0 |
| 192 | + altitude_source = "relative" |
| 193 | + |
| 194 | + recorded = _extract_speed(body) |
| 195 | + if recorded is not None: |
| 196 | + speed_mps = recorded |
| 197 | + speed_source = "recorded" |
| 198 | + else: |
| 199 | + speed_mps = None |
| 200 | + speed_source = "unknown" |
| 201 | + |
| 202 | + frames.append( |
| 203 | + TelemetryFrame( |
| 204 | + index=index, |
| 205 | + start_ms=start_ms, |
| 206 | + end_ms=end_ms, |
| 207 | + latitude=lat, |
| 208 | + longitude=lon, |
| 209 | + altitude_m=altitude_m, |
| 210 | + altitude_source=altitude_source, |
| 211 | + speed_mps=speed_mps, |
| 212 | + speed_source=speed_source, |
| 213 | + datetime_utc=_parse_datetime(body), |
| 214 | + rel_alt_m=rel_alt, |
| 215 | + abs_alt_m=abs_alt, |
| 216 | + raw_text=body, |
| 217 | + ) |
| 218 | + ) |
| 219 | + |
| 220 | + estimate_speeds(frames) |
| 221 | + return frames |
| 222 | + |
| 223 | + |
| 224 | +def parse_srt_file(path: PathLike) -> List[TelemetryFrame]: |
| 225 | + """Read and parse a .SRT file (read-only).""" |
| 226 | + p = Path(path) |
| 227 | + if not p.is_file(): |
| 228 | + raise FileNotFoundError(f"SRT file not found: {p}") |
| 229 | + # Try utf-8 then latin-1 for odd exports |
| 230 | + try: |
| 231 | + content = p.read_text(encoding="utf-8") |
| 232 | + except UnicodeDecodeError: |
| 233 | + content = p.read_text(encoding="latin-1") |
| 234 | + return parse_srt_text(content) |
| 235 | + |
| 236 | + |
| 237 | +# Re-export haversine for tests |
| 238 | +__all__ = ["parse_srt_text", "parse_srt_file", "haversine_m"] |
0 commit comments