Create main.py

main.py

@@ -0,0 +1,51 @@
+from fastapi import FastAPI
+from pymongo import MongoClient
+from sentence_transformers import SentenceTransformer
+from PIL import Image
+from io import BytesIO
+import numpy as np
+import requests
+from bson import ObjectId
+import os
+
+# ===============================================================
+# 🚀 FastAPI App
+# ===============================================================
+app = FastAPI(title="Educational Placemat Embedding API")
+
+# ===============================================================
+# 🔐 Database Connection
+# ===============================================================
+MONGO_URI = "mongodb+srv://anna_db_user:6zxpOoyMUqnpxrBS@similaritysearch.xblvd4g.mongodb.net/"
+client = MongoClient(MONGO_URI)
+db = client["similaritysearch"]
+
+# ===============================================================
+# 🧩 Load Lightweight CLIP Model
+# ===============================================================
+os.makedirs("/tmp/model_cache", exist_ok=True)
+model = SentenceTransformer(
+    "sentence-transformers/clip-ViT-B-16",
+    cache_folder="/tmp/model_cache"
+)
+
+# ===============================================================
+# 🖼️ Endpoint — Generate Image Embedding
+# ===============================================================
+@app.post("/generate_embedding")
+def generate_embedding(data: dict):
+    """
+    Generate a CLIP embedding for an image from URL and store it in MongoDB.
+    """
+    img_url = data["thumbnail"]
+    image = Image.open(BytesIO(requests.get(img_url).content)).convert("RGB").resize((512, 512))
+    emb = model.encode(image, convert_to_numpy=True, normalize_embeddings=True)
+    db.images.update_one({"_id": ObjectId(data["_id"])}, {"$set": {"embedding": emb.tolist()}})
+    return {"message": "✅ Embedding added successfully"}
+
+# ===============================================================
+# 🏠 Root route
+# ===============================================================
+@app.get("/")
+def home():
+    return {"status": "running", "model": "clip-ViT-B-16"}