Spaces:
No application file
No application file
File size: 4,157 Bytes
d188b4c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 |
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)
|