Update game/game.py

game/game.py

@@ -173,6 +173,28 @@ def load_exported_model(model_dir, model_type):
         session = ort.InferenceSession(model_path, providers=['CPUExecutionProvider'])
     
     # Create a wrapper class to mimic the PyTorch model interface
+    class ONNXModelWrapper:
+        """
+        - predict(image_path=..., threshold=...) -> {"refined_probabilities": np.ndarray[N, C]}
+        - dataset.get_tag_info(idx) -> (tag, category)
+        - Keeps your signature compatible, ignores 'threshold' in the wrapper.
+        """
+        def __init__(self, session, metadata: dict):
+            self.session = session
+            self.metadata = metadata or {}
+            dataset_info = self.metadata.get("dataset_info", {})
+            self.total_tags = dataset_info.get("total_tags", 0)
+
+            # idx <-> tag mapping
+            tag_mapping = dataset_info.get("tag_mapping", {})
+            if "idx_to_tag" in tag_mapping:
+                self.idx_to_tag = {int(k): v for k, v in tag_mapping["idx_to_tag"].items()}
+            elif "tag_to_idx" in tag_mapping:
+                t2i = tag_mapping["tag_to_idx"]
+                self.idx_to_tag = {v: k for k, v in t2i.items()}
+            else:
+                self.idx_to_tag = {}
+
     class ONNXModelWrapper:
         """
         - predict(image_path=..., threshold=...) -> {"refined_probabilities": np.ndarray[N, C]}
@@ -196,8 +218,13 @@ def load_exported_model(model_dir, model_type):
                 self.idx_to_tag = {}
 
             self.tag_to_category = self.metadata.get("tag_to_category", {})
-            if not self.tag_to_category and "tag_to_category" in dataset_info:
-                self.tag_to_category = dataset_info["tag_to_category"]
+            if not self.tag_to_category:
+                # Try to get from dataset_info.tag_mapping.tag_to_category (correct path)
+                if "tag_mapping" in dataset_info and "tag_to_category" in dataset_info["tag_mapping"]:
+                    self.tag_to_category = dataset_info["tag_mapping"]["tag_to_category"]
+                # Fallback to direct path in case structure varies
+                elif "tag_to_category" in dataset_info:
+                    self.tag_to_category = dataset_info["tag_to_category"]
 
             # Compatibility shim for scan_handler: model.dataset.get_tag_info(...)
             self.dataset = self