import os
import torch
from flask import Flask
from sentence_transformers.cross_encoder import CrossEncoder 

FEATURE_WEIGHTS = {"shape": 0.4, "color": 0.5, "texture": 0.1}
FINAL_SCORE_THRESHOLD = 0.5

# create flask app
app = Flask(__name__)

# load models
print("="*50)
print("🚀 Initializing application and loading models...")
device_name = os.environ.get("device", "cpu")
device = torch.device('cuda' if 'cuda' in device_name and torch.cuda.is_available() else 'cpu')
print(f"🧠 Using device: {device}")

from transformers import AutoProcessor, AutoModelForZeroShotObjectDetection, AutoTokenizer, AutoModel
from segment_anything import SamPredictor, sam_model_registry

print("...Loading Grounding DINO model...")
gnd_model_id = "IDEA-Research/grounding-dino-tiny"
processor_gnd = AutoProcessor.from_pretrained(gnd_model_id)
model_gnd = AutoModelForZeroShotObjectDetection.from_pretrained(gnd_model_id).to(device)

print("...Loading Segment Anything (SAM) model...")
sam_checkpoint = "sam_vit_b_01ec64.pth"
sam_model = sam_model_registry["vit_b"](checkpoint=sam_checkpoint).to(device)
predictor = SamPredictor(sam_model)

print("...Loading BGE model for text embeddings...")
bge_model_id = "BAAI/bge-small-en-v1.5"
tokenizer_text = AutoTokenizer.from_pretrained(bge_model_id)
model_text = AutoModel.from_pretrained(bge_model_id).to(device)


models = {
    "processor_gnd": processor_gnd,
    "model_gnd": model_gnd,
    "predictor": predictor,
    "tokenizer_text": tokenizer_text,
    "model_text": model_text,
    "device": device
}


print("...Loading Cross-Encoder model for re-ranking...")

cross_encoder_model = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2', device=device)


models["cross_encoder"] = cross_encoder_model


print("✅ All models loaded successfully.")
print("="*50)

from pipeline import routes