| 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) |
