Add IndexWrapper for yolo_io

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

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)