Upload api.py

api.py

@@ -0,0 +1,107 @@
+from flask import Flask, request, jsonify
+import cv2
+import numpy as np
+from pathlib import Path
+import tempfile
+import os
+import logging
+from hloc import logger as hloc_logger
+from common.api import ImageMatchingAPI
+from common.utils import (
+    DEVICE,
+    get_matcher_zoo,
+    load_config,
+    ROOT,
+)
+
+app = Flask(__name__)
+
+# Load configuration and matcher
+config = load_config(ROOT / "common/config.yaml")
+matcher_zoo_restored = get_matcher_zoo(config["matcher_zoo"])
+conf = matcher_zoo_restored.get("superpoint+superglue")
+api = ImageMatchingAPI(conf=conf, device=DEVICE)
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+def match_images(img_path1, img_path2):
+    logger.info(f"Matching images with {DEVICE}")
+    try:
+        # Read the images using OpenCV
+        image0 = cv2.imread(img_path1)
+        image1 = cv2.imread(img_path2)
+
+        if image0 is None or image1 is None:
+            raise ValueError("One or both images could not be read. Ensure the files are valid images.")
+
+        # Convert BGR to RGB
+        image0 = cv2.cvtColor(image0, cv2.COLOR_BGR2RGB)
+        image1 = cv2.cvtColor(image1, cv2.COLOR_BGR2RGB)
+
+        results = api(image0, image1)
+
+        return {
+            "num_keypoints0": len(results["keypoints0_orig"]),
+            "num_keypoints1": len(results["keypoints1_orig"]),
+            "num_matches": len(results["mkeypoints0_orig"]),
+            "num_inliers": len(results["mmkeypoints0_orig"]),
+            "keypoints0": results["keypoints0_orig"].tolist(),
+            "keypoints1": results["keypoints1_orig"].tolist(),
+            "matches0": results["mkeypoints0_orig"].tolist(),
+            "matches1": results["mmkeypoints0_orig"].tolist(),
+            "inliers0": results["mmkeypoints0_orig"].tolist(),
+            "inliers1": results["mmkeypoints1_orig"].tolist(),
+            "confidence": results["mconf"].tolist(),
+            "inlier_confidence": results["mmconf"].tolist(),
+        }
+    except Exception as e:
+        logger.error("Matching failed: %s", str(e))
+        return {"error": str(e)}
+    finally:
+        try:
+            os.remove(img_path1)
+            os.remove(img_path2)
+        except Exception as cleanup_error:
+            logger.error("Failed to clean up temporary files: %s", str(cleanup_error))
+
+@app.route('/match-images', methods=['POST'])
+def match_images_endpoint():
+    if 'img1' not in request.files or 'img2' not in request.files:
+        return jsonify({"error": "Please upload both images."}), 400
+
+    img1_file = request.files['img1']
+    img2_file = request.files['img2']
+
+    # Validate file types
+    valid_extensions = {'png', 'jpg', 'jpeg', 'bmp', 'tiff'}
+    if not (img1_file.filename.split('.')[-1].lower() in valid_extensions and
+            img2_file.filename.split('.')[-1].lower() in valid_extensions):
+        return jsonify({"error": "Invalid file type. Please upload images in one of the following formats: png, jpg, jpeg, bmp, tiff"}), 400
+
+    try:
+        # Save the uploaded files to a temporary directory with unique filenames
+        temp_dir = tempfile.mkdtemp()
+        img_path1 = os.path.join(temp_dir, next(tempfile._get_candidate_names()) + os.path.splitext(img1_file.filename)[1])
+        img_path2 = os.path.join(temp_dir, next(tempfile._get_candidate_names()) + os.path.splitext(img2_file.filename)[1])
+        img1_file.save(img_path1)
+        img2_file.save(img_path2)
+
+        # Run the matching task synchronously
+        result = match_images(img_path1, img_path2)
+    except Exception as e:
+        logger.error("Error processing images: %s", str(e))
+        return jsonify({"error": str(e)}), 500
+    finally:
+        # Clean up the temporary directory
+        if os.path.exists(temp_dir):
+            try:
+                os.rmdir(temp_dir)
+            except Exception as cleanup_error:
+                logger.error("Failed to clean up temporary directory: %s", str(cleanup_error))
+
+    return jsonify(result), 200
+
+if __name__ == '__main__':
+    app.run(debug=True)