Encoder

LabelEncoder.py

@@ -0,0 +1,46 @@
+import numpy as np
+import json
+
+class LabelEncoder(object):
+    """Label encoder for tag labels."""
+    def __init__(self, class_to_index={}):
+        self.class_to_index = class_to_index
+        self.index_to_class = {v: k for k, v in self.class_to_index.items()}
+        self.classes = list(self.class_to_index.keys())
+
+    def __len__(self):
+        return len(self.class_to_index)
+
+    def __str__(self):
+        return f"<LabelEncoder(num_classes={len(self)})>"
+
+    def fit(self, y):
+        classes = np.unique(y)
+        for i, class_ in enumerate(classes):
+            self.class_to_index[class_] = i
+        self.index_to_class = {v: k for k, v in self.class_to_index.items()}
+        self.classes = list(self.class_to_index.keys())
+        return self
+
+    def encode(self, y):
+        encoded = np.zeros((len(y)), dtype=int)
+        for i, item in enumerate(y):
+            encoded[i] = self.class_to_index[item]
+        return encoded
+
+    def decode(self, y):
+        classes = []
+        for i, item in enumerate(y):
+            classes.append(self.index_to_class[item])
+        return classes
+
+    def save(self, fp):
+        with open(fp, "w") as fp:
+            contents = {'class_to_index': self.class_to_index}
+            json.dump(contents, fp, indent=4, sort_keys=False)
+
+    @classmethod
+    def load(cls, fp):
+        with open(fp, "r") as fp:
+            kwargs = json.load(fp=fp)
+        return cls(**kwargs)