2026001

classes/patient_identifier.py

@@ -53,8 +53,8 @@ class PatientIdentifier:
         self.labels_ = labels
         groups = self._run_clustering(
             labels, filenames,
-            data_loader_fn=lambda class_name, files: self._load_thumbnails(
-                base_path, class_name, files),
+            data_loader_fn=lambda class_name, idx: self._load_thumbnails(
+                base_path, class_name, filenames[idx]),
             valid_extensions=valid_extensions,
         )
         self.groups_ = groups
@@ -82,21 +82,28 @@ class PatientIdentifier:
         self.groups_ = groups
         return groups
 
-    def build_assignment_dict(self, groups, filenames, labels):
+    def build_assignment_dict(self, groups, filenames, labels,
+                              class_prefixes=None):
         """Convert flat group array into a human-readable dict.
 
+        Parameters
+        ----------
+        class_prefixes : dict or None
+            Optional mapping from class_name -> short prefix used in the
+            patient IDs (e.g. {"Bengin cases": "benign"}). If None, a prefix
+            is derived automatically from each class name, so the method works
+            for any dataset rather than only IQ-OTH's three classes.
+
         Returns
         -------
         dict mapping patient_id (e.g. "benign_0") -> list of image filenames.
         """
         self.labels_ = labels
-        class_map = {
-            "Bengin cases":    "benign",
-            "Malignant cases": "malig",
-            "Normal cases":    "normal",
-        }
+        if class_prefixes is None:
+            class_prefixes = {c: self._slugify(c)
+                              for c in self.patient_estimates}
         assignment = {}
-        for class_name, prefix in class_map.items():
+        for class_name, prefix in class_prefixes.items():
             idx_class = np.where(labels == class_name)[0]
             # Get unique cluster IDs within this class (preserve order)
             seen = []
@@ -115,6 +122,12 @@ class PatientIdentifier:
         self.assignments_ = assignment
         return assignment
 
+    @staticmethod
+    def _slugify(class_name):
+        """Turn a class name into a short lowercase prefix for patient IDs."""
+        token = str(class_name).strip().lower().split()[0]
+        return "".join(ch for ch in token if ch.isalnum()) or "group"
+
     def get_unique_groups(self):
         if self.groups_ is None:
             raise RuntimeError("Call identify_*() first.")