api/__init__.py

import sys
from typing import List
from pydantic import BaseModel
import base64
import io
import numpy as np
from fastapi.exceptions import HTTPException
from PIL import Image
from pathlib import Path

sys.path.append(str(Path(__file__).parents[1]))
from hloc import logger


class ImagesInput(BaseModel):
    data: List[str] = []
    max_keypoints: List[int] = []
    timestamps: List[str] = []
    grayscale: bool = False
    image_hw: List[List[int]] = [[], []]
    feature_type: int = 0
    rotates: List[float] = []
    scales: List[float] = []
    reference_points: List[List[float]] = []
    binarize: bool = False


def decode_base64_to_image(encoding):
    if encoding.startswith("data:image/"):
        encoding = encoding.split(";")[1].split(",")[1]
    try:
        image = Image.open(io.BytesIO(base64.b64decode(encoding)))
        return image
    except Exception as e:
        logger.warning(f"API cannot decode image: {e}")
        raise HTTPException(
            status_code=500, detail="Invalid encoded image"
        ) from e


def to_base64_nparray(encoding: str) -> np.ndarray:
    return np.array(decode_base64_to_image(encoding)).astype("uint8")