228 lines
8.6 KiB
Python
228 lines
8.6 KiB
Python
"""
|
|
Vase Matcher -- Flask backend.
|
|
|
|
Upload a photo of a vase; it's matched against the fixed template set using
|
|
four methods (SIFT, ORB, SuperPoint+LightGlue, LoFTR), each scored
|
|
independently and timed. Template-side data is precomputed once at startup
|
|
so a request only ever processes the single uploaded image.
|
|
"""
|
|
|
|
import logging
|
|
import os
|
|
from logging.handlers import RotatingFileHandler
|
|
|
|
import config # noqa: F401 (must be imported first -- caps thread envs)
|
|
|
|
|
|
def _configure_logging():
|
|
"""All app + pipeline.* + werkzeug loggers propagate to the root logger,
|
|
so attaching handlers here is enough to capture everything in one file --
|
|
useful for correlating the last thing that happened before a crash (e.g.
|
|
an OOM kill, which the process itself never gets to log) against
|
|
`dmesg`/`journalctl -k` timestamps."""
|
|
os.makedirs(config.LOG_DIR, exist_ok=True)
|
|
fmt = logging.Formatter("%(asctime)s %(levelname)-7s [%(name)s] %(message)s")
|
|
|
|
root = logging.getLogger()
|
|
if root.handlers:
|
|
return # already configured (e.g. module imported twice)
|
|
root.setLevel(logging.INFO)
|
|
|
|
file_handler = RotatingFileHandler(config.LOG_FILE, maxBytes=5 * 1024 * 1024,
|
|
backupCount=5)
|
|
file_handler.setFormatter(fmt)
|
|
root.addHandler(file_handler)
|
|
|
|
console_handler = logging.StreamHandler()
|
|
console_handler.setFormatter(fmt)
|
|
root.addHandler(console_handler)
|
|
|
|
|
|
_configure_logging()
|
|
|
|
import requests
|
|
from flask import Flask, jsonify, render_template, request, send_from_directory
|
|
from werkzeug.utils import secure_filename
|
|
|
|
from pipeline import engine, utils, verify
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
app = Flask(__name__)
|
|
app.config["MAX_CONTENT_LENGTH"] = config.MAX_CONTENT_LENGTH_BYTES
|
|
|
|
|
|
@app.route("/")
|
|
def index():
|
|
templates = []
|
|
for fname in utils.list_template_files():
|
|
name = utils.template_name(fname)
|
|
templates.append({"name": name, "filename": fname})
|
|
return render_template("index.html", templates=templates,
|
|
methods=list(config.METHOD_LABELS.items()))
|
|
|
|
|
|
@app.route("/template_image/<path:filename>")
|
|
def template_image(filename):
|
|
safe = secure_filename(filename)
|
|
return send_from_directory(config.TEMPLATE_IMAGES_DIR, safe)
|
|
|
|
|
|
@app.route("/uploads/<request_id>/<path:filename>")
|
|
def uploaded_file(request_id, filename):
|
|
safe_id = secure_filename(request_id)
|
|
safe_name = secure_filename(filename)
|
|
request_dir = os.path.join(config.UPLOADS_DIR, safe_id)
|
|
return send_from_directory(request_dir, safe_name)
|
|
|
|
|
|
@app.errorhandler(413)
|
|
def too_large(_e):
|
|
logger.warning("Upload rejected: exceeds %d byte limit", config.MAX_CONTENT_LENGTH_BYTES)
|
|
return jsonify({"error": "Image too large (15 MB limit)."}), 413
|
|
|
|
|
|
@app.route("/api/match", methods=["POST"])
|
|
def api_match():
|
|
file = request.files.get("image")
|
|
if file is None or file.filename == "":
|
|
logger.warning("Match request rejected: no image uploaded")
|
|
return jsonify({"error": "No image uploaded."}), 400
|
|
|
|
ext = os.path.splitext(file.filename)[1].lower()
|
|
if ext not in config.VALID_EXTS:
|
|
logger.warning("Match request rejected: unsupported file type %r", ext)
|
|
return jsonify({"error": f"Unsupported file type: {ext or 'unknown'}"}), 400
|
|
|
|
image_bytes = file.read()
|
|
if not image_bytes:
|
|
logger.warning("Match request rejected: empty file")
|
|
return jsonify({"error": "Uploaded file is empty."}), 400
|
|
|
|
filename = file.filename
|
|
try:
|
|
image_bytes, was_compressed = utils.compress_image_bytes(
|
|
image_bytes, config.COMPRESS_ABOVE_BYTES, config.MAX_IMAGE_DIM
|
|
)
|
|
except Exception:
|
|
logger.exception("Compression failed for %r, using original bytes", filename)
|
|
was_compressed = False
|
|
|
|
if was_compressed:
|
|
# Re-encoded as JPEG regardless of the original format, so the
|
|
# filename extension needs to match -- otherwise the saved
|
|
# "original.*" file gets served with the wrong Content-Type later.
|
|
filename = os.path.splitext(filename)[0] + ".jpg"
|
|
|
|
try:
|
|
result = engine.process_upload(image_bytes, filename)
|
|
except Exception as e:
|
|
logger.exception("Match pipeline failed")
|
|
return jsonify({"error": f"Processing failed: {e}"}), 500
|
|
|
|
return jsonify(result)
|
|
|
|
|
|
@app.route("/api/count_flowers", methods=["POST"])
|
|
def api_count_flowers():
|
|
"""Opt-in, heavyweight: runs SAM instance segmentation over an
|
|
already-processed upload to count individual flowers and cluster them
|
|
by color as a rough proxy for distinct "kinds". Kept as its own request
|
|
(rather than folded into /api/match) since it unloads the matching
|
|
pipeline's GPU-resident models first -- it's meant to be triggered
|
|
explicitly, not on every upload."""
|
|
data = request.get_json(silent=True) or {}
|
|
request_id = data.get("request_id")
|
|
template = data.get("template")
|
|
if not request_id:
|
|
return jsonify({"error": "Missing request_id."}), 400
|
|
|
|
safe_id = secure_filename(request_id)
|
|
try:
|
|
result = engine.count_flowers(safe_id, template=template)
|
|
except FileNotFoundError:
|
|
return jsonify({"error": "Upload not found (it may have expired)."}), 404
|
|
except Exception as e:
|
|
logger.exception("[%s] flower count request failed", safe_id)
|
|
return jsonify({"error": f"Flower counting failed: {e}"}), 500
|
|
|
|
result["request_id"] = safe_id
|
|
return jsonify(result)
|
|
|
|
|
|
@app.route("/api/flower_summary", methods=["POST"])
|
|
def api_flower_summary():
|
|
"""Lightweight, auto-triggered (not opt-in) companion to the weighted
|
|
verdict: flower count comparison + CLIP-only flower similarity. Fired
|
|
automatically right after a confident match, shown beside the weighted
|
|
card while /api/count_flowers stays behind its own button for the
|
|
heavier SAM+YOLO-World+DINOv2/vase comparison."""
|
|
data = request.get_json(silent=True) or {}
|
|
request_id = data.get("request_id")
|
|
template = data.get("template")
|
|
if not request_id or not template:
|
|
return jsonify({"error": "Missing request_id or template."}), 400
|
|
|
|
safe_id = secure_filename(request_id)
|
|
try:
|
|
result = engine.flower_summary(safe_id, template)
|
|
except FileNotFoundError:
|
|
return jsonify({"error": "Upload not found (it may have expired)."}), 404
|
|
except ValueError as e:
|
|
return jsonify({"error": str(e)}), 400
|
|
except Exception as e:
|
|
logger.exception("[%s] flower summary request failed", safe_id)
|
|
return jsonify({"error": f"Flower summary failed: {e}"}), 500
|
|
|
|
return jsonify(result)
|
|
|
|
|
|
@app.route("/api/verify", methods=["POST"])
|
|
def api_verify():
|
|
"""Optional follow-up step: cross-checks the winning template against
|
|
the user's original upload via an external vision-LLM endpoint. Kept
|
|
separate from /api/match so a slow or unreachable verification service
|
|
never affects the core (fast, local) matching result."""
|
|
data = request.get_json(silent=True) or {}
|
|
request_id = data.get("request_id")
|
|
template = data.get("template")
|
|
|
|
if not request_id or not template:
|
|
return jsonify({"error": "Missing request_id or template."}), 400
|
|
|
|
safe_id = secure_filename(request_id)
|
|
request_dir = os.path.join(config.UPLOADS_DIR, safe_id)
|
|
if not os.path.isdir(request_dir):
|
|
return jsonify({"error": "Upload not found (it may have expired)."}), 404
|
|
|
|
originals = [f for f in os.listdir(request_dir) if f.startswith("original.")]
|
|
if not originals:
|
|
return jsonify({"error": "Original upload not found."}), 404
|
|
upload_path = os.path.join(request_dir, originals[0])
|
|
|
|
templates = engine.templates_meta()
|
|
if template not in templates:
|
|
return jsonify({"error": f"Unknown template: {template}"}), 400
|
|
template_path = os.path.join(config.TEMPLATE_IMAGES_DIR,
|
|
templates[template]["original_filename"])
|
|
|
|
logger.info("[%s] verifying against %s via external endpoint", safe_id, template)
|
|
try:
|
|
result = verify.verify_images(template_path, upload_path)
|
|
except requests.exceptions.RequestException as e:
|
|
logger.warning("[%s] verification endpoint unreachable: %s", safe_id, e)
|
|
return jsonify({"error": f"Verification service unavailable: {e}"}), 502
|
|
except Exception as e:
|
|
logger.exception("[%s] verification failed", safe_id)
|
|
return jsonify({"error": f"Verification failed: {e}"}), 500
|
|
|
|
return jsonify(result)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
logger.info("Starting Vase Matcher on %s:%d (logs: %s)",
|
|
config.HOST, config.PORT, config.LOG_FILE)
|
|
engine.bootstrap()
|
|
app.run(host=config.HOST, port=config.PORT, debug=False, threaded=True,
|
|
use_reloader=False)
|