keyword filter

indexer.py

@@ -14,25 +14,15 @@ from PIL import Image
 from tqdm import tqdm
 from transformers import SiglipModel, SiglipProcessor
 
+from keyword_filters import (
+    CATEGORY_SYNONYMS,
+    COLOR_SYNONYMS,
+    VIBE_SYNONYMS,
+    extract_keywords,
+)
+
 IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".gif"}
 T = TypeVar("T")
-CATEGORY_SYNONYMS = {
-    "모자": ["모자", "헬름", "헬멧", "햇", "보닛", "캡"],
-    "신발": ["신발", "슈즈", "부츠", "샌들"],
-    "장갑": ["장갑", "글러브"],
-    "무기": ["무기", "검", "소드", "대검", "스태프", "완드", "활", "석궁", "창", "스피어", "폴암", "도끼", "단검", "너클", "건", "총", "클로"],
-    "상의": ["상의", "셔츠", "자켓", "코트", "로브", "블라우스"],
-    "하의": ["하의", "바지", "팬츠", "스커트"],
-    "망토": ["망토", "케이프", "cape"],
-    "귀걸이": ["귀걸이", "귀고리", "이어링"],
-    "반지": ["반지", "링"],
-    "목걸이": ["목걸이", "펜던트", "네클리스"],
-    "벨트": ["벨트"],
-    "얼굴장식": ["얼굴장식", "얼굴 장식"],
-    "눈장식": ["눈장식", "눈 장식"],
-    "보조무기": ["보조무기", "보조 무기"],
-    "방패": ["방패", "쉴드", "실드"],
-}
 
 
 def parse_args() -> argparse.Namespace:
@@ -154,12 +144,39 @@ def detect_category(texts: List[str]) -> Optional[str]:
     return None
 
 
-def load_labels(labels_path: Path) -> Dict[str, Dict[str, str]]:
+def collect_label_texts(
+    item_name: Optional[str],
+    label_ko: Optional[str],
+    tags: List[str],
+    query_variants: List[str],
+    attributes: Dict[str, object],
+    item_type_guess: Optional[str],
+) -> List[str]:
+    texts: List[str] = []
+    for value in (item_name, label_ko, item_type_guess):
+        if value:
+            texts.append(value)
+    texts.extend(tag for tag in tags if tag)
+    texts.extend(variant for variant in query_variants if variant)
+    for value in attributes.values():
+        if isinstance(value, list):
+            for entry in value:
+                entry_norm = normalize_label(entry)
+                if entry_norm:
+                    texts.append(entry_norm)
+        else:
+            entry_norm = normalize_label(value)
+            if entry_norm:
+                texts.append(entry_norm)
+    return texts
+
+
+def load_labels(labels_path: Path) -> Dict[str, Dict[str, object]]:
     if not labels_path.exists():
         print(f"Labels file not found, continuing without labels: {labels_path}")
         return {}
 
-    label_map: Dict[str, Dict[str, str]] = {}
+    label_map: Dict[str, Dict[str, object]] = {}
     with labels_path.open("r", encoding="utf-8") as file:
         for line_no, line in enumerate(file, start=1):
             line = line.strip()
@@ -180,6 +197,15 @@ def load_labels(labels_path: Path) -> Dict[str, Dict[str, str]]:
             tags = record.get("tags_ko") or []
             tag_texts = [normalize_label(tag) for tag in tags if tag is not None]
             tag_texts = [tag for tag in tag_texts if tag]
+            query_variants = record.get("query_variants_ko") or []
+            variant_texts = [
+                normalize_label(variant)
+                for variant in query_variants
+                if variant is not None
+            ]
+            variant_texts = [variant for variant in variant_texts if variant]
+            attributes = record.get("attributes") or {}
+            item_type_guess = normalize_label(attributes.get("item_type_guess"))
             if not item_name and not label_ko and not tag_texts:
                 continue
 
