From d9d8711e4df6c9e976fe23e669bc208b56dd5028 Mon Sep 17 00:00:00 2001 From: Shiroki Satsuki Date: Thu, 7 Apr 2022 17:51:56 +0800 Subject: [PATCH] Add IndexWrapper for yolo_io --- libs/utils.py | 12 ++++++++++++ libs/yolo_io.py | 2 ++ 2 files changed, 14 insertions(+) diff --git a/libs/utils.py b/libs/utils.py index d511c151b..2a41730ff 100644 --- a/libs/utils.py +++ b/libs/utils.py @@ -115,3 +115,15 @@ def trimmed(text): else: def trimmed(text): return text.trimmed() + +class IndexWrapper: + def __init__(self, content) -> None: + self.content = content + + def __gettitem__(self, index): + try: + result = self.content[index] + except IndexError: + return str(index) + else: + return result \ No newline at end of file diff --git a/libs/yolo_io.py b/libs/yolo_io.py index 192e2c785..ea2227a20 100644 --- a/libs/yolo_io.py +++ b/libs/yolo_io.py @@ -2,6 +2,7 @@ # -*- coding: utf8 -*- import codecs import os +from utils import IndexWrapper from libs.constants import DEFAULT_ENCODING @@ -96,6 +97,7 @@ def __init__(self, file_path, image, class_list_path=None): classes_file = open(self.class_list_path, 'r') self.classes = classes_file.read().strip('\n').split('\n') + self.classes = IndexWrapper(self.classes) # print (self.classes)