meta-sam-parser 0.0.2__py3-none-any.whl

meta_sam_parser/_segmentation.py

@@ -0,0 +1,533 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved.
+
+from __future__ import annotations
+
+import math
+import re
+from dataclasses import dataclass
+from typing import cast
+
+from ._errors import InvalidSegmentationMaskError
+from ._mask_codec import decode_mask_to_raster
+from ._types import (
+    FrameReference,
+    ImageSegmentationResult,
+    ImageSegmentationSnapshot,
+    ParserFinish,
+    ResponseFormat,
+    ResponseStreamOutcome,
+    SegmentationBoxRecord,
+    SegmentationDiagnostic,
+    SegmentationMask,
+    SegmentationMaskBounds,
+    SegmentationMaskIdentity,
+    SegmentationMaskRecord,
+    SegmentationMedia,
+    SegmentationPointRecord,
+    SegmentationRecord,
+    SegmentationResult,
+    SegmentationSnapshot,
+    SegmentationTextRecord,
+    VideoSegmentationResult,
+    VideoSegmentationSnapshot,
+)
+
+_MAXIMUM_SAFE_INTEGER = (1 << 53) - 1
+_JS_WHITESPACE = (
+    "\t\n\v\f\r \u00a0\u1680\u2000\u2001\u2002\u2003\u2004\u2005\u2006"
+    "\u2007\u2008\u2009\u200a\u2028\u2029\u202f\u205f\u3000\ufeff"
+)
+_JS_WHITESPACE_CLASS = (
+    r"\t\n\v\f\r \u00a0\u1680\u2000-\u200a\u2028\u2029\u202f\u205f\u3000\ufeff"
+)
+_JS_WHITESPACE_PATTERN = rf"[{_JS_WHITESPACE_CLASS}]"
+_JS_NON_WHITESPACE_PATTERN = rf"[^{_JS_WHITESPACE_CLASS}]"
+_JS_DOT_PATTERN = r"[^\n\r\u2028\u2029]"
+_RECORD_PATTERN = re.compile(
+    rf"^(?:\[frame=([0-9]+)\]{_JS_WHITESPACE_PATTERN}+)?"
+    rf"object=([A-Za-z0-9._-]+){_JS_WHITESPACE_PATTERN}+"
+    rf"(point|box|mask)=({_JS_DOT_PATTERN}+)$"
+)
+_COMPACT_API_HEADER_PATTERN = re.compile(rf"^<([0-9]+)f>({_JS_DOT_PATTERN}*)$")
+_COMPACT_API_OBJECT_PATTERN = re.compile(
+    r"^(?:,)?([0-9]+)"
+    r"<\|box;x1=(-?[0-9]+);y1=(-?[0-9]+);x2=(-?[0-9]+);"
+    r"y2=(-?[0-9]+);w=([0-9]+);h=([0-9]+)\|>"
+    r"<\|mask;x=0;y=0;data=([0-9]+),([0-9]+),([!~][^|]+)\|>"
+)
+_POINT_PATTERN = re.compile(r"^\(([^,]+),([^,]+)\)$")
+_BOX_PATTERN = re.compile(r"^\(([^,]+),([^,]+),([^,]+),([^,]+)\)$")
+_MASK_PATTERN = re.compile(
+    rf"^([^;{_JS_WHITESPACE_CLASS}]+);"
+    rf"size=([0-9]+)x([0-9]+);data=({_JS_NON_WHITESPACE_PATTERN}+)$"
+)
+_DECIMAL_NUMBER_PATTERN = re.compile(
+    r"^[+-]?(?:(?:[0-9]+(?:\.[0-9]*)?)|(?:\.[0-9]+))"
+    r"(?:[eE][+-]?[0-9]+)?$"
+)
+_HEXADECIMAL_NUMBER_PATTERN = re.compile(r"^0[xX][0-9A-Fa-f]+$")
+_BINARY_NUMBER_PATTERN = re.compile(r"^0[bB][01]+$")
+_OCTAL_NUMBER_PATTERN = re.compile(r"^0[oO][0-7]+$")
+
+
+def _trim(value: str) -> str:
+    return value.strip(_JS_WHITESPACE)
+
+
+def _javascript_number(value: str) -> float | None:
+    normalized = _trim(value)
+    if not normalized:
+        return None
+    try:
+        if _HEXADECIMAL_NUMBER_PATTERN.fullmatch(normalized):
+            parsed = float(int(normalized[2:], 16))
+        elif _BINARY_NUMBER_PATTERN.fullmatch(normalized):
+            parsed = float(int(normalized[2:], 2))
+        elif _OCTAL_NUMBER_PATTERN.fullmatch(normalized):
+            parsed = float(int(normalized[2:], 8))
+        elif _DECIMAL_NUMBER_PATTERN.fullmatch(normalized):
+            parsed = float(normalized)
+        else:
+            return None
+    except (OverflowError, ValueError):
+        return None
+    return parsed if math.isfinite(parsed) else None
+
+
+def _integer_text(value: str) -> int | None:
+    parsed = _javascript_number(value)
+    if parsed is None or not parsed.is_integer() or abs(parsed) > _MAXIMUM_SAFE_INTEGER:
+        return None
+    return int(parsed)
+
+
+class _SegmentationParser:
+    __slots__ = (
+        "_buffer_parts",
+        "_diagnostics",
+        "_line",
+        "_media",
+        "_raw_output_parts",
+        "_records",
+        "_revision",
+        "_revisions",
+    )
+
+    def __init__(self, media: SegmentationMedia) -> None:
+        self._media = media
+        self._records: list[SegmentationRecord] = []
+        self._diagnostics: list[SegmentationDiagnostic] = []
+        self._revisions: dict[SegmentationMaskIdentity, int] = {}
+        self._raw_output_parts: list[str] = []
+        self._buffer_parts: list[str] = []
+        self._line = 0
+        self._revision = 0
+
+    def push(
+        self, chunk: str, *, emit: bool = True
+    ) -> tuple[SegmentationSnapshot, ...]:
+        if not isinstance(chunk, str):
+            raise TypeError("Parser chunks must be strings.")
+        self._raw_output_parts.append(chunk)
+
+        changed = False
+        segments = chunk.split("\n")
+        for segment in segments[:-1]:
+            self._append_buffer(segment)
+            raw = "".join(self._buffer_parts)
+            if raw.endswith("\r"):
+                raw = raw[:-1]
+            self._clear_buffer()
+            prior = (len(self._records), len(self._diagnostics))
+            self._accept_line(raw)
+            changed = changed or prior != (
+                len(self._records),
+                len(self._diagnostics),
+            )
+        self._append_buffer(segments[-1])
+
+        if changed:
+            self._revision += 1
+        return (self._snapshot(),) if changed and emit is not False else ()
+
+    def finish(
+        self, outcome: ResponseStreamOutcome
+    ) -> ParserFinish[SegmentationSnapshot, SegmentationResult]:
+        events: tuple[SegmentationSnapshot, ...] = ()
+        if self._buffer_parts:
+            raw = "".join(self._buffer_parts)
+            if raw.endswith("\r"):
+                raw = raw[:-1]
+            self._clear_buffer()
+            prior = (len(self._records), len(self._diagnostics))
+            self._accept_line(raw)
+            if prior != (len(self._records), len(self._diagnostics)):
+                self._revision += 1
+                events = (self._snapshot(),)
+        return ParserFinish(events=events, result=self._result(outcome))
+
+    def _append_buffer(self, value: str) -> None:
+        if value:
+            self._buffer_parts.append(value)
+
+    def _clear_buffer(self) -> None:
+        self._buffer_parts.clear()
+
+    def _accept_line(self, raw: str) -> bool:
+        self._line += 1
+        line = _trim(raw)
+        if not line:
+            return False
+
+        if line.startswith("<"):
+            accepted = self._accept_compact_api_line(line, raw)
+            if accepted is not None:
+                return accepted
+
+        match = _RECORD_PATTERN.fullmatch(line)
+        if match is None:
+            if line.startswith("[frame=") or line.startswith("object="):
+                self._diagnose("malformed_record", "Malformed structured record.", raw)
+            else:
+                self._add_record(
+                    SegmentationTextRecord(order=len(self._records), text=raw)
+                )
+            return True
+
+        frame_index = None if match.group(1) is None else _integer_text(match.group(1))
+        if match.group(1) is not None and (frame_index is None or frame_index < 0):
+            self._diagnose(
+                "invalid_frame",
+                "Frame references must be non-negative safe integers.",
+                raw,
+            )
+            return True
+        if self._media == "image" and frame_index is not None:
+            self._diagnose(
+                "unexpected_frame",
+                "Image segmentation records cannot include a frame reference.",
+                raw,
+            )
+            return True
+
+        object_id = match.group(2)
+        kind = match.group(3)
+        value = match.group(4)
+        frame = None if frame_index is None else FrameReference(frame_index)
+        if kind == "point":
+            self._accept_point(object_id, frame, value, raw)
+        elif kind == "box":
+            self._accept_box(object_id, frame, value, raw)
+        else:
+            self._accept_mask(object_id, frame, value, raw)
+        return True
+
+    def _accept_compact_api_line(self, line: str, raw: str) -> bool | None:
+        header = _COMPACT_API_HEADER_PATTERN.fullmatch(line)
+        if header is None:
+            return None
+        frame_index = _integer_text(header.group(1))
+        if frame_index is None or frame_index < 0:
+            self._diagnose(
+                "invalid_frame",
+                "Frame references must be non-negative safe integers.",
+                raw,
+            )
+            return True
+        if self._media == "image" and frame_index != 0:
+            self._diagnose(
+                "unexpected_frame",
+                "Image segmentation records require frame zero.",
+                raw,
+            )
+            return True
+
+        frame = FrameReference(frame_index) if self._media == "video" else None
+        remainder = header.group(2)
+        if not remainder:
+            self._diagnose(
+                "malformed_record",
+                "Malformed compact SAM API object record.",
+                raw,
+            )
+            return True
+        accepted = False
+        while remainder:
+            match = _COMPACT_API_OBJECT_PATTERN.match(remainder)
+            if match is None:
+                self._diagnose(
+                    "malformed_record",
+                    "Malformed compact SAM API object record.",
+                    raw,
+                )
+                return True
+            object_id = match.group(1)
+            coordinate_values = tuple(
+                _integer_text(value) for value in match.group(2, 3, 4, 5, 6, 7)
+            )
+            if any(value is None for value in coordinate_values):
+                self._diagnose(
+                    "invalid_box", "Canonical SAM box coordinates are invalid.", raw
+                )
+                return True
+            (
+                left,
+                top,
+                inclusive_right,
+                inclusive_bottom,
+                source_width,
+                source_height,
+            ) = cast(tuple[int, int, int, int, int, int], coordinate_values)
+            if (
+                source_width <= 0
+                or source_height <= 0
+                or left < 0
+                or top < 0
+                or inclusive_right < left
+                or inclusive_bottom < top
+                or inclusive_right >= source_width
+                or inclusive_bottom >= source_height
+            ):
+                self._diagnose(
+                    "invalid_box", "Canonical SAM box coordinates are invalid.", raw
+                )
+                return True
+            bounds = SegmentationMaskBounds(
+                left=left,
+                top=top,
+                right=inclusive_right + 1,
+                bottom=inclusive_bottom + 1,
+            )
+            self._add_record(
+                SegmentationBoxRecord(
+                    order=len(self._records),
+                    object_id=object_id,
+                    frame=frame,
+                    left=left,
+                    top=top,
+                    right=inclusive_right + 1,
+                    bottom=inclusive_bottom + 1,
+                )
+            )
+            mask_height = match.group(8)
+            mask_width = match.group(9)
+            payload = match.group(10)
+            encoding = "lossless" if payload.startswith("~") else "one_bit"
+            self._accept_mask(
+                object_id,
+                frame,
+                f"{encoding};size={mask_width}x{mask_height};data={payload}",
+                raw,
+                bounds,
+            )
+            accepted = True
+            remainder = remainder[match.end() :]
+        return accepted
+
+    def _accept_point(
+        self,
+        object_id: str,
+        frame: FrameReference | None,
+        value: str,
+        raw: str,
+    ) -> None:
+        match = _POINT_PATTERN.fullmatch(_trim(value))
+        x = None if match is None else _javascript_number(match.group(1))
+        y = None if match is None else _javascript_number(match.group(2))
+        if x is None or y is None:
+            self._diagnose(
+                "malformed_point", "Point coordinates must be finite numbers.", raw
+            )
+            return
+        self._add_record(
+            SegmentationPointRecord(
+                order=len(self._records),
+                object_id=object_id,
+                frame=frame,
+                x=x,
+                y=y,
+            )
+        )
+
+    def _accept_box(
+        self,
+        object_id: str,
+        frame: FrameReference | None,
+        value: str,
+        raw: str,
+    ) -> None:
+        match = _BOX_PATTERN.fullmatch(_trim(value))
+        coordinates = (
+            None
+            if match is None
+            else tuple(_javascript_number(item) for item in match.group(1, 2, 3, 4))
+        )
+        if coordinates is None or any(item is None for item in coordinates):
+            self._diagnose(
+                "malformed_box", "Box coordinates must be finite numbers.", raw
+            )
+            return
+        left, top, right, bottom = cast(tuple[float, float, float, float], coordinates)
+        if right < left or bottom < top:
+            self._diagnose("invalid_box", "Box coordinates must be ordered.", raw)
+            return
+        self._add_record(
+            SegmentationBoxRecord(
+                order=len(self._records),
+                object_id=object_id,
+                frame=frame,
+                left=left,
+                top=top,
+                right=right,
+                bottom=bottom,
+            )
+        )
+
+    def _accept_mask(
+        self,
+        object_id: str,
+        frame: FrameReference | None,
+        value: str,
+        raw: str,
+        bounds: SegmentationMaskBounds | None = None,
+    ) -> None:
+        match = _MASK_PATTERN.fullmatch(_trim(value))
+        if match is None:
+            self._diagnose(
+                "malformed_mask",
+                "Mask records require encoding, size, and one complete payload.",
+                raw,
+            )
+            return
+        encoding = match.group(1)
+        width_number = _javascript_number(match.group(2))
+        height_number = _javascript_number(match.group(3))
+        width = int(width_number) if width_number is not None else 0
+        height = int(height_number) if height_number is not None else 0
+        payload = match.group(4)
+        if encoding not in ("one_bit", "lossless"):
+            self._diagnose(
+                "unsupported_mask_encoding",
+                f"Unsupported complete mask encoding: {encoding}.",
+                raw,
+            )
+            return
+        area = float(width) * float(height)
+        if (
+            width <= 0
+            or height <= 0
+            or not math.isfinite(area)
+            or not area.is_integer()
+            or abs(area) > _MAXIMUM_SAFE_INTEGER
+        ):
+            self._diagnose(
+                "invalid_mask_size", "Mask dimensions must be positive.", raw
+            )
+            return
+        mask = SegmentationMask(
+            encoding=encoding, payload=payload, width=width, height=height
+        )
+        try:
+            decode_mask_to_raster(mask)
+        except InvalidSegmentationMaskError as error:
+            self._diagnose("invalid_mask_payload", str(error), raw)
+            return
+
+        identity = SegmentationMaskIdentity(
+            media=self._media,
+            frame_index=None if frame is None else frame.frame_index,
+            object_id=object_id,
+        )
+        revision = self._revisions.get(identity, 0) + 1
+        self._revisions[identity] = revision
+        self._add_record(
+            SegmentationMaskRecord(
+                order=len(self._records),
+                object_id=object_id,
+                frame=frame,
+                identity=identity,
+                revision=revision,
+                mask=mask,
+                bounds=bounds,
+            )
+        )
+
+    def _add_record(self, record: SegmentationRecord) -> None:
+        self._records.append(record)
+
+    def _diagnose(self, code: str, message: str, raw: str) -> None:
+        self._diagnostics.append(
+            SegmentationDiagnostic(
+                severity="error",
+                code=code,
+                message=message,
+                line=self._line,
+                raw=raw,
+            )
+        )
+
+    def _snapshot(self) -> SegmentationSnapshot:
+        records = tuple(self._records)
+        diagnostics = tuple(self._diagnostics)
+        if self._media == "image":
+            return ImageSegmentationSnapshot(
+                revision=self._revision,
+                records=records,
+                diagnostics=diagnostics,
+                raw_output="".join(self._raw_output_parts),
+            )
+        return VideoSegmentationSnapshot(
+            revision=self._revision,
+            records=records,
+            diagnostics=diagnostics,
+            raw_output="".join(self._raw_output_parts),
+        )
+
+    def _result(self, outcome: ResponseStreamOutcome) -> SegmentationResult:
+        records = tuple(self._records)
+        diagnostics = tuple(self._diagnostics)
+        if self._media == "image":
+            return ImageSegmentationResult(
+                revision=self._revision,
+                records=records,
+                diagnostics=diagnostics,
+                raw_output="".join(self._raw_output_parts),
+                outcome=outcome,
+            )
+        return VideoSegmentationResult(
+            revision=self._revision,
+            records=records,
+            diagnostics=diagnostics,
+            raw_output="".join(self._raw_output_parts),
+            outcome=outcome,
+        )
+
+
+@dataclass(frozen=True, slots=True)
+class _SegmentationFormat:
+    _media: SegmentationMedia
+
+    def create_parser(self) -> _SegmentationParser:
+        return _SegmentationParser(self._media)
+
+
+def image_segmentation_format() -> ResponseFormat[
+    ImageSegmentationSnapshot, ImageSegmentationResult
+]:
+    """Create an image format whose parsers have isolated incremental state."""
+
+    return cast(
+        ResponseFormat[ImageSegmentationSnapshot, ImageSegmentationResult],
+        _SegmentationFormat("image"),
+    )
+
+
+def video_segmentation_format() -> ResponseFormat[
+    VideoSegmentationSnapshot, VideoSegmentationResult
+]:
+    """Create a video format whose parsers have isolated incremental state."""
+
+    return cast(
+        ResponseFormat[VideoSegmentationSnapshot, VideoSegmentationResult],
+        _SegmentationFormat("video"),
+    )