Update app.py

app.py

@@ -1,10 +1,9 @@
 # -------------------------------------------------------------------------- #
-# UNIFIED AI SERVICE V3.3 (Added Markings Comparison)
+# UNIFIED AI SERVICE V3.4 (Color-Enhanced Segmentation)
 # -------------------------------------------------------------------------- #
 # This service uses DINOv2 for image embeddings and BGE for text embeddings.
-# - Filtering is handled by the Node.js backend.
-# - For debugging, segmented images are uploaded to Uploadcare and the URL
-#   is printed to the console log.
+# - The segmentation prompt now includes colors for better accuracy.
+# - For debugging, segmented images are uploaded to Uploadcare.
 # --------------------------------------------------------------------------
 import sys
 sys.stdout.reconfigure(line_buffering=True)
@@ -30,7 +29,6 @@ from transformers import AutoProcessor as AutoGndProcessor, AutoModelForZeroShot
 
 app = Flask(__name__)
 
-# --- UPDATED: Added "markings" to the list of fields to compare ---
 TEXT_FIELDS_TO_EMBED = ["brand", "material", "markings"]
 SCORE_WEIGHTS = {
     "text_score": 0.4,
@@ -113,9 +111,15 @@ def jaccard_similarity(set1, set2):
         return 1.0 if not intersection else 0.0
     return len(intersection) / len(union)
 
-def segment_guided_object(image: Image.Image, object_label: str) -> Image.Image:
-    prompt = f"a {object_label}."
-    print(f"  [Segment] Using simple prompt: '{prompt}'")
+def segment_guided_object(image: Image.Image, object_label: str, colors: list = []) -> Image.Image:
+    # --- UPDATED: Create a more descriptive prompt using colors ---
+    color_str = " ".join(c.lower() for c in colors if c)
+    if color_str:
+        prompt = f"a {color_str} {object_label}."
+    else:
+        prompt = f"a {object_label}."
+    
+    print(f"  [Segment] Using prompt: '{prompt}'")
     image_rgb = image.convert("RGB")
     image_np = np.array(image_rgb)
     h, w = image_np.shape[:2]
@@ -179,7 +183,6 @@ def process_item():
         data = request.json
         print(f"\n[PROCESS] Received request for: {data.get('objectName')}")
 
-        # --- UPDATED: Added markings_embedding ---
         response = {
             "canonicalLabel": data.get('objectName', '').lower().strip(),
             "brand_embedding": get_text_embedding(data.get('brand')),
@@ -196,7 +199,8 @@ def process_item():
                     img_response.raise_for_status()
                     image = Image.open(BytesIO(img_response.content))
 
-                    segmented_image = segment_guided_object(image, data['objectName'])
+                    # --- UPDATED: Pass colors to the segmentation function ---
+                    segmented_image = segment_guided_object(image, data['objectName'], data.get('colors', []))
                     debug_url = upload_to_uploadcare(segmented_image)
                     print(f"    - 🐞 DEBUG URL: {debug_url}")