api/server.py

# server.py
import base64
import io
import sys
import warnings
from pathlib import Path
from typing import Any, Dict, Optional, Union

import cv2
import matplotlib.pyplot as plt
import numpy as np
import torch
import uvicorn
from fastapi import FastAPI, File, UploadFile
from fastapi.exceptions import HTTPException
from fastapi.responses import JSONResponse
from PIL import Image

sys.path.append(str(Path(__file__).parents[1]))

from api.types import ImagesInput
from hloc import DEVICE, extract_features, logger, match_dense, match_features
from hloc.utils.viz import add_text, plot_keypoints
from ui import get_version
from ui.utils import filter_matches, get_feature_model, get_model
from ui.viz import display_matches, fig2im, plot_images

warnings.simplefilter("ignore")


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")


class ImageMatchingAPI(torch.nn.Module):
    default_conf = {
        "ransac": {
            "enable": True,
            "estimator": "poselib",
            "geometry": "homography",
            "method": "RANSAC",
            "reproj_threshold": 3,
            "confidence": 0.9999,
            "max_iter": 10000,
        },
    }

    def __init__(
        self,
        conf: dict = {},
        device: str = "cpu",
        detect_threshold: float = 0.015,
        max_keypoints: int = 1024,
        match_threshold: float = 0.2,
    ) -> None:
        """
        Initializes an instance of the ImageMatchingAPI class.

        Args:
            conf (dict): A dictionary containing the configuration parameters.
            device (str, optional): The device to use for computation. Defaults to "cpu".
            detect_threshold (float, optional): The threshold for detecting keypoints. Defaults to 0.015.
            max_keypoints (int, optional): The maximum number of keypoints to extract. Defaults to 1024.
            match_threshold (float, optional): The threshold for matching keypoints. Defaults to 0.2.

        Returns:
            None
        """
        super().__init__()
        self.device = device
        self.conf = {**self.default_conf, **conf}
        self._updata_config(detect_threshold, max_keypoints, match_threshold)
        self._init_models()
        if device == "cuda":
            memory_allocated = torch.cuda.memory_allocated(device)
            memory_reserved = torch.cuda.memory_reserved(device)
            logger.info(
                f"GPU memory allocated: {memory_allocated / 1024**2:.3f} MB"
            )
            logger.info(
                f"GPU memory reserved: {memory_reserved / 1024**2:.3f} MB"
            )
        self.pred = None

    def parse_match_config(self, conf):
        if conf["dense"]:
            return {
                **conf,
                "matcher": match_dense.confs.get(
                    conf["matcher"]["model"]["name"]
                ),
                "dense": True,
            }
        else:
            return {
                **conf,
                "feature": extract_features.confs.get(
                    conf["feature"]["model"]["name"]
                ),
                "matcher": match_features.confs.get(
                    conf["matcher"]["model"]["name"]
                ),
                "dense": False,
            }

    def _updata_config(
        self,
        detect_threshold: float = 0.015,
        max_keypoints: int = 1024,
        match_threshold: float = 0.2,
    ):
        self.dense = self.conf["dense"]
        if self.conf["dense"]:
            try:
                self.conf["matcher"]["model"][
                    "match_threshold"
                ] = match_threshold
            except TypeError as e:
                logger.error(e)
        else:
            self.conf["feature"]["model"]["max_keypoints"] = max_keypoints
            self.conf["feature"]["model"][
                "keypoint_threshold"
            ] = detect_threshold
            self.extract_conf = self.conf["feature"]

        self.match_conf = self.conf["matcher"]

    def _init_models(self):
        # initialize matcher
        self.matcher = get_model(self.match_conf)
        # initialize extractor
        if self.dense:
            self.extractor = None
        else:
            self.extractor = get_feature_model(self.conf["feature"])

    def _forward(self, img0, img1):
        if self.dense:
            pred = match_dense.match_images(
                self.matcher,
                img0,
                img1,
                self.match_conf["preprocessing"],
                device=self.device,
            )
            last_fixed = "{}".format(  # noqa: F841
                self.match_conf["model"]["name"]
            )
        else:
            pred0 = extract_features.extract(
                self.extractor, img0, self.extract_conf["preprocessing"]
            )
            pred1 = extract_features.extract(
                self.extractor, img1, self.extract_conf["preprocessing"]
            )
            pred = match_features.match_images(self.matcher, pred0, pred1)
        return pred

    @torch.inference_mode()
    def extract(self, img0: np.ndarray, **kwargs) -> Dict[str, np.ndarray]:
        """Extract features from a single image.

        Args:
            img0 (np.ndarray): image

        Returns:
            Dict[str, np.ndarray]: feature dict
        """

        # setting prams
        self.extractor.conf["max_keypoints"] = kwargs.get("max_keypoints", 512)
        self.extractor.conf["keypoint_threshold"] = kwargs.get(
            "keypoint_threshold", 0.0
        )

        pred = extract_features.extract(
            self.extractor, img0, self.extract_conf["preprocessing"]
        )
        pred = {
            k: v.cpu().detach()[0].numpy() if isinstance(v, torch.Tensor) else v
            for k, v in pred.items()
        }
        # back to origin scale
        s0 = pred["original_size"] / pred["size"]
        pred["keypoints_orig"] = (
            match_features.scale_keypoints(pred["keypoints"] + 0.5, s0) - 0.5
        )
        # TODO: rotate back

        binarize = kwargs.get("binarize", False)
        if binarize:
            assert "descriptors" in pred
            pred["descriptors"] = (pred["descriptors"] > 0).astype(np.uint8)
            pred["descriptors"] = pred["descriptors"].T  # N x DIM
        return pred

    @torch.inference_mode()
    def forward(
        self,
        img0: np.ndarray,
        img1: np.ndarray,
    ) -> Dict[str, np.ndarray]:
        """
        Forward pass of the image matching API.

        Args:
            img0: A 3D NumPy array of shape (H, W, C) representing the first image.
                  Values are in the range [0, 1] and are in RGB mode.
            img1: A 3D NumPy array of shape (H, W, C) representing the second image.
                  Values are in the range [0, 1] and are in RGB mode.

        Returns:
            A dictionary containing the following keys:
            - image0_orig: The original image 0.
            - image1_orig: The original image 1.
            - keypoints0_orig: The keypoints detected in image 0.
            - keypoints1_orig: The keypoints detected in image 1.
            - mkeypoints0_orig: The raw matches between image 0 and image 1.
            - mkeypoints1_orig: The raw matches between image 1 and image 0.
            - mmkeypoints0_orig: The RANSAC inliers in image 0.
            - mmkeypoints1_orig: The RANSAC inliers in image 1.
            - mconf: The confidence scores for the raw matches.
            - mmconf: The confidence scores for the RANSAC inliers.
        """
        # Take as input a pair of images (not a batch)
        assert isinstance(img0, np.ndarray)
        assert isinstance(img1, np.ndarray)
        self.pred = self._forward(img0, img1)
        if self.conf["ransac"]["enable"]:
            self.pred = self._geometry_check(self.pred)
        return self.pred

    def _geometry_check(
        self,
        pred: Dict[str, Any],
    ) -> Dict[str, Any]:
        """
        Filter matches using RANSAC. If keypoints are available, filter by keypoints.
        If lines are available, filter by lines. If both keypoints and lines are
        available, filter by keypoints.

        Args:
            pred (Dict[str, Any]): dict of matches, including original keypoints.
                                  See :func:`filter_matches` for the expected keys.

        Returns:
            Dict[str, Any]: filtered matches
        """
        pred = filter_matches(
            pred,
            ransac_method=self.conf["ransac"]["method"],
            ransac_reproj_threshold=self.conf["ransac"]["reproj_threshold"],
            ransac_confidence=self.conf["ransac"]["confidence"],
            ransac_max_iter=self.conf["ransac"]["max_iter"],
        )
        return pred

    def visualize(
        self,
        log_path: Optional[Path] = None,
    ) -> None:
        """
        Visualize the matches.

        Args:
            log_path (Path, optional): The directory to save the images. Defaults to None.

        Returns:
            None
        """
        if self.conf["dense"]:
            postfix = str(self.conf["matcher"]["model"]["name"])
        else:
            postfix = "{}_{}".format(
                str(self.conf["feature"]["model"]["name"]),
                str(self.conf["matcher"]["model"]["name"]),
            )
        titles = [
            "Image 0 - Keypoints",
            "Image 1 - Keypoints",
        ]
        pred: Dict[str, Any] = self.pred
        image0: np.ndarray = pred["image0_orig"]
        image1: np.ndarray = pred["image1_orig"]
        output_keypoints: np.ndarray = plot_images(
            [image0, image1], titles=titles, dpi=300
        )
        if (
            "keypoints0_orig" in pred.keys()
            and "keypoints1_orig" in pred.keys()
        ):
            plot_keypoints([pred["keypoints0_orig"], pred["keypoints1_orig"]])
            text: str = (
                f"# keypoints0: {len(pred['keypoints0_orig'])} \n"
                + f"# keypoints1: {len(pred['keypoints1_orig'])}"
            )
            add_text(0, text, fs=15)
        output_keypoints = fig2im(output_keypoints)
        # plot images with raw matches
        titles = [
            "Image 0 - Raw matched keypoints",
            "Image 1 - Raw matched keypoints",
        ]
        output_matches_raw, num_matches_raw = display_matches(
            pred, titles=titles, tag="KPTS_RAW"
        )
        # plot images with ransac matches
        titles = [
            "Image 0 - Ransac matched keypoints",
            "Image 1 - Ransac matched keypoints",
        ]
        output_matches_ransac, num_matches_ransac = display_matches(
            pred, titles=titles, tag="KPTS_RANSAC"
        )
        if log_path is not None:
            img_keypoints_path: Path = log_path / f"img_keypoints_{postfix}.png"
            img_matches_raw_path: Path = (
                log_path / f"img_matches_raw_{postfix}.png"
            )
            img_matches_ransac_path: Path = (
                log_path / f"img_matches_ransac_{postfix}.png"
            )
            cv2.imwrite(
                str(img_keypoints_path),
                output_keypoints[:, :, ::-1].copy(),  # RGB -> BGR
            )
            cv2.imwrite(
                str(img_matches_raw_path),
                output_matches_raw[:, :, ::-1].copy(),  # RGB -> BGR
            )
            cv2.imwrite(
                str(img_matches_ransac_path),
                output_matches_ransac[:, :, ::-1].copy(),  # RGB -> BGR
            )
            plt.close("all")


