PCGao/MatchAnything

1import argparse2import base643import os4import pickle5import time6from typing import Dict, List7 8import cv29import numpy as np10import requests11 12ENDPOINT = "http://127.0.0.1:8001"13if "REMOTE_URL_RAILWAY" in os.environ:14    ENDPOINT = os.environ["REMOTE_URL_RAILWAY"]15 16print(f"API ENDPOINT: {ENDPOINT}")17 18API_VERSION = f"{ENDPOINT}/version"19API_URL_MATCH = f"{ENDPOINT}/v1/match"20API_URL_EXTRACT = f"{ENDPOINT}/v1/extract"21 22 23def read_image(path: str) -> str:24    """25    Read an image from a file, encode it as a JPEG and then as a base64 string.26 27    Args:28        path (str): The path to the image to read.29 30    Returns:31        str: The base64 encoded image.32    """33    # Read the image from the file34    img = cv2.imread(path, cv2.IMREAD_GRAYSCALE)35 36    # Encode the image as a png, NO COMPRESSION!!!37    retval, buffer = cv2.imencode(".png", img)38 39    # Encode the JPEG as a base64 string40    b64img = base64.b64encode(buffer).decode("utf-8")41 42    return b64img43 44 45def do_api_requests(url=API_URL_EXTRACT, **kwargs):46    """47    Helper function to send an API request to the image matching service.48 49    Args:50        url (str): The URL of the API endpoint to use. Defaults to the51            feature extraction endpoint.52        **kwargs: Additional keyword arguments to pass to the API.53 54    Returns:55        List[Dict[str, np.ndarray]]: A list of dictionaries containing the56            extracted features. The keys are "keypoints", "descriptors", and57            "scores", and the values are ndarrays of shape (N, 2), (N, ?),58            and (N,), respectively.59    """60    # Set up the request body61    reqbody = {62        # List of image data base64 encoded63        "data": [],64        # List of maximum number of keypoints to extract from each image65        "max_keypoints": [100, 100],66        # List of timestamps for each image (not used?)67        "timestamps": ["0", "1"],68        # Whether to convert the images to grayscale69        "grayscale": 0,70        # List of image height and width71        "image_hw": [[640, 480], [320, 240]],72        # Type of feature to extract73        "feature_type": 0,74        # List of rotation angles for each image75        "rotates": [0.0, 0.0],76        # List of scale factors for each image77        "scales": [1.0, 1.0],78        # List of reference points for each image (not used)79        "reference_points": [[640, 480], [320, 240]],80        # Whether to binarize the descriptors81        "binarize": True,82    }83    # Update the request body with the additional keyword arguments84    reqbody.update(kwargs)85    try:86        # Send the request87        r = requests.post(url, json=reqbody)88        if r.status_code == 200:89            # Return the response90            return r.json()91        else:92            # Print an error message if the response code is not 20093            print(f"Error: Response code {r.status_code} - {r.text}")94    except Exception as e:95        # Print an error message if an exception occurs96        print(f"An error occurred: {e}")97 98 99def send_request_match(path0: str, path1: str) -> Dict[str, np.ndarray]:100    """101    Send a request to the API to generate a match between two images.102 103    Args:104        path0 (str): The path to the first image.105        path1 (str): The path to the second image.106 107    Returns:108        Dict[str, np.ndarray]: A dictionary containing the generated matches.109            The keys are "keypoints0", "keypoints1", "matches0", and "matches1",110            and the values are ndarrays of shape (N, 2), (N, 2), (N, 2), and111            (N, 2), respectively.112    """113    files = {"image0": open(path0, "rb"), "image1": open(path1, "rb")}114    try:115        # TODO: replace files with post json116        response = requests.post(API_URL_MATCH, files=files)117        pred = {}118        if response.status_code == 200:119            pred = response.json()120            for key in list(pred.keys()):121                pred[key] = np.array(pred[key])122        else:123            print(f"Error: Response code {response.status_code} - {response.text}")124    finally:125        files["image0"].close()126        files["image1"].close()127    return pred128 129 130def send_request_extract(131    input_images: str, viz: bool = False132) -> List[Dict[str, np.ndarray]]:133    """134    Send a request to the API to extract features from an image.135 136    Args:137        input_images (str): The path to the image.138 139    Returns:140        List[Dict[str, np.ndarray]]: A list of dictionaries containing the141            extracted features. The keys are "keypoints", "descriptors", and142            "scores", and the values are ndarrays of shape (N, 2), (N, 128),143            and (N,), respectively.144    """145    image_data = read_image(input_images)146    inputs = {147        "data": [image_data],148    }149    response = do_api_requests(150        url=API_URL_EXTRACT,151        **inputs,152    )153    # breakpoint()154    # print("Keypoints detected: {}".format(len(response[0]["keypoints"])))155 156    # draw matching, debug only157    if viz:158        from hloc.utils.viz import plot_keypoints159        from ui.viz import fig2im, plot_images160 161        kpts = np.array(response[0]["keypoints_orig"])162        if "image_orig" in response[0].keys():163            img_orig = np.array(["image_orig"])164 165            output_keypoints = plot_images([img_orig], titles="titles", dpi=300)166            plot_keypoints([kpts])167            output_keypoints = fig2im(output_keypoints)168            cv2.imwrite(169                "demo_match.jpg",170                output_keypoints[:, :, ::-1].copy(),  # RGB -> BGR171            )172    return response173 174 175def get_api_version():176    try:177        response = requests.get(API_VERSION).json()178        print("API VERSION: {}".format(response["version"]))179    except Exception as e:180        print(f"An error occurred: {e}")181 182 183if __name__ == "__main__":184    from pathlib import Path185 186    parser = argparse.ArgumentParser(187        description="Send text to stable audio server and receive generated audio."188    )189    parser.add_argument(190        "--image0",191        required=False,192        help="Path for the file's melody",193        default=str(194            Path(__file__).parents[1]195            / "datasets/sacre_coeur/mapping_rot/02928139_3448003521_rot45.jpg"196        ),197    )198    parser.add_argument(199        "--image1",200        required=False,201        help="Path for the file's melody",202        default=str(203            Path(__file__).parents[1]204            / "datasets/sacre_coeur/mapping_rot/02928139_3448003521_rot90.jpg"205        ),206    )207    args = parser.parse_args()208 209    # get api version210    get_api_version()211 212    # request match213    # for i in range(10):214    #     t1 = time.time()215    #     preds = send_request_match(args.image0, args.image1)216    #     t2 = time.time()217    #     print(218    #         "Time cost1: {} seconds, matched: {}".format(219    #             (t2 - t1), len(preds["mmkeypoints0_orig"])220    #         )221    #     )222 223    # request extract224    for i in range(1000):225        t1 = time.time()226        preds = send_request_extract(args.image0)227        t2 = time.time()228        print(f"Time cost2: {(t2 - t1)} seconds")229 230    # dump preds231    with open("preds.pkl", "wb") as f:232        pickle.dump(preds, f)233