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