Upload evaluation/annex94_search_demo.py with huggingface_hub

evaluation/annex94_search_demo.py

@@ -32,12 +32,8 @@ PROJECT_ROOT = Path(__file__).resolve().parents[1]
 if str(PROJECT_ROOT) not in sys.path:
     sys.path.insert(0, str(PROJECT_ROOT))
 
-# Import custom models
-try:
-    from training.color_model import CLIPModel as ColorModel
-except ModuleNotFoundError:
-    ColorModel = None
-from training.hierarchy_model import Model as HierarchyModel, HierarchyExtractor
+# Import custom models via shared loaders
+from evaluation.utils.model_loader import load_color_model, load_hierarchy_model
 import config
 
 warnings.filterwarnings("ignore")
@@ -77,39 +73,28 @@ class FashionSearchEngine:
         print("✅ Fashion Search Engine ready!")
 
     def _load_models(self):
-        """Load all required models"""
+        """Load all required models."""
         print("📦 Loading models...")
 
         # Load color model (optional for search in this script).
         self.color_model = None
         color_model_path = getattr(config, "color_model_path", None)
-        if ColorModel is None:
-            print("⚠️ color_model.py not found; continuing without color model.")
-        elif not color_model_path or not Path(color_model_path).exists():
+        if not color_model_path or not Path(color_model_path).exists():
             print("⚠️ color model checkpoint not found; continuing without color model.")
         else:
-            color_checkpoint = torch.load(
-                color_model_path, map_location=self.device, weights_only=True
-            )
-            self.color_model = ColorModel(embed_dim=self.color_dim).to(self.device)
-            self.color_model.load_state_dict(color_checkpoint)
-            self.color_model.eval()
+            try:
+                self.color_model, _ = load_color_model(
+                    color_model_path=config.color_model_path,
+                    device=self.device,
+                )
+            except Exception as e:
+                print(f"⚠️ Failed to load color model: {e}; continuing without it.")
 
         # Load hierarchy model
-        hierarchy_checkpoint = torch.load(
-            config.hierarchy_model_path, map_location=self.device
+        self.hierarchy_model = load_hierarchy_model(
+            hierarchy_model_path=config.hierarchy_model_path,
+            device=self.device,
         )
-        self.hierarchy_classes = hierarchy_checkpoint.get("hierarchy_classes", [])
-        self.hierarchy_model = HierarchyModel(
-            num_hierarchy_classes=len(self.hierarchy_classes),
-            embed_dim=self.hierarchy_dim,
-        ).to(self.device)
-        self.hierarchy_model.load_state_dict(hierarchy_checkpoint["model_state"])
-
-        # Set hierarchy extractor
-        hierarchy_extractor = HierarchyExtractor(self.hierarchy_classes, verbose=False)
-        self.hierarchy_model.set_hierarchy_extractor(hierarchy_extractor)
-        self.hierarchy_model.eval()
 
         # Load main CLIP model (baseline or fine-tuned GAP-CLIP)
         if self.use_baseline:
@@ -409,7 +394,7 @@ if __name__ == "__main__":
 
     if args.queries:
         all_results = {}
-        figures_dir = Path(args.save).parent if args.save else Path("evaluation")
+        figures_dir = Path("evaluation")
         figures_dir.mkdir(parents=True, exist_ok=True)
         (figures_dir / "figures").mkdir(parents=True, exist_ok=True)