@@ -190,9 +216,27 @@ def load_labels(labels_path: Path) -> Dict[str, Dict[str, str]]:
             if label_ko:
                 label_map[normalized_path]["label_ko"] = label_ko
                 label_map[normalized_path]["label"] = label_ko
-            category = detect_category([item_name or "", label_ko or "", *tag_texts])
+            texts = collect_label_texts(
+                item_name,
+                label_ko,
+                tag_texts,
+                variant_texts,
+                attributes,
+                item_type_guess,
+            )
+            category = detect_category(texts)
             if category:
                 label_map[normalized_path]["category"] = category
+            colors = extract_keywords(texts, COLOR_SYNONYMS)
+            if colors:
+                label_map[normalized_path]["colors"] = colors
+                for color in colors:
+                    label_map[normalized_path][f"color_{color}"] = True
+            vibes = extract_keywords(texts, VIBE_SYNONYMS)
+            if vibes:
+                label_map[normalized_path]["vibes"] = vibes
+                for vibe in vibes:
+                    label_map[normalized_path][f"vibe_{vibe}"] = True
 
     print(f"Loaded labels for {len(label_map)} images from {labels_path}")
     return label_map

keyword_filters.py

@@ -0,0 +1,118 @@
+from __future__ import annotations
+
+from typing import Dict, Iterable, List
+
+
+CATEGORY_SYNONYMS = {
+    "모자": ["모자", "헬름", "헬멧", "햇", "보닛", "캡", "hat", "cap", "helmet"],
+    "신발": ["신발", "슈즈", "부츠", "샌들", "shoes", "shoe", "boots", "sandal"],
+    "장갑": ["장갑", "글러브", "glove", "gloves"],
+    "무기": [
+        "무기",
+        "검",
+        "소드",
+        "대검",
+        "스태프",
+        "완드",
+        "활",
+        "석궁",
+        "창",
+        "스피어",
+        "폴암",
+        "도끼",
+        "단검",
+        "너클",
+        "건",
+        "총",
+        "클로",
+        "weapon",
+        "sword",
+        "staff",
+        "wand",
+        "bow",
+        "spear",
+        "axe",
+        "dagger",
+        "gun",
+        "claw",
+    ],
+    "상의": [
+        "상의",
+        "셔츠",
+        "자켓",
+        "코트",
+        "로브",
+        "블라우스",
+        "top",
+        "shirt",
+        "jacket",
+        "coat",
+        "robe",
+        "blouse",
+    ],
+    "하의": ["하의", "바지", "팬츠", "스커트", "bottom", "pants", "skirt"],
+    "망토": ["망토", "케이프", "cape", "날개", "윙", "wing", "wings"],
+    "귀걸이": ["귀걸이", "귀고리", "이어링", "earring", "earrings"],
+    "반지": ["반지", "링", "ring"],
+    "목걸이": ["목걸이", "펜던트", "네클리스", "necklace", "pendant"],
+    "벨트": ["벨트", "belt"],
+    "얼굴장식": ["얼굴장식", "얼굴 장식"],
+    "눈장식": ["눈장식", "눈 장식"],
+    "보조무기": ["보조무기", "보조 무기", "sub weapon", "subweapon", "offhand"],
+    "방패": ["방패", "쉴드", "실드", "shield"],
+}
+
+COLOR_SYNONYMS = {
+    "black": ["검은", "검정", "블랙", "black"],
+    "white": ["흰", "하얀", "화이트", "white"],
+    "gray": ["회색", "그레이", "gray"],
+    "silver": ["은색", "실버", "silver"],
+    "gold": ["금색", "골드", "gold"],
+    "red": ["빨간", "빨강", "레드", "적색", "붉은", "red"],
+    "pink": ["핑크", "분홍", "분홍색", "핑크색", "pink"],
+    "orange": ["주황", "오렌지", "orange"],
+    "yellow": ["노란", "노랑", "옐로", "yellow"],
+    "green": ["초록", "녹색", "그린", "green"],
+    "blue": ["파란", "파랑", "블루", "blue", "하늘색", "스카이", "sky"],
+    "purple": ["보라", "퍼플", "purple"],
+    "brown": ["갈색", "브라운", "brown"],
+    "beige": ["베이지", "beige"],
+    "mint": ["민트", "mint"],
+    "teal": ["청록", "teal", "터쿼이즈", "turquoise"],
+    "navy": ["남색", "네이비", "navy"],
+}
+
+VIBE_SYNONYMS = {
+    "cute": ["귀여움", "귀여운", "귀엽", "큐트", "cute", "사랑스러운", "lovely"],
+    "sporty": ["스포티", "스포츠", "sporty", "sports", "스포티한"],
+    "casual": ["캐주얼", "casual"],
+    "luxury": ["고급스러움", "고급", "luxury", "classy", "품격", "vip", "VIP"],
+    "elegant": ["우아", "elegant", "고상", "세련", "세련된"],
+    "playful": ["유쾌한", "funny", "playful", "장난", "발랄"],
+    "bright": ["빛나는", "sparkle", "glitter", "반짝", "sparkling"],
+    "powerful": ["강력한", "전투적인", "전투용", "powerful", "강인"],
+    "romantic": ["로맨틱", "romance", "romantic", "설렘", "사랑"],
+    "mysterious": ["신비", "mysterious", "묘한"],
+    "retro": ["레트로", "retro", "빈티지", "vintage", "클래식", "classic", "고전적인"],
+    "futuristic": ["futuristic", "미래", "사이버", "sf"],
+    "sweet": ["달달", "달콤", "sweet", "상큼"],
+    "unique": ["유니크", "unique", "독특", "개성", "특별"],
+    "calm": ["고요한", "차분", "calm"],
+    "dark": ["다크", "dark", "ダーク", "어두운"],
+}
+
+
+def extract_keywords(
+    texts: Iterable[str], synonyms: Dict[str, List[str]]
+) -> List[str]:
+    lowered_texts = [text.lower() for text in texts if text]
+    if not lowered_texts:
+        return []
+    hits: List[str] = []
+    for canonical, variants in synonyms.items():
+        for variant in variants:
+            variant_lower = variant.lower()
+            if any(variant_lower in text for text in lowered_texts):
+                hits.append(canonical)
+                break
+    return hits

