cores/vision/decode.py

"""Decode / encode — bytes ↔ numpy ↔ base64 ↔ URL."""

from __future__ import annotations

import base64
from typing import Optional

import cv2
import numpy as np
import requests


# --------------------------------------------------------------------------- #
# Bytes ↔ numpy
# --------------------------------------------------------------------------- #
def bytes_to_numpy(image_bytes: bytes) -> np.ndarray:
    """Decode raw image bytes into an OpenCV BGR numpy array."""
    nparr = np.frombuffer(image_bytes, np.uint8)
    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
    if img is None:
        raise ValueError("Could not decode image bytes. Unsupported format or corrupted data.")
    return img


def numpy_to_bytes(img: np.ndarray, fmt: str = ".jpg", quality: int = 90) -> bytes:
    """Encode a BGR numpy array to raw bytes."""
    params = [cv2.IMWRITE_JPEG_QUALITY, quality] if fmt.lower() in (".jpg", ".jpeg") else []
    ok, buffer = cv2.imencode(fmt, img, params)
    if not ok:
        raise ValueError("Could not encode image.")
    return buffer.tobytes()


# --------------------------------------------------------------------------- #
# Base64
# --------------------------------------------------------------------------- #
def base64_to_numpy(b64_string: str) -> np.ndarray:
    """Decode a base64-encoded image string into a BGR numpy array."""
    if "," in b64_string:
        b64_string = b64_string.split(",", 1)[1]
    raw = base64.b64decode(b64_string)
    return bytes_to_numpy(raw)


def numpy_to_base64(img: np.ndarray, fmt: str = ".jpg", quality: int = 85) -> str:
    """Encode a BGR numpy array as a base64 string."""
    return base64.b64encode(numpy_to_bytes(img, fmt, quality)).decode("utf-8")


# --------------------------------------------------------------------------- #
# URL
# --------------------------------------------------------------------------- #
_DEFAULT_HEADERS = {"User-Agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36"}


def url_to_bytes(url: str, timeout: int = 15) -> bytes:
    """Download a URL and return raw bytes."""
    resp = requests.get(url, headers=_DEFAULT_HEADERS, timeout=timeout, stream=True)
    resp.raise_for_status()
    return resp.content


def url_to_numpy(url: str, timeout: int = 15) -> np.ndarray:
    """Download an image from a URL and return it as a BGR numpy array."""
    return bytes_to_numpy(url_to_bytes(url, timeout))


# --------------------------------------------------------------------------- #
# Format sniffing (magic bytes)
# --------------------------------------------------------------------------- #
_IMAGE_SIGNATURES = {
    b"\xff\xd8\xff": "jpeg",
    b"\x89PNG\r\n\x1a\n": "png",
    b"GIF87a": "gif",
    b"GIF89a": "gif",
    b"BM": "bmp",
    b"II*\x00": "tiff",
    b"MM\x00*": "tiff",
}


def sniff_format(data: bytes) -> Optional[str]:
    """Identify image format from magic bytes.  Returns None if unknown."""
    if not data or len(data) < 12:
        return None
    for sig, fmt in _IMAGE_SIGNATURES.items():
        if data.startswith(sig):
            if sig == b"RIFF" and data[8:12] != b"WEBP":
                continue
            return fmt
    # RIFF/WEBP (4-byte prefix overlap with RIFF)
    if data[:4] == b"RIFF" and data[8:12] == b"WEBP":
        return "webp"
    return None