PCGao/MatchAnything

1# server.py2import warnings3from pathlib import Path4from typing import Union5 6import numpy as np7import ray8import torch9import yaml10from fastapi import FastAPI, File, UploadFile11from fastapi.responses import JSONResponse12from PIL import Image13from ray import serve14 15from . import ImagesInput, to_base64_nparray16from .core import ImageMatchingAPI17from ..hloc import DEVICE18from ..ui import get_version19 20warnings.simplefilter("ignore")21app = FastAPI()22if ray.is_initialized():23    ray.shutdown()24ray.init(25    dashboard_port=8265,26    ignore_reinit_error=True,27)28serve.start(29    http_options={"host": "0.0.0.0", "port": 8001},30)31 32num_gpus = 1 if torch.cuda.is_available() else 033 34 35@serve.deployment(36    num_replicas=4, ray_actor_options={"num_cpus": 2, "num_gpus": num_gpus}37)38@serve.ingress(app)39class ImageMatchingService:40    def __init__(self, conf: dict, device: str):41        self.conf = conf42        self.api = ImageMatchingAPI(conf=conf, device=device)43 44    @app.get("/")45    def root(self):46        return "Hello, world!"47 48    @app.get("/version")49    async def version(self):50        return {"version": get_version()}51 52    @app.post("/v1/match")53    async def match(54        self, image0: UploadFile = File(...), image1: UploadFile = File(...)55    ):56        """57        Handle the image matching request and return the processed result.58 59        Args:60            image0 (UploadFile): The first image file for matching.61            image1 (UploadFile): The second image file for matching.62 63        Returns:64            JSONResponse: A JSON response containing the filtered match results65                            or an error message in case of failure.66        """67        try:68            # Load the images from the uploaded files69            image0_array = self.load_image(image0)70            image1_array = self.load_image(image1)71 72            # Perform image matching using the API73            output = self.api(image0_array, image1_array)74 75            # Keys to skip in the output76            skip_keys = ["image0_orig", "image1_orig"]77 78            # Postprocess the output to filter unwanted data79            pred = self.postprocess(output, skip_keys)80 81            # Return the filtered prediction as a JSON response82            return JSONResponse(content=pred)83        except Exception as e:84            # Return an error message with status code 500 in case of exception85            return JSONResponse(content={"error": str(e)}, status_code=500)86 87    @app.post("/v1/extract")88    async def extract(self, input_info: ImagesInput):89        """90        Extract keypoints and descriptors from images.91 92        Args:93            input_info: An object containing the image data and options.94 95        Returns:96            A list of dictionaries containing the keypoints and descriptors.97        """98        try:99            preds = []100            for i, input_image in enumerate(input_info.data):101                # Load the image from the input data102                image_array = to_base64_nparray(input_image)103                # Extract keypoints and descriptors104                output = self.api.extract(105                    image_array,106                    max_keypoints=input_info.max_keypoints[i],107                    binarize=input_info.binarize,108                )109                # Do not return the original image and image_orig110                # skip_keys = ["image", "image_orig"]111                skip_keys = []112 113                # Postprocess the output114                pred = self.postprocess(output, skip_keys)115                preds.append(pred)116            # Return the list of extracted features117            return JSONResponse(content=preds)118        except Exception as e:119            # Return an error message if an exception occurs120            return JSONResponse(content={"error": str(e)}, status_code=500)121 122    def load_image(self, file_path: Union[str, UploadFile]) -> np.ndarray:123        """124        Reads an image from a file path or an UploadFile object.125 126        Args:127            file_path: A file path or an UploadFile object.128 129        Returns:130            A numpy array representing the image.131        """132        if isinstance(file_path, str):133            file_path = Path(file_path).resolve(strict=False)134        else:135            file_path = file_path.file136        with Image.open(file_path) as img:137            image_array = np.array(img)138        return image_array139 140    def postprocess(self, output: dict, skip_keys: list, binarize: bool = True) -> dict:141        pred = {}142        for key, value in output.items():143            if key in skip_keys:144                continue145            if isinstance(value, np.ndarray):146                pred[key] = value.tolist()147        return pred148 149    def run(self, host: str = "0.0.0.0", port: int = 8001):150        import uvicorn151 152        uvicorn.run(app, host=host, port=port)153 154 155def read_config(config_path: Path) -> dict:156    with open(config_path, "r") as f:157        conf = yaml.safe_load(f)158    return conf159 160 161# api server162conf = read_config(Path(__file__).parent / "config/api.yaml")163service = ImageMatchingService.bind(conf=conf["api"], device=DEVICE)164handle = serve.run(service, route_prefix="/")165 166# serve run api.server_ray:service167 168# build to generate config file169# serve build api.server_ray:service -o api/config/ray.yaml170# serve run api/config/ray.yaml171