Spaces:
Sleeping
Sleeping
File size: 3,025 Bytes
4b4aedc ac8561f 4b4aedc 3426b11 4b4aedc 3426b11 4b4aedc 9ca2c04 4b4aedc 9ca2c04 4b4aedc 9ca2c04 |
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 |
import os
import json
import base64
import numpy as np
from flask import Flask, request, jsonify, render_template
from langchain_experimental.open_clip.open_clip import OpenCLIPEmbeddings
from sklearn.metrics.pairwise import cosine_similarity
from io import BytesIO
from PIL import Image
# from matplotlib.offsetbox import OffsetImage, AnnotationBbox
from io import BytesIO
from pathlib import Path
# ============================== #
# INITIALIZE APP #
# ============================== #
app = Flask(__name__)
clip_embd = OpenCLIPEmbeddings()
BASE_DIR = Path("/app")
BLOCKS_DIR = BASE_DIR / "blocks"
# STATIC_DIR = BASE_DIR / "static"
# GEN_PROJECT_DIR = BASE_DIR / "generated_projects"
BACKDROP_DIR = BLOCKS_DIR / "Backdrops"
SPRITE_DIR = BLOCKS_DIR / "sprites"
CODE_BLOCKS_DIR = BLOCKS_DIR / "code_blocks"
# === new: outputs rooted under BASE_DIR ===
OUTPUT_DIR = BASE_DIR / "outputs"
# ============================== #
# LOAD PRE-COMPUTED EMBEDS #
# ============================== #
with open(f"{BLOCKS_DIR}/embeddings.json", "r") as f:
embedding_json = json.load(f)
image_paths = [item["file-path"] for item in embedding_json]
image_embeds = np.array([item["embeddings"] for item in embedding_json])
# ============================== #
# HELPER: Decode Base64 Image #
# ============================== #
def decode_base64_image(b64_string):
img_data = base64.b64decode(b64_string)
img = Image.open(BytesIO(img_data)).convert("RGB")
return img
# ============================== #
# API ROUTE #
# ============================== #
@app.route("/match", methods=["POST"])
def match_image():
data = request.get_json()
if "images" not in data:
return jsonify({"error": "No images provided"}), 400
results = []
for b64_img in data["images"]:
try:
# Convert Base64 → BytesIO
b_io = BytesIO(base64.b64decode(b64_img))
# Embed the query image
query_embed = np.array(clip_embd.embed_image([b_io]))
# Cosine similarity with stored embeddings
sims = cosine_similarity(query_embed, image_embeds)[0]
best_idx = np.argmax(sims)
results.append({
"input": b64_img[:50] + "...",
"best_match": {
"name": os.path.basename(image_paths[best_idx]),
"path": image_paths[best_idx],
"similarity": float(sims[best_idx])
}
})
except Exception as e:
results.append({"error": str(e)})
return jsonify(results)
# ============================== #
# SIMPLE HTML UI #
# ============================== #
@app.route("/", methods=["GET", "POST"])
def index():
return render_template("index.html")
# ============================== #
# MAIN ENTRY #
# ============================== #
if __name__ == "__main__":
app.run(debug=True, port=7860) |