Spaces:
Sleeping
Sleeping
| # shared_resources.py | |
| from pymongo import MongoClient | |
| from transformers import Dinov2Model, AutoImageProcessor | |
| import torch.nn as nn | |
| import torch | |
| from core.config import MONGO_URI, DB_NAME, DURING_COLLECTION, MODEL_NAME, MODEL_PATH, DEVICE, TEMP_HISTORY_COLLECTION | |
| import os | |
| from PIL import Image | |
| from typing import Dict, Any, List | |
| model = None | |
| processor = None | |
| db = None | |
| client = None | |
| class FineTunedDINOv2(nn.Module): | |
| def __init__(self): | |
| super(FineTunedDINOv2, self).__init__() | |
| self.backbone = Dinov2Model.from_pretrained(MODEL_NAME) | |
| def forward(self, x): | |
| outputs = self.backbone(x) | |
| cls_token = outputs.last_hidden_state[:, 0] | |
| return nn.functional.normalize(cls_token, p=2, dim=1) | |
| def load_resources(): | |
| global model, processor, db, client | |
| print("🚀 Starting up application...") | |
| # DB connect | |
| try: | |
| client = MongoClient(MONGO_URI) | |
| client.admin.command('ping') | |
| db = client[DB_NAME] | |
| print("✅ Đã kết nối thành công đến MongoDB Atlas.") | |
| except Exception as e: | |
| db = None | |
| client = None | |
| print(f"❌ Lỗi kết nối MongoDB: {type(e).__name__}: {e}") | |
| # không raise để app vẫn khởi động (endpoints kiểm tra db is None) | |
| try: | |
| temp_col = db[TEMP_HISTORY_COLLECTION] | |
| # Tạo index cho trường "created_at", hết hạn sau 24h | |
| temp_col.create_index("created_at", expireAfterSeconds=86400) | |
| print("✅ TTL Index for Temp History ensured (24h).") | |
| except Exception as e: | |
| print(f"⚠️ Could not create TTL index: {e}") | |
| # Model load | |
| try: | |
| model_instance = FineTunedDINOv2() | |
| if os.path.exists(MODEL_PATH): | |
| checkpoint = torch.load(MODEL_PATH, map_location=DEVICE) | |
| state_dict = checkpoint.get("model_state_dict", checkpoint) | |
| new_state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()} | |
| model_instance.load_state_dict(new_state_dict, strict=False) | |
| model_instance.to(DEVICE) | |
| model_instance.eval() | |
| model = model_instance | |
| processor = AutoImageProcessor.from_pretrained(MODEL_NAME) | |
| print("✅ Fine-tuned DINOv2 ready!") | |
| except Exception as e: | |
| model = None | |
| processor = None | |
| print(f"❌ Model loading failed: {type(e).__name__}: {e}") | |
| def embed_image(img: Image.Image) -> List[float]: | |
| if model is None: | |
| raise RuntimeError("Model not loaded.") | |
| inputs = processor(images=img, return_tensors="pt").pixel_values.to(DEVICE) | |
| with torch.no_grad(): | |
| embedding = model(inputs) | |
| return embedding.cpu().numpy().flatten().tolist() | |
| def search_similar_landmark(query_vector: List[float], collection_name: str, limit: int = 1): | |
| if db is None: | |
| raise RuntimeError("Database not connected.") | |
| during_col = db[collection_name] | |
| pipeline = [ | |
| { | |
| "$vectorSearch": { | |
| "index": "vector_index", | |
| "path": "embedding", | |
| "queryVector": query_vector, | |
| "numCandidates": 1024, | |
| "limit": limit | |
| } | |
| }, | |
| {"$project": {"_id": 0, "landmark_id": 1, "image_url": 1, "score": {"$meta": "vectorSearchScore"}}} | |
| ] | |
| try: | |
| return list(during_col.aggregate(pipeline)) | |
| except Exception as e: | |
| # Nếu index tên khác, log thông báo rõ ràng | |
| print(f"❌ Vector search failed: {type(e).__name__}: {e}") | |
| raise | |