final commit 5000

app.py

@@ -278,9 +278,10 @@ def _background_bootstrap():
         if not os.path.exists(vit_ckpt):
             BOOT_STATUS = "training-vit"
             subprocess.run([
-                "python", "train_vit_triplet.py", "--data_root", ds_root, "--epochs", "3",
-                "--batch_size", "4", "--lr", "5e-4", "--early_stopping_patience", "3",
-                "--skip_validation", "--max_samples", "200", "--export", os.path.join(export_dir, "vit_outfit_model.pth")
+                "python", "train_vit_triplet.py", "--data_root", ds_root, "--epochs", "10",
+                "--batch_size", "4", "--lr", "5e-4", "--early_stopping_patience", "5",
+                "--max_samples", "5000", "--triplet_margin", "0.5", "--gradient_clip", "1.0",
+                "--warmup_epochs", "2", "--export", os.path.join(export_dir, "vit_outfit_model.pth")
             ], check=False)
         service.reload_models()
         BOOT_STATUS = "ready"
@@ -389,7 +390,7 @@ def _stitch_strip(imgs: List[Image.Image], height: int = 256, pad: int = 6, bg=(
     return out
 
 
-def gradio_recommend(files: List[str], occasion: str, weather: str, num_outfits: int):
+def gradio_recommend(files: List[str], occasion: str, weather: str, num_outfits: int, outfit_style: str = "casual"):
     # Check model status first
     model_status = service.get_model_status()
     if not model_status["can_recommend"]:
@@ -415,7 +416,7 @@ def gradio_recommend(files: List[str], occasion: str, weather: str, num_outfits:
         {"id": f"item_{i}", "image": images[i], "category": None}
         for i in range(len(images))
     ]
-    res = service.compose_outfits(items, context={"occasion": occasion, "weather": weather, "num_outfits": int(num_outfits)})
+    res = service.compose_outfits(items, context={"occasion": occasion, "weather": weather, "num_outfits": int(num_outfits), "outfit_style": outfit_style})
     
     # Check if compose_outfits returned an error
     if res and isinstance(res[0], dict) and "error" in res[0]:
@@ -710,8 +711,9 @@ def start_training_simple(dataset_size: str, res_epochs: int, vit_epochs: int):
             log_message += f"\n🚀 Starting ViT training on {dataset_size} samples...\n"
             vit_result = subprocess.run([
                 "python", "train_vit_triplet.py", "--data_root", DATASET_ROOT, "--epochs", str(vit_epochs),
-                "--batch_size", "4", "--lr", "5e-4", "--early_stopping_patience", "3",
-                "--skip_validation", "--max_samples", "200", "--export", os.path.join(export_dir, "vit_outfit_model.pth")
+                "--batch_size", "4", "--lr", "5e-4", "--early_stopping_patience", "5",
+                "--max_samples", "5000", "--triplet_margin", "0.5", "--gradient_clip", "1.0",
+                "--warmup_epochs", "2", "--export", os.path.join(export_dir, "vit_outfit_model.pth")
             ] + dataset_args, capture_output=True, text=True, check=False)
             
             if vit_result.returncode == 0:
@@ -770,11 +772,13 @@ with gr.Blocks(fill_height=True, title="Dressify - Advanced Outfit Recommendatio
         with gr.Row():
             occasion = gr.Dropdown(choices=["casual", "business", "formal", "sport"], value="casual", label="Occasion")
             weather = gr.Dropdown(choices=["any", "hot", "mild", "cold", "rain"], value="any", label="Weather")
+            outfit_style = gr.Dropdown(choices=["casual", "smart_casual", "formal", "sporty"], value="casual", label="Outfit Style")
+        with gr.Row():
             num_outfits = gr.Slider(minimum=1, maximum=8, step=1, value=3, label="Number of outfits")
         out_gallery = gr.Gallery(label="Recommended Outfits", columns=1, height=320)
         out_json = gr.JSON(label="Outfit Details")
         btn2 = gr.Button("Generate Outfits", variant="primary")
-        btn2.click(fn=gradio_recommend, inputs=[inp2, occasion, weather, num_outfits], outputs=[out_gallery, out_json])
+        btn2.click(fn=gradio_recommend, inputs=[inp2, occasion, weather, num_outfits, outfit_style], outputs=[out_gallery, out_json])
     
     with gr.Tab("🔬 Advanced Training"):
         gr.Markdown("### 🎯 Comprehensive Training Parameter Control\nCustomize every aspect of model training for research and experimentation.")

inference.py

@@ -233,25 +233,148 @@ class InferenceService:
         if len(proc_items) < 2:
             return []
 
-        # 2) Candidate generation
+        # 2) Candidate generation with outfit templates
         rng = np.random.default_rng(int(context.get("seed", 42)))
         num_outfits = int(context.get("num_outfits", 3))
         min_size, max_size = 4, 6
         ids = list(range(len(proc_items)))
+        
+        # Outfit templates for cohesive styling
+        outfit_templates = {
+            "casual": {
+                "style": "relaxed, comfortable, everyday",
+                "preferred_categories": ["tshirt", "jean", "sneaker", "hoodie", "sweatpant"],
+                "color_palette": ["neutral", "denim", "white", "black", "gray"],
+                "accessory_limit": 2
+            },
+            "smart_casual": {
+                "style": "polished but relaxed, business casual",
+                "preferred_categories": ["shirt", "chino", "loafer", "blazer", "polo"],
+                "color_palette": ["navy", "white", "khaki", "brown", "gray"],
+                "accessory_limit": 3
+            },
+            "formal": {
+                "style": "professional, elegant, sophisticated",
+                "preferred_categories": ["blazer", "dress shirt", "dress pant", "oxford", "suit"],
+                "color_palette": ["navy", "black", "white", "gray", "charcoal"],
+                "accessory_limit": 4
+            },
+            "sporty": {
+                "style": "athletic, active, performance",
+                "preferred_categories": ["athletic shirt", "jogger", "running shoe", "tank", "legging"],
+                "color_palette": ["bright", "neon", "white", "black", "primary colors"],
+                "accessory_limit": 1
+            }
+        }
+        
+        # Select outfit template (can be passed in context or randomly selected)
+        template_name = context.get("outfit_style", rng.choice(list(outfit_templates.keys())))
+        template = outfit_templates[template_name]
 
         # Enhanced category-aware pools with diversity checks
         def cat_str(i: int) -> str:
             return (proc_items[i].get("category") or "").lower()
+        
+        def extract_color_from_category(category: str) -> str:
+            """Extract color information from category name"""
+            category_lower = category.lower()
+            color_keywords = {
+                "black": ["black", "dark", "charcoal", "navy"],
+                "white": ["white", "cream", "ivory", "off-white"],
+                "gray": ["gray", "grey", "silver", "ash"],
+                "brown": ["brown", "tan", "beige", "khaki", "camel"],
+                "blue": ["blue", "navy", "denim", "indigo", "royal"],
+                "red": ["red", "burgundy", "maroon", "crimson"],
+                "green": ["green", "olive", "emerald", "forest"],
+                "yellow": ["yellow", "gold", "mustard", "lemon"],
+                "pink": ["pink", "rose", "coral", "salmon"],
+                "purple": ["purple", "violet", "lavender", "plum"],
+                "orange": ["orange", "peach", "apricot", "tangerine"],
+                "neutral": ["neutral", "nude", "natural", "earth"]
+            }
+            
+            for color, keywords in color_keywords.items():
+                if any(kw in category_lower for kw in keywords):
+                    return color
+            return "unknown"
+        
+        def calculate_color_consistency_score(items: List[int]) -> float:
+            """Calculate color consistency score for outfit items"""
+            colors = [extract_color_from_category(cat_str(i)) for i in items]
+            color_counts = {}
+            for color in colors:
+                color_counts[color] = color_counts.get(color, 0) + 1
+            
+            # Prefer outfits with 2-3 dominant colors
+            dominant_colors = [c for c, count in color_counts.items() if count >= 2]
+            if len(dominant_colors) == 0:
+                return 0.5  # Neutral score for all different colors
+            elif len(dominant_colors) == 1:
+                return 0.8  # Good consistency
+            elif len(dominant_colors) == 2:
+                return 1.0  # Perfect balance
+            else:
+                return 0.3  # Too many dominant colors
+            
+        def calculate_style_consistency_score(items: List[int]) -> float:
+            """Calculate style consistency based on template preferences"""
+            categories = [cat_str(i) for i in items]
+            preferred_cats = template["preferred_categories"]
+            
+            matches = 0
+            for cat in categories:
+                if any(pref in cat for pref in preferred_cats):
+                    matches += 1
+            
+            return matches / len(categories) if categories else 0.0
 
         def get_category_type(cat: str) -> str:
-            """Map category to outfit slot type"""
-            if any(k in cat for k in ["top", "shirt", "tshirt", "blouse", "jacket", "hoodie", "sweater", "cardigan"]):
+            """Map category to outfit slot type with comprehensive taxonomy"""
+            cat_lower = cat.lower().strip()
+            
+            # Upper body items (tops, outerwear)
+            upper_keywords = [
+                "top", "shirt", "tshirt", "t-shirt", "blouse", "tank", "camisole", "cami",
+                "jacket", "blazer", "coat", "hoodie", "sweater", "pullover", "cardigan",
+                "vest", "waistcoat", "windbreaker", "bomber", "denim jacket", "leather jacket",
+                "polo", "henley", "tunic", "crop top", "bodysuit", "romper", "jumpsuit"
+            ]
+            
+            # Bottom items
+            bottom_keywords = [
+                "pant", "pants", "trouser", "trousers", "jean", "jeans", "denim",
+                "skirt", "short", "shorts", "legging", "leggings", "tights",
+                "chino", "khaki", "cargo", "jogger", "sweatpant", "sweatpants",
+                "culotte", "palazzo", "mini skirt", "midi skirt", "maxi skirt",
+                "bermuda", "capri", "bike short", "bike shorts"
+            ]
+            
+            # Footwear
+            shoe_keywords = [
+                "shoe", "shoes", "sneaker", "sneakers", "boot", "boots", "heel", "heels",
+                "sandal", "sandals", "flat", "flats", "loafer", "loafers", "oxford",
+                "pump", "pumps", "stiletto", "wedge", "ankle boot", "knee high boot",
+                "combat boot", "hiking boot", "running shoe", "athletic shoe",
+                "mule", "mules", "clog", "clogs", "espadrille", "espadrilles"
+            ]
+            
+            # Accessories (can have multiple)
+            accessory_keywords = [
+                "watch", "belt", "ring", "rings", "bracelet", "bracelets", "necklace", "necklaces",
+                "earring", "earrings", "bag", "bags", "handbag", "purse", "clutch", "tote",
+                "hat", "cap", "beanie", "scarf", "scarves", "glove", "gloves", "sunglass", "sunglasses",
+                "tie", "bow tie", "pocket square", "cufflink", "cufflinks", "brooch", "pin",
+                "hair accessory", "headband", "hair clip", "barrette", "scrunchy", "scrunchies"
+            ]
+            
+            # Check each category
+            if any(k in cat_lower for k in upper_keywords):
                 return "upper"
-            elif any(k in cat for k in ["pant", "trouser", "jean", "skirt", "short", "legging"]):
+            elif any(k in cat_lower for k in bottom_keywords):
                 return "bottom"
-            elif any(k in cat for k in ["shoe", "sneaker", "boot", "heel", "sandal", "flat"]):
+            elif any(k in cat_lower for k in shoe_keywords):
                 return "shoe"
-            elif any(k in cat for k in ["watch", "belt", "ring", "bracelet", "accessor", "bag", "hat", "scarf", "necklace"]):
+            elif any(k in cat_lower for k in accessory_keywords):
                 return "accessory"
             else:
                 return "other"
@@ -275,52 +398,50 @@ class InferenceService:
         
         for _ in range(num_samples):
             subset = []
-            used_categories = set()
             
-            # Ensure one item from each major category (no duplicates)
-            if uppers:
+            # EXACT SLOT CONSTRAINTS: Exactly 1 upper, 1 bottom, 1 shoe, ≤2 accessories
+            if uppers and bottoms and shoes:
+                # Core outfit: exactly 1 of each required slot
                 subset.append(int(rng.choice(uppers)))
-                used_categories.add("upper")
-            
-            if bottoms:
                 subset.append(int(rng.choice(bottoms)))
-                used_categories.add("bottom")
-            
-            if shoes:
                 subset.append(int(rng.choice(shoes)))
-                used_categories.add("shoe")
-            
-            # Add accessories (can have multiple, but ensure diversity)
-            if accs:
-                # Add 1-2 accessories
-                num_accs = rng.integers(1, min(3, len(accs) + 1))
-                available_accs = [i for i in accs if i not in subset]
-                if available_accs:
-                    selected_accs = rng.choice(available_accs, size=min(num_accs, len(available_accs)), replace=False)
-                    subset.extend(selected_accs.tolist())
-                    used_categories.add("accessory")
-            
-            # Add other items if available and we have space
-            if others and len(subset) < max_size:
-                available_others = [i for i in others if i not in subset]
-                if available_others:
-                    num_others = rng.integers(0, min(2, len(available_others) + 1))
-                    if num_others > 0:
-                        selected_others = rng.choice(available_others, size=min(num_others, len(available_others)), replace=False)
-                        subset.extend(selected_others.tolist())
-                        used_categories.add("other")
-            
-            # Ensure we have at least min_size items
-            if len(subset) < min_size:
-                remaining = [i for i in ids if i not in subset]
-                needed = min_size - len(subset)
-                if remaining:
-                    additional = rng.choice(remaining, size=min(needed, len(remaining)), replace=False)
-                    subset.extend(additional.tolist())
+                
+                # Add accessories based on template limit
+                if accs:
+                    max_accs = template["accessory_limit"]
+                    num_accs = rng.integers(1, min(max_accs + 1, len(accs) + 1))
+                    available_accs = [i for i in accs if i not in subset]
+                    if available_accs:
+                        selected_accs = rng.choice(available_accs, size=min(num_accs, len(available_accs)), replace=False)
+                        subset.extend(selected_accs.tolist())
+                
+                # Add 0-1 other items for variety (but not if it would exceed max_size)
+                if others and len(subset) < max_size:
+                    available_others = [i for i in others if i not in subset]
+                    if available_others and rng.random() < 0.3:  # 30% chance to add other item
+                        subset.append(int(rng.choice(available_others)))
+            else:
+                # Fallback: ensure we have at least 3 items with category diversity
+                required_categories = []
+                if uppers: required_categories.append(("upper", uppers))
+                if bottoms: required_categories.append(("bottom", bottoms))
+                if shoes: required_categories.append(("shoe", shoes))
+                
+                # Add one from each available required category
+                for cat_type, cat_items in required_categories:
+                    subset.append(int(rng.choice(cat_items)))
+                
+                # Add accessories if available
+                if accs and len(subset) < max_size:
+                    num_accs = rng.integers(1, min(3, len(accs) + 1))
+                    available_accs = [i for i in accs if i not in subset]
+                    if available_accs:
+                        selected_accs = rng.choice(available_accs, size=min(num_accs, len(available_accs)), replace=False)
+                        subset.extend(selected_accs.tolist())
             
-            # Remove duplicates and ensure valid outfit
+            # Remove duplicates and validate
             subset = list(set(subset))
-            if len(subset) >= 2:  # At least 2 items for a valid outfit
+            if len(subset) >= 3:  # At least 3 items for a valid outfit
                 candidates.append(subset)
 
         # 3) Score using ViT
@@ -334,46 +455,132 @@ class InferenceService:
             s = self.vit.score_compatibility(embs).item()
             return float(s)
 
-        # Filter out invalid outfits (duplicate categories)
+        # Enhanced validation with strict slot constraints
         def is_valid_outfit(subset: List[int]) -> bool:
-            """Check if outfit has no duplicate categories"""
+            """Check if outfit meets exact slot requirements"""
             categories = [get_category_type(cat_str(i)) for i in subset]
-            # Allow multiple accessories, but only one of each other category
             category_counts = {}
+            
             for cat in categories:
-                if cat == "accessory":
-                    continue  # Allow multiple accessories
-                if cat in category_counts:
-                    return False  # Duplicate non-accessory category
-                category_counts[cat] = 1
+                category_counts[cat] = category_counts.get(cat, 0) + 1
+            
+            # STRICT VALIDATION:
+            # - Exactly 1 upper, 1 bottom, 1 shoe
+            # - ≤2 accessories
+            # - No other duplicates
+            if category_counts.get("upper", 0) != 1:
+                return False
+            if category_counts.get("bottom", 0) != 1:
+                return False
+            if category_counts.get("shoe", 0) != 1:
+                return False
+            if category_counts.get("accessory", 0) > 2:
+                return False
+            if category_counts.get("other", 0) > 1:
+                return False
+            
             return True
         
-        # Score and filter valid outfits
+        def calculate_outfit_penalty(subset: List[int], base_score: float) -> float:
+            """Calculate penalty-adjusted score for outfit quality with style/color bonuses"""
+            categories = [get_category_type(cat_str(i)) for i in subset]
+            category_counts = {}
+            
+            for cat in categories:
+                category_counts[cat] = category_counts.get(cat, 0) + 1
+            
+            penalty = 0.0
+            bonus = 0.0
+            
+            # Missing core slots: -∞ penalty
+            if category_counts.get("upper", 0) == 0:
+                penalty += -1000.0
+            if category_counts.get("bottom", 0) == 0:
+                penalty += -1000.0
+            if category_counts.get("shoe", 0) == 0:
+                penalty += -1000.0
+            
+            # Duplicate non-accessory categories: -∞ penalty
+            for cat, count in category_counts.items():
+                if cat != "accessory" and count > 1:
+                    penalty += -1000.0
+            
+            # Too many accessories: moderate penalty
+            max_accs = template["accessory_limit"]
+            if category_counts.get("accessory", 0) > max_accs:
+                penalty += -2.0
+            
+            # Unbalanced outfit: small penalty
+            if len(subset) < 3:
+                penalty += -1.0
+            elif len(subset) > 6:
+                penalty += -0.5
+            
+            # Style consistency bonus
+            style_score = calculate_style_consistency_score(subset)
+            bonus += style_score * 0.5  # Up to 0.5 bonus for style consistency
+            
+            # Color consistency bonus
+            color_score = calculate_color_consistency_score(subset)
+            bonus += color_score * 0.3  # Up to 0.3 bonus for color consistency
+            
+            # Template adherence bonus
+            if style_score > 0.6:  # Good style match
+                bonus += 0.2
+            
+            return base_score + penalty + bonus
+        
+        # Score and filter valid outfits with penalty adjustment
         valid_candidates = [subset for subset in candidates if is_valid_outfit(subset)]
         if not valid_candidates:
             # Fallback: use all candidates if no valid ones found
             valid_candidates = candidates
         
-        scored = [(subset, score_subset(subset)) for subset in valid_candidates]
+        # Score with penalty adjustment
+        scored = []
+        for subset in valid_candidates:
+            base_score = score_subset(subset)
+            adjusted_score = calculate_outfit_penalty(subset, base_score)
+            scored.append((subset, adjusted_score, base_score))
+        
+        # Sort by penalty-adjusted score
         scored.sort(key=lambda x: x[1], reverse=True)
         topk = scored[:num_outfits]
         
         results = []
-        for subset, score in topk:
+        for subset, adjusted_score, base_score in topk:
             # Double-check validity and get item details
             outfit_items = []
             for i in subset:
                 item = proc_items[i]
                 outfit_items.append({
                     "id": item["id"],
-                    "category": item.get("category", "unknown")
+                    "category": item.get("category", "unknown"),
+                    "category_type": get_category_type(item.get("category", ""))
                 })
             
+            # Calculate additional metrics
+            style_score = calculate_style_consistency_score(subset)
+            color_score = calculate_color_consistency_score(subset)
+            colors = [extract_color_from_category(cat_str(i)) for i in subset]
+            
             results.append({
                 "item_ids": [item["id"] for item in outfit_items],
                 "items": outfit_items,
-                "score": float(score),
-                "categories": [item["category"] for item in outfit_items]
+                "score": float(adjusted_score),
+                "base_score": float(base_score),
+                "categories": [item["category"] for item in outfit_items],
+                "category_types": [item["category_type"] for item in outfit_items],
+                "outfit_size": len(outfit_items),
+                "is_valid": is_valid_outfit(subset),
+                "template": {
+                    "name": template_name,
+                    "style": template["style"],
+                    "style_score": float(style_score),
+                    "color_score": float(color_score),
+                    "colors": colors,
+                    "accessory_limit": template["accessory_limit"]
+                }
             })
         
         return results

train_vit_triplet.py

@@ -26,13 +26,15 @@ def parse_args() -> argparse.Namespace:
     p.add_argument("--batch_size", type=int, default=4)
     p.add_argument("--lr", type=float, default=5e-4)
     p.add_argument("--embedding_dim", type=int, default=512)
-    p.add_argument("--triplet_margin", type=float, default=0.3)
+    p.add_argument("--triplet_margin", type=float, default=0.5)
     p.add_argument("--export", type=str, default="models/exports/vit_outfit_model.pth")
     p.add_argument("--eval_every", type=int, default=1)
     p.add_argument("--skip_validation", action="store_true", help="Skip validation for faster training")
-    p.add_argument("--max_samples", type=int, default=500, help="Maximum number of training samples (for faster testing)")
-    p.add_argument("--early_stopping_patience", type=int, default=10, help="Early stopping patience")
+    p.add_argument("--max_samples", type=int, default=5000, help="Maximum number of training samples (for better quality)")
+    p.add_argument("--early_stopping_patience", type=int, default=5, help="Early stopping patience")
     p.add_argument("--min_delta", type=float, default=1e-4, help="Minimum change to qualify as improvement")
+    p.add_argument("--gradient_clip", type=float, default=1.0, help="Gradient clipping value")
+    p.add_argument("--warmup_epochs", type=int, default=2, help="Learning rate warmup epochs")
     return p.parse_args()
 
 
@@ -113,6 +115,17 @@ def main() -> None:
 
     optimizer = optim.AdamW(model.parameters(), lr=args.lr, weight_decay=5e-2)
     triplet = nn.TripletMarginWithDistanceLoss(distance_function=lambda x, y: 1 - nn.functional.cosine_similarity(x, y), margin=args.triplet_margin)
+    
+    # Learning rate scheduler with warmup
+    total_steps = len(loader) * args.epochs
+    warmup_steps = len(loader) * args.warmup_epochs
+    scheduler = torch.optim.lr_scheduler.OneCycleLR(
+        optimizer, 
+        max_lr=args.lr, 
+        total_steps=total_steps,
+        pct_start=warmup_steps/total_steps,
+        anneal_strategy='cos'
+    )
 
     export_dir = ensure_export_dir(os.path.dirname(args.export) or "models/exports")
     best_loss = float("inf")
@@ -159,7 +172,13 @@ def main() -> None:
                 
                 optimizer.zero_grad(set_to_none=True)
                 loss.backward()
+                
+                # Gradient clipping for stability
+                if args.gradient_clip > 0:
+                    torch.nn.utils.clip_grad_norm_(model.parameters(), args.gradient_clip)
+                
                 optimizer.step()
+                scheduler.step()  # Update learning rate
                 
                 # Collect metrics (simplified for ViT training)
                 # Note: ViT training uses outfit-level embeddings, not classification predictions