Fix majority voting result mapping to class names

synapse/hardware/cpu.py

@@ -69,8 +69,9 @@ def step(self, program):
             self.running = False
             if self.neural_ip.last_result is not None:
                 result = self.get_reg("$t9")
-                label_map = {0: "A", 1: "B", 2: "Unknown"}
-                print(f"Final classification: {label_map.get(result, result)}")
+                class_names = getattr(self.neural_ip, "class_names", [])
+                label = class_names[result] if 0 <= result < len(class_names) else result
+                print(f"Final classification: {label}")
         elif instr == "ADDI":
             rd = parts[1].rstrip(",")
             rs = parts[2].rstrip(",")
@@ -99,6 +100,6 @@ def step(self, program):
         elif instr == "OP_NEUR":
             subcmd = " ".join(parts[1:])
             self.neural_ip.run_instruction(subcmd, memory=self.memory)
-            if subcmd.upper().startswith("INFER_ANN") and self.neural_ip.last_result is not None:
+            if any(subcmd.upper().startswith(cmd) for cmd in ("INFER_ANN", "GET_ARGMAX", "GET_NUM_CLASSES")) and self.neural_ip.last_result is not None:
                 self.set_reg("$t9", int(self.neural_ip.last_result))
 

synapse/models/redundant_ip.py

@@ -55,6 +55,12 @@ def __init__(self, train_data_dir: str | None = None) -> None:
         # Metrics and figures generated during training keyed by ANN ID
         self.metrics_by_ann: Dict[int, Dict[str, float]] = {}
         self.figures_by_ann: Dict[int, List] = {}
+        # Class names inferred from the training directory for label mapping
+        self.class_names: List[str] = []
+        if train_data_dir:
+            root = Path(train_data_dir)
+            if root.is_dir():
+                self.class_names = sorted([d.name for d in root.iterdir() if d.is_dir()])
 
     # ------------------------------------------------------------------
     # Assembly interface
@@ -84,6 +90,15 @@ def run_instruction(self, subcmd: str, memory=None) -> None:
                     ann.load(f"{prefix}_{ann_id}.pt")
                 except FileNotFoundError:
                     pass
+        elif op == "GET_ARGMAX":
+            # Argmax already computed during INFER_ANN; expose via last_result
+            pass
+        elif op == "GET_NUM_CLASSES":
+            if not self.class_names and self.train_data_dir:
+                root = Path(self.train_data_dir)
+                if root.is_dir():
+                    self.class_names = sorted([d.name for d in root.iterdir() if d.is_dir()])
+            self.last_result = len(self.class_names) if self.class_names else hp.num_classes
         elif op == "SAVE_PROJECT":
             json_path = tokens[1] if len(tokens) > 1 else "project.json"
             prefix = tokens[2] if len(tokens) > 2 else "weights"
@@ -231,6 +246,10 @@ def predict_majority(self, X: np.ndarray):
     # Internal helpers
     # ------------------------------------------------------------------
     def _load_dataset(self):
+        if not self.class_names and self.train_data_dir:
+            root = Path(self.train_data_dir)
+            if root.is_dir():
+                self.class_names = sorted([d.name for d in root.iterdir() if d.is_dir()])
         if self._cached_dataset is None:
             if not self.train_data_dir:
                 print("No training data directory specified; aborting training.")

tests/test_cpu_class_labels.py

@@ -0,0 +1,27 @@
+import os
+import sys
+
+import pytest
+
+sys.path.append(os.getcwd())
+from synapse.models.redundant_ip import RedundantNeuralIP
+from synapse.hardware.cpu import CPU
+
+
+def test_get_num_classes(tmp_path):
+    (tmp_path / "class0").mkdir()
+    (tmp_path / "class1").mkdir()
+    ip = RedundantNeuralIP(train_data_dir=str(tmp_path))
+    ip.run_instruction("GET_NUM_CLASSES")
+    assert ip.last_result == 2
+
+
+def test_cpu_halt_uses_class_names(capsys):
+    ip = RedundantNeuralIP()
+    ip.class_names = ["class0", "class1"]
+    ip.last_result = 1
+    cpu = CPU("CPU1", None, ip, None, None)
+    cpu.set_reg("$t9", 1)
+    cpu.step(["HALT"])
+    captured = capsys.readouterr().out
+    assert "Final classification: class1" in captured