main.py

@@ -16,25 +16,14 @@ from peft import PeftModel
 from pydantic import BaseModel, Field
 from transformers import SiglipModel, SiglipProcessor
 
+from keyword_filters import (
+    CATEGORY_SYNONYMS,
+    COLOR_SYNONYMS,
+    VIBE_SYNONYMS,
+    extract_keywords,
+)
 
 DATA_DIR = (Path(__file__).resolve().parent / "data/2026-01-11").resolve()
-CATEGORY_SYNONYMS = {
-    "모자": ["모자", "헬름", "헬멧", "햇", "보닛", "캡"],
-    "신발": ["신발", "슈즈", "부츠", "샌들"],
-    "장갑": ["장갑", "글러브"],
-    "무기": ["무기", "검", "소드", "대검", "스태프", "완드", "활", "석궁", "창", "스피어", "폴암", "도끼", "단검", "너클", "건", "총", "클로"],
-    "상의": ["상의", "셔츠", "자켓", "코트", "로브", "블라우스"],
-    "하의": ["하의", "바지", "팬츠", "스커트"],
-    "망토": ["망토", "케이프", "cape"],
-    "귀걸이": ["귀걸이", "귀고리", "이어링"],
-    "반지": ["반지", "링"],
-    "목걸이": ["목걸이", "펜던트", "네클리스"],
-    "벨트": ["벨트"],
-    "얼굴장식": ["얼굴장식", "얼굴 장식"],
-    "눈장식": ["눈장식", "눈 장식"],
-    "보조무기": ["보조무기", "보조 무기"],
-    "방패": ["방패", "쉴드", "실드"],
-}
 
 
 class SearchRequest(BaseModel):
