mapillary-tools 0.13.3a1__py3-none-any.whl → 0.14.0a1__py3-none-any.whl

mapillary_tools/camm/camm_parser.py

@@ -1,27 +1,29 @@
 # pyre-ignore-all-errors[5, 11, 16, 21, 24, 58]
+from __future__ import annotations
 
+import abc
 import dataclasses
 import io
 import logging
-import pathlib
 import typing as T
 from enum import Enum
 
 import construct as C
+from typing_extensions import TypeIs
 
 from .. import geo, telemetry
-from ..mp4 import simple_mp4_parser as sparser
 from ..mp4.mp4_sample_parser import MovieBoxParser, Sample, TrackBoxParser
 
 
 LOG = logging.getLogger(__name__)
+# All fields are little-endian
+_Float = C.Float32l
+_Double = C.Float64l
 
 
 TelemetryMeasurement = T.Union[
     geo.Point,
-    telemetry.AccelerationData,
-    telemetry.GyroscopeData,
-    telemetry.MagnetometerData,
+    telemetry.TimestampedMeasurement,
 ]
 
 
@@ -37,99 +39,399 @@ class CAMMType(Enum):
     MAGNETIC_FIELD = 7
 
 
-# All fields are little-endian
-Float = C.Float32l
-Double = C.Float64l
+TTelemetry = T.TypeVar("TTelemetry", bound=TelemetryMeasurement)
 
-_SWITCH: T.Dict[int, C.Struct] = {
-    # angle_axis
-    CAMMType.ANGLE_AXIS.value: Float[3],
-    CAMMType.EXPOSURE_TIME.value: C.Struct(
-        "pixel_exposure_time" / C.Int32sl,
-        "rolling_shutter_skew_time" / C.Int32sl,
-    ),
-    # gyro
-    CAMMType.GYRO.value: Float[3],
-    # acceleration
-    CAMMType.ACCELERATION.value: Float[3],
-    # position
-    CAMMType.POSITION.value: Float[3],
-    # lat, lon, alt
-    CAMMType.MIN_GPS.value: Double[3],
-    CAMMType.GPS.value: C.Struct(
-        "time_gps_epoch" / Double,
-        "gps_fix_type" / C.Int32sl,
-        "latitude" / Double,
-        "longitude" / Double,
-        "altitude" / Float,
-        "horizontal_accuracy" / Float,
-        "vertical_accuracy" / Float,
-        "velocity_east" / Float,
-        "velocity_north" / Float,
-        "velocity_up" / Float,
-        "speed_accuracy" / Float,
-    ),
-    # magnetic_field
-    CAMMType.MAGNETIC_FIELD.value: Float[3],
-}
 
