1st

pipeline/__init__.py

@@ -0,0 +1,51 @@
+import os
+import torch
+from flask import Flask
+
+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...")
+# IMPORTANT: The path is now relative to the root of the project
+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)
+
+# Store models in a dictionary to pass to logic functions
+models = {
+    "processor_gnd": processor_gnd,
+    "model_gnd": model_gnd,
+    "predictor": predictor,
+    "tokenizer_text": tokenizer_text,
+    "model_text": model_text,
+    "device": device
+}
+
+print("✅ All models loaded successfully.")
+print("="*50)
+
+# Import routes after app and models are defined to avoid circular imports
+from pipeline import routes