class ImageMatchingService:
    def __init__(self, conf: dict, device: str):
        self.conf = conf
        self.api = ImageMatchingAPI(conf=conf, device=device)
        self.app = FastAPI()
        self.register_routes()

    def register_routes(self):

        @self.app.get("/version")
        async def version():
            return {"version": get_version()}

        @self.app.post("/v1/match")
        async def match(
            image0: UploadFile = File(...), image1: UploadFile = File(...)
        ):
            """
            Handle the image matching request and return the processed result.

            Args:
                image0 (UploadFile): The first image file for matching.
                image1 (UploadFile): The second image file for matching.

            Returns:
                JSONResponse: A JSON response containing the filtered match results
                              or an error message in case of failure.
            """
            try:
                # Load the images from the uploaded files
                image0_array = self.load_image(image0)
                image1_array = self.load_image(image1)

                # Perform image matching using the API
                output = self.api(image0_array, image1_array)

                # Keys to skip in the output
                skip_keys = ["image0_orig", "image1_orig"]

                # Postprocess the output to filter unwanted data
                pred = self.postprocess(output, skip_keys)

                # Return the filtered prediction as a JSON response
                return JSONResponse(content=pred)
            except Exception as e:
                # Return an error message with status code 500 in case of exception
                return JSONResponse(content={"error": str(e)}, status_code=500)

        @self.app.post("/v1/extract")
        async def extract(input_info: ImagesInput):
            """
            Extract keypoints and descriptors from images.

            Args:
                input_info: An object containing the image data and options.

            Returns:
                A list of dictionaries containing the keypoints and descriptors.
            """
            try:
                preds = []
                for i, input_image in enumerate(input_info.data):
                    # Load the image from the input data
                    image_array = to_base64_nparray(input_image)
                    # Extract keypoints and descriptors
                    output = self.api.extract(
                        image_array,
                        max_keypoints=input_info.max_keypoints[i],
                        binarize=input_info.binarize,
                    )
                    # Do not return the original image and image_orig
                    # skip_keys = ["image", "image_orig"]
                    skip_keys = []

                    # Postprocess the output
                    pred = self.postprocess(output, skip_keys)
                    preds.append(pred)
                # Return the list of extracted features
                return JSONResponse(content=preds)
            except Exception as e:
                # Return an error message if an exception occurs
                return JSONResponse(content={"error": str(e)}, status_code=500)

    def load_image(self, file_path: Union[str, UploadFile]) -> np.ndarray:
        """
        Reads an image from a file path or an UploadFile object.

        Args:
            file_path: A file path or an UploadFile object.

        Returns:
            A numpy array representing the image.
        """
        if isinstance(file_path, str):
            file_path = Path(file_path).resolve(strict=False)
        else:
            file_path = file_path.file
        with Image.open(file_path) as img:
            image_array = np.array(img)
        return image_array

    def postprocess(
        self, output: dict, skip_keys: list, binarize: bool = True
    ) -> dict:
        pred = {}
        for key, value in output.items():
            if key in skip_keys:
                continue
            if isinstance(value, np.ndarray):
                pred[key] = value.tolist()
        return pred

    def run(self, host: str = "0.0.0.0", port: int = 8001):
        uvicorn.run(self.app, host=host, port=port)


if __name__ == "__main__":
    conf = {
        "feature": {
            "output": "feats-superpoint-n4096-rmax1600",
            "model": {
                "name": "superpoint",
                "nms_radius": 3,
                "max_keypoints": 4096,
                "keypoint_threshold": 0.005,
            },
            "preprocessing": {
                "grayscale": True,
                "force_resize": True,
                "resize_max": 1600,
                "width": 640,
                "height": 480,
                "dfactor": 8,
            },
        },
        "matcher": {
            "output": "matches-NN-mutual",
            "model": {
                "name": "nearest_neighbor",
                "do_mutual_check": True,
                "match_threshold": 0.2,
            },
        },
        "dense": False,
    }

    service = ImageMatchingService(conf=conf, device=DEVICE)
    service.run()