diff --git a/frigate/util.py b/frigate/util.py index f11c0b0f9b..f26ddd5549 100755 --- a/frigate/util.py +++ b/frigate/util.py @@ -604,6 +604,24 @@ def add_mask(mask, mask_img): ) cv2.fillPoly(mask_img, pts=[contour], color=(0)) +def load_labels(path, encoding="utf-8"): + """Loads labels from file (with or without index numbers). + Args: + path: path to label file. + encoding: label file encoding. + Returns: + Dictionary mapping indices to labels. + """ + with open(path, "r", encoding=encoding) as f: + lines = f.readlines() + if not lines: + return {} + + if lines[0].split(" ", maxsplit=1)[0].isdigit(): + pairs = [line.split(" ", maxsplit=1) for line in lines] + return {int(index): label.strip() for index, label in pairs} + else: + return {index: line.strip() for index, line in enumerate(lines)} def load_labels(path, encoding="utf-8"): """Loads labels from file (with or without index numbers).