@@ -51,21 +40,59 @@ def resolve_adapter_path(adapter_path: Path) -> Path:
     return adapter_path
 
 
-def extract_category_keywords(query: str) -> List[str]:
-    keywords: List[str] = []
-    lowered_query = query.lower()
-    for category, variants in CATEGORY_SYNONYMS.items():
-        for variant in variants:
-            if variant.lower() in lowered_query and category not in keywords:
-                keywords.append(category)
-                break
-    return keywords
+def extract_query_filters(query: str) -> Dict[str, List[str]]:
+    texts = [query]
+    return {
+        "categories": extract_keywords(texts, CATEGORY_SYNONYMS),
+        "colors": extract_keywords(texts, COLOR_SYNONYMS),
+        "vibes": extract_keywords(texts, VIBE_SYNONYMS),
+    }
 
 
-def build_metadata_filter(keywords: List[str]) -> Dict[str, Any] | None:
-    if not keywords:
+def build_where_filter(
+    categories: List[str], colors: List[str], vibes: List[str]
+) -> Dict[str, Any] | None:
+    clauses: List[Dict[str, Any]] = []
+    if categories:
+        clauses.append({"category": {"$in": categories}})
+    if colors:
+        clauses.append({"$and": [{f"color_{color}": True} for color in colors]})
+    if vibes:
+        clauses.append({"$and": [{f"vibe_{vibe}": True} for vibe in vibes]})
+    if not clauses:
         return None
-    return {"category": {"$in": keywords}}
+    if len(clauses) == 1:
+        return clauses[0]
+    return {"$and": clauses}
+
+
+def build_filter_candidates(filters: Dict[str, List[str]]) -> List[Dict[str, Any]]:
+    parts = {
+        "category": filters.get("categories") or [],
+        "color": filters.get("colors") or [],
+        "vibe": filters.get("vibes") or [],
+    }
+    candidates: List[Dict[str, Any]] = []
+    combos = [
+        ("category", "color", "vibe"),
+        ("category", "color"),
+        ("category", "vibe"),
+        ("color", "vibe"),
+        ("category",),
+        ("color",),
+        ("vibe",),
+    ]
+    for combo in combos:
+        if not all(parts[facet] for facet in combo):
+            continue
+        where_filter = build_where_filter(
+            parts["category"] if "category" in combo else [],
+            parts["color"] if "color" in combo else [],
+            parts["vibe"] if "vibe" in combo else [],
+        )
+        if where_filter:
+            candidates.append(where_filter)
+    return candidates
 
 
 @asynccontextmanager
@@ -150,11 +177,11 @@ def search(payload: SearchRequest) -> Dict[str, Any]:
 
     query_embedding = text_embeds[0].detach().cpu().tolist()
 
-    filter_keywords = extract_category_keywords(query)
-    where_filter = build_metadata_filter(filter_keywords)
+    filter_parts = extract_query_filters(query)
+    where_candidates = build_filter_candidates(filter_parts)
 
     results = None
-    if where_filter:
+    for where_filter in where_candidates:
         try:
             results = collection.query(
                 query_embeddings=[query_embedding],
@@ -163,8 +190,11 @@ def search(payload: SearchRequest) -> Dict[str, Any]:
                 include=["distances", "metadatas"],
             )
         except Exception as exc:  # noqa: BLE001
-            print(f"Filtered query failed ({exc}); falling back to vector-only.")
+            print(f"Filtered query failed ({exc}); trying less strict.")
             results = None
+            continue
+        if results and results.get("ids") and results["ids"][0]:
+            break
 
     if not results or not results.get("ids") or not results["ids"][0]:
         results = collection.query(