File size: 3,332 Bytes
892fa81 | 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 | """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
|