Update pipeline/__init__.py

pipeline/__init__.py

@@ -1,22 +1,29 @@
 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
+# Create Flask app
 app = Flask(__name__)
 
-# load models
-print("="*50)
+# 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')
+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 transformers import (
+    AutoProcessor,
+    AutoModelForZeroShotObjectDetection,
+    AutoTokenizer,
+    AutoModel
+)
 from segment_anything import SamPredictor, sam_model_registry
 
 print("...Loading Grounding DINO model...")
@@ -25,6 +32,7 @@ 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)
@@ -34,7 +42,7 @@ 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,
@@ -44,16 +52,8 @@ models = {
     "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)
+print("=" * 50)
 
-from pipeline import routes
+# Import routes after app and models are defined to avoid circular imports
+from pipeline import routes