-CAMMSampleData = C.Struct(
-    C.Padding(2),
-    "type" / C.Int16ul,
-    "data"
-    / C.Switch(
-        C.this.type,
-        _SWITCH,
-    ),
-)
+@dataclasses.dataclass
+class CAMMInfo:
+    # None indicates the data has been extracted,
+    # while [] indicates extracetd but no data point found
+    mini_gps: list[geo.Point] | None = None
+    gps: list[telemetry.CAMMGPSPoint] | None = None
+    accl: list[telemetry.AccelerationData] | None = None
+    gyro: list[telemetry.GyroscopeData] | None = None
+    magn: list[telemetry.MagnetometerData] | None = None
+    make: str = ""
+    model: str = ""
 
 
-def _parse_telemetry_from_sample(
-    fp: T.BinaryIO, sample: Sample
-) -> T.Optional[TelemetryMeasurement]:
-    fp.seek(sample.raw_sample.offset, io.SEEK_SET)
-    data = fp.read(sample.raw_sample.size)
-    box = CAMMSampleData.parse(data)
-    if box.type == CAMMType.MIN_GPS.value:
+def extract_camm_info(fp: T.BinaryIO, telemetry_only: bool = False) -> CAMMInfo | None:
+    moov = MovieBoxParser.parse_stream(fp)
+
+    make, model = "", ""
+    if not telemetry_only:
+        udta_boxdata = moov.extract_udta_boxdata()
+        if udta_boxdata is not None:
+            make, model = _extract_camera_make_and_model_from_utda_boxdata(udta_boxdata)
+
+    gps_only_construct = _construct_with_selected_camm_types(
+        [CAMMType.MIN_GPS, CAMMType.GPS]
+    )
+    # Optimization: skip parsing sample data smaller than 16 bytes
+    # because we are only interested in MIN_GPS and GPS which are larger than 16 bytes
+    MIN_GPS_SAMPLE_SIZE = 17
+
+    for track in moov.extract_tracks():
+        if _contains_camm_description(track):
+            if telemetry_only:
+                maybe_measurements = (
+                    _parse_telemetry_from_sample(fp, sample)
+                    for sample in track.extract_samples()
+                    if _is_camm_description(sample.description)
+                )
+                measurements = _filter_telemetry_by_track_elst(
+                    moov, track, (m for m in maybe_measurements if m is not None)
+                )
+
+                accl: list[telemetry.AccelerationData] = []
+                gyro: list[telemetry.GyroscopeData] = []
+                magn: list[telemetry.MagnetometerData] = []
+
+                for measurement in measurements:
+                    if isinstance(measurement, telemetry.AccelerationData):
+                        accl.append(measurement)
+                    elif isinstance(measurement, telemetry.GyroscopeData):
+                        gyro.append(measurement)
+                    elif isinstance(measurement, telemetry.MagnetometerData):
+                        magn.append(measurement)
+
+                return CAMMInfo(accl=accl, gyro=gyro, magn=magn)
+            else:
+                maybe_measurements = (
+                    _parse_telemetry_from_sample(fp, sample, gps_only_construct)
+                    for sample in track.extract_samples()
+                    if _is_camm_description(sample.description)
+                    and sample.raw_sample.size >= MIN_GPS_SAMPLE_SIZE
+                )
+                measurements = _filter_telemetry_by_track_elst(
+                    moov, track, (m for m in maybe_measurements if m is not None)
+                )
+
+                mini_gps: list[geo.Point] = []
+                gps: list[telemetry.CAMMGPSPoint] = []
+
+                for measurement in measurements:
+                    if isinstance(measurement, geo.Point):
+                        mini_gps.append(measurement)
+                    elif isinstance(measurement, telemetry.CAMMGPSPoint):
+                        gps.append(measurement)
+
+                return CAMMInfo(mini_gps=mini_gps, gps=gps, make=make, model=model)
+
+    return None
+
+
+def extract_camera_make_and_model(fp: T.BinaryIO) -> tuple[str, str]:
+    moov = MovieBoxParser.parse_stream(fp)
+    udta_boxdata = moov.extract_udta_boxdata()
+    if udta_boxdata is None:
+        return "", ""
+    return _extract_camera_make_and_model_from_utda_boxdata(udta_boxdata)
+
+
+class CAMMSampleEntry(abc.ABC, T.Generic[TTelemetry]):
+    serialized_camm_type: CAMMType
+
+    telemetry_cls_type: T.Type[TTelemetry]
+
+    construct: C.Struct
+
+    @classmethod
+    def serializable(cls, data: T.Any, throw: bool = False) -> TypeIs[TTelemetry]:
+        # Use "is" for exact type match, instead of isinstance
+        if type(data) is cls.telemetry_cls_type:
+            return True
+
+        if throw:
+            raise TypeError(
+                f"{cls} can not serialize {type(data)}: expect {cls.telemetry_cls_type}"
+            )
+
+        return False
+
+    @classmethod
+    @abc.abstractmethod
+    def serialize(cls, data: TTelemetry) -> bytes:
+        raise NotImplementedError
+
+    @classmethod
+    @abc.abstractmethod
+    def deserialize(cls, sample: Sample, data: T.Any) -> TTelemetry:
+        raise NotImplementedError
+
+
+class MinGPSSampleEntry(CAMMSampleEntry):
+    serialized_camm_type = CAMMType.MIN_GPS
+
+    telemetry_cls_type = geo.Point
+
+    construct = _Double[3]  # type: ignore
+
+    @classmethod
+    def deserialize(cls, sample: Sample, data: T.Any) -> geo.Point:
         return geo.Point(
             time=sample.exact_time,
-            lat=box.data[0],
-            lon=box.data[1],
-            alt=box.data[2],
+            lat=data[0],
+            lon=data[1],
+            alt=data[2],
             angle=None,
         )
-    elif box.type == CAMMType.GPS.value:
-        # Not using box.data.time_gps_epoch as the point timestamp
-        # because it is from another clock
-        return geo.Point(
+
+    @classmethod
+    def serialize(cls, data: geo.Point) -> bytes:
+        cls.serializable(data, throw=True)
+
+        return CAMMSampleData.build(
+            {
+                "type": cls.serialized_camm_type.value,
+                "data": [
+                    data.lat,
+                    data.lon,
+                    -1.0 if data.alt is None else data.alt,
+                ],
+            }
+        )
+
+
+class GPSSampleEntry(CAMMSampleEntry):
+    serialized_camm_type: CAMMType = CAMMType.GPS
+
+    telemetry_cls_type = telemetry.CAMMGPSPoint
+
+    construct = C.Struct(
+        "time_gps_epoch" / _Double,  # type: ignore
+        "gps_fix_type" / C.Int32sl,  # type: ignore
+        "latitude" / _Double,  # type: ignore
+        "longitude" / _Double,  # type: ignore
+        "altitude" / _Float,  # type: ignore
+        "horizontal_accuracy" / _Float,  # type: ignore
+        "vertical_accuracy" / _Float,  # type: ignore
+        "velocity_east" / _Float,  # type: ignore
+        "velocity_north" / _Float,  # type: ignore
+        "velocity_up" / _Float,  # type: ignore
+        "speed_accuracy" / _Float,  # type: ignore
+    )
+
+    @classmethod
+    def deserialize(cls, sample: Sample, data: T.Any) -> telemetry.CAMMGPSPoint:
+        return telemetry.CAMMGPSPoint(
             time=sample.exact_time,
-            lat=box.data.latitude,
-            lon=box.data.longitude,
-            alt=box.data.altitude,
+            lat=data.latitude,
+            lon=data.longitude,
+            alt=data.altitude,
             angle=None,
+            time_gps_epoch=data.time_gps_epoch,
+            gps_fix_type=data.gps_fix_type,
+            horizontal_accuracy=data.horizontal_accuracy,
+            vertical_accuracy=data.vertical_accuracy,
+            velocity_east=data.velocity_east,
+            velocity_north=data.velocity_north,
+            velocity_up=data.velocity_up,
+            speed_accuracy=data.speed_accuracy,
+        )
+
+    @classmethod
+    def serialize(cls, data: telemetry.CAMMGPSPoint) -> bytes:
+        cls.serializable(data, throw=True)
+
+        return CAMMSampleData.build(
+            {
+                "type": cls.serialized_camm_type.value,
+                "data": {
+                    "time_gps_epoch": data.time_gps_epoch,
+                    "gps_fix_type": data.gps_fix_type,
+                    "latitude": data.lat,
+                    "longitude": data.lon,
+                    "altitude": -1.0 if data.alt is None else data.alt,
+                    "horizontal_accuracy": data.horizontal_accuracy,
+                    "vertical_accuracy": data.vertical_accuracy,
+                    "velocity_east": data.velocity_east,
+                    "velocity_north": data.velocity_north,
+                    "velocity_up": data.velocity_up,
+                    "speed_accuracy": data.speed_accuracy,
+                },
+            }
         )
-    elif box.type == CAMMType.ACCELERATION.value:
+
+
+class GoProGPSSampleEntry(CAMMSampleEntry):
+    serialized_camm_type: CAMMType = CAMMType.MIN_GPS
+
+    telemetry_cls_type = telemetry.GPSPoint
+
+    construct = _Double[3]  # type: ignore
+
+    @classmethod
+    def deserialize(cls, sample: Sample, data: T.Any) -> telemetry.GPSPoint:
+        raise NotImplementedError("Deserializing GoPro GPS Point is not supported")
+
+    @classmethod
+    def serialize(cls, data: telemetry.GPSPoint) -> bytes:
+        cls.serializable(data, throw=True)
+
+        return CAMMSampleData.build(
+            {
+                "type": cls.serialized_camm_type.value,
+                "data": [
+                    data.lat,
+                    data.lon,
+                    -1.0 if data.alt is None else data.alt,
+                ],
+            }
+        )
+
+
+class AccelerationSampleEntry(CAMMSampleEntry):
+    serialized_camm_type: CAMMType = CAMMType.ACCELERATION
+
+    telemetry_cls_type = telemetry.AccelerationData
+
+    construct: C.Struct = _Float[3]  # type: ignore
+
+    @classmethod
+    def deserialize(cls, sample: Sample, data: T.Any) -> telemetry.AccelerationData:
         return telemetry.AccelerationData(
             time=sample.exact_time,
-            x=box.data[0],
-            y=box.data[1],
-            z=box.data[2],
+            x=data[0],
+            y=data[1],
+            z=data[2],
         )
-    elif box.type == CAMMType.GYRO.value:
+
+    @classmethod
+    def serialize(cls, data: telemetry.AccelerationData) -> bytes:
+        cls.serializable(data, throw=True)
+
+        return CAMMSampleData.build(
+            {
+                "type": cls.serialized_camm_type.value,
+                "data": [data.x, data.y, data.z],
+            }
+        )
+
+
+class GyroscopeSampleEntry(CAMMSampleEntry):
+    serialized_camm_type: CAMMType = CAMMType.GYRO
+
+    telemetry_cls_type = telemetry.GyroscopeData
+
+    construct: C.Struct = _Float[3]  # type: ignore
+
+    @classmethod
+    def deserialize(cls, sample: Sample, data: T.Any) -> telemetry.GyroscopeData:
         return telemetry.GyroscopeData(
             time=sample.exact_time,
-            x=box.data[0],
-            y=box.data[1],
-            z=box.data[2],
+            x=data[0],
+            y=data[1],
+            z=data[2],
+        )
+
+    @classmethod
+    def serialize(cls, data: telemetry.GyroscopeData) -> bytes:
+        cls.serializable(data)
+
+        return CAMMSampleData.build(
+            {
+                "type": cls.serialized_camm_type.value,
+                "data": [data.x, data.y, data.z],
+            }
         )
-    elif box.type == CAMMType.MAGNETIC_FIELD.value:
+
+
+class MagnetometerSampleEntry(CAMMSampleEntry):
+    serialized_camm_type: CAMMType = CAMMType.MAGNETIC_FIELD
+
+    telemetry_cls_type = telemetry.MagnetometerData
+
+    construct: C.Struct = _Float[3]  # type: ignore
+
+    @classmethod
+    def deserialize(cls, sample: Sample, data: T.Any) -> telemetry.MagnetometerData:
         return telemetry.MagnetometerData(
             time=sample.exact_time,
-            x=box.data[0],
-            y=box.data[1],
-            z=box.data[2],
+            x=data[0],
+            y=data[1],
+            z=data[2],
         )
-    return None
+
+    @classmethod
+    def serialize(cls, data: telemetry.MagnetometerData) -> bytes:
+        cls.serializable(data)
+
+        return CAMMSampleData.build(
+            {
+                "type": cls.serialized_camm_type.value,
+                "data": [data.x, data.y, data.z],
+            }
+        )
+
+
+SAMPLE_ENTRY_CLS_BY_CAMM_TYPE = {
+    sample_entry_cls.serialized_camm_type: sample_entry_cls
+    for sample_entry_cls in CAMMSampleEntry.__subclasses__()
+    if sample_entry_cls not in [GoProGPSSampleEntry]
+}
+assert len(SAMPLE_ENTRY_CLS_BY_CAMM_TYPE) == 5, SAMPLE_ENTRY_CLS_BY_CAMM_TYPE.keys()
+
+
+_SWITCH: T.Dict[int, C.Struct] = {
+    # Angle_axis
+    CAMMType.ANGLE_AXIS.value: _Float[3],  # type: ignore
+    # Exposure time
+    CAMMType.EXPOSURE_TIME.value: C.Struct(
+        "pixel_exposure_time" / C.Int32sl,  # type: ignore
+        "rolling_shutter_skew_time" / C.Int32sl,  # type: ignore
+    ),
+    # Position
+    CAMMType.POSITION.value: _Float[3],  # type: ignore
+    # Serializable types
+    **{t.value: cls.construct for t, cls in SAMPLE_ENTRY_CLS_BY_CAMM_TYPE.items()},
+}
+
+
+def _construct_with_selected_camm_types(
+    selected_camm_types: T.Container[CAMMType] | None = None,
+) -> C.Struct:
+    if selected_camm_types is None:
+        switch = _SWITCH
+    else:
+        switch = {
+            k: v for k, v in _SWITCH.items() if CAMMType(k) in selected_camm_types
+        }
+
+    return C.Struct(
+        C.Padding(2),
+        "type" / C.Int16ul,
+        "data" / C.Switch(C.this.type, switch),
+    )
+
+
+CAMMSampleData = _construct_with_selected_camm_types()
+
+
+def _parse_telemetry_from_sample(
+    fp: T.BinaryIO,
+    sample: Sample,
+    construct: C.Struct | None = None,
+) -> TelemetryMeasurement | None:
+    if construct is None:
+        construct = CAMMSampleData
+
+    fp.seek(sample.raw_sample.offset, io.SEEK_SET)
+    data = fp.read(sample.raw_sample.size)
+
+    box = construct.parse(data)
+
+    # boxdata=None when the construct is unable to parse the data
+    # (CAMM type not in the switch)
+    if box.data is None:
+        return None
+
+    camm_type = CAMMType(box.type)  # type: ignore
+    SampleKlass = SAMPLE_ENTRY_CLS_BY_CAMM_TYPE.get(camm_type)
+    if SampleKlass is None:
+        return None
+
+    return SampleKlass.deserialize(sample, box.data)
 
 
 def _filter_telemetry_by_elst_segments(
@@ -188,7 +490,7 @@ def _filter_telemetry_by_track_elst(
     moov: MovieBoxParser,
     track: TrackBoxParser,
     measurements: T.Iterable[TelemetryMeasurement],
-) -> T.List[TelemetryMeasurement]:
+) -> list[TelemetryMeasurement]:
     elst_boxdata = track.extract_elst_boxdata()
 
     if elst_boxdata is not None:
@@ -216,127 +518,73 @@ def _filter_telemetry_by_track_elst(
     return list(measurements)
 
 
-def extract_points(fp: T.BinaryIO) -> T.Optional[T.List[geo.Point]]:
-    """
-    Return a list of points (could be empty) if it is a valid CAMM video,
-    otherwise None
-    """
-
-    moov = MovieBoxParser.parse_stream(fp)
-
-    for track in moov.extract_tracks():
-        if _contains_camm_description(track):
-            maybe_measurements = (
-                _parse_telemetry_from_sample(fp, sample)
-                for sample in track.extract_samples()
-                if _is_camm_description(sample.description)
-            )
-            points = [m for m in maybe_measurements if isinstance(m, geo.Point)]
-
-            return T.cast(
-                T.List[geo.Point], _filter_telemetry_by_track_elst(moov, track, points)
-            )
-
-    return None
-
-
-def extract_telemetry_data(fp: T.BinaryIO) -> T.Optional[T.List[TelemetryMeasurement]]:
-    moov = MovieBoxParser.parse_stream(fp)
-
-    for track in moov.extract_tracks():
-        if _contains_camm_description(track):
-            maybe_measurements = (
-                _parse_telemetry_from_sample(fp, sample)
-                for sample in track.extract_samples()
-                if _is_camm_description(sample.description)
-            )
-            measurements = [m for m in maybe_measurements if m is not None]
-
-            measurements = _filter_telemetry_by_track_elst(moov, track, measurements)
-
-            return measurements
-
-    return None
-
-
-def parse_gpx(path: pathlib.Path) -> T.List[geo.Point]:
-    with path.open("rb") as fp:
-        points = extract_points(fp)
-    if points is None:
-        return []
-    return points
-
-
-MakeOrModel = C.Struct(
+_MakeOrModel = C.Struct(
     "size" / C.Int16ub,
     C.Padding(2),
     "data" / C.FixedSized(C.this.size, C.GreedyBytes),
 )
 
 
-def _decode_quietly(data: bytes, h: sparser.Header) -> str:
+def _decode_quietly(data: bytes, type: bytes) -> str:
     try:
         return data.decode("utf-8")
     except UnicodeDecodeError:
-        LOG.warning("Failed to decode %s: %s", h, data[:512])
+        LOG.warning("Failed to decode %s: %s", type, data[:512])
         return ""
 
 
-def _parse_quietly(data: bytes, h: sparser.Header) -> bytes:
+def _parse_quietly(data: bytes, type: bytes) -> bytes:
     try:
-        parsed = MakeOrModel.parse(data)
+        parsed = _MakeOrModel.parse(data)
     except C.ConstructError:
-        LOG.warning("Failed to parse %s: %s", h, data[:512])
+        LOG.warning("Failed to parse %s: %s", type, data[:512])
         return b""
-    return parsed["data"]
 
+    if parsed is None:
+        return b""
 
-def extract_camera_make_and_model(fp: T.BinaryIO) -> T.Tuple[str, str]:
-    header_and_stream = sparser.parse_path(
-        fp,
-        [
-            b"moov",
-            b"udta",
-            [
-                # Insta360 Titan
-                b"\xa9mak",
-                b"\xa9mod",
-                # RICHO THETA V
-                b"@mod",
-                b"@mak",
-                # RICHO THETA V
-                b"manu",
-                b"modl",
-            ],
-        ],
-    )
+    return parsed["data"]
 
-    make: T.Optional[str] = None
-    model: T.Optional[str] = None
 
-    try:
-        for h, s in header_and_stream:
-            data = s.read(h.maxsize)
-            if h.type == b"\xa9mak":
-                make_data = _parse_quietly(data, h)
+def _extract_camera_make_and_model_from_utda_boxdata(
+    utda_boxdata: dict,
+) -> tuple[str, str]:
+    make: str = ""
+    model: str = ""
+
+    for box in utda_boxdata:
+        # Insta360 Titan
+        if box.type == b"\xa9mak":
+            if not make:
+                make_data = _parse_quietly(box.data, box.type)
                 make_data = make_data.rstrip(b"\x00")
-                make = _decode_quietly(make_data, h)
-            elif h.type == b"\xa9mod":
-                model_data = _parse_quietly(data, h)
+                make = _decode_quietly(make_data, box.type)
+
+        # Insta360 Titan
+        elif box.type == b"\xa9mod":
+            if not model:
+                model_data = _parse_quietly(box.data, box.type)
                 model_data = model_data.rstrip(b"\x00")
-                model = _decode_quietly(model_data, h)
-            elif h.type in [b"@mak", b"manu"]:
-                make = _decode_quietly(data, h)
-            elif h.type in [b"@mod", b"modl"]:
-                model = _decode_quietly(data, h)
-            # quit when both found
-            if make and model:
-                break
-    except sparser.ParsingError:
-        pass
+                model = _decode_quietly(model_data, box.type)
+
+        # RICHO THETA V
+        elif box.type in [b"@mak", b"manu"]:
+            if not make:
+                make = _decode_quietly(box.data, box.type)
+
+        # RICHO THETA V
+        elif box.type in [b"@mod", b"modl"]:
+            if not model:
+                model = _decode_quietly(box.data, box.type)
+
+        # quit when both found
+        if make and model:
+            break
 
     if make:
         make = make.strip()
+
     if model:
         model = model.strip()
-    return make or "", model or ""
+
+    return make, model