update voc_annotation.py

bubbliiiing · May 4, 2022 · 7d74410 · 7d74410
1 parent 8bc7a96
commit 7d74410
Showing 1 changed file with 35 additions and 25 deletions.
diff --git a/voc_annotation.py b/voc_annotation.py
@@ -39,8 +39,9 @@
 
 #-------------------------------------------------------#
 #   统计目标数量
-#-------------------------------------------------------#,
-nums = np.zeros(len(classes))
+#-------------------------------------------------------#
+photo_nums  = np.zeros(len(VOCdevkit_sets))
+nums        = np.zeros(len(classes))
 def convert_annotation(year, image_id, list_file):
     in_file = open(os.path.join(VOCdevkit_path, 'VOC%s/Annotations/%s.xml'%(year, image_id)), encoding='utf-8')
     tree=ET.parse(in_file)
@@ -62,6 +63,9 @@ def convert_annotation(year, image_id, list_file):
 
 if __name__ == "__main__":
     random.seed(0)
+    if " " in os.path.abspath(VOCdevkit_path):
+        raise ValueError("数据集存放的文件夹路径与图片名称中不可以存在空格，否则会影响正常的模型训练，请注意修改。")
+
     if annotation_mode == 0 or annotation_mode == 1:
         print("Generate txt in ImageSets.")
         xmlfilepath     = os.path.join(VOCdevkit_path, 'VOC2007/Annotations')
@@ -105,6 +109,7 @@ def convert_annotation(year, image_id, list_file):
 
     if annotation_mode == 0 or annotation_mode == 2:
         print("Generate 2007_train.txt and 2007_val.txt for train.")
+        type_index = 0
         for year, image_set in VOCdevkit_sets:
             image_ids = open(os.path.join(VOCdevkit_path, 'VOC%s/ImageSets/Main/%s.txt'%(year, image_set)), encoding='utf-8').read().strip().split()
             list_file = open('%s_%s.txt'%(year, image_set), 'w', encoding='utf-8')
@@ -113,31 +118,36 @@ def convert_annotation(year, image_id, list_file):
 
                 convert_annotation(year, image_id, list_file)
                 list_file.write('\n')
+            photo_nums[type_index] = len(image_ids)
+            type_index += 1
             list_file.close()
         print("Generate 2007_train.txt and 2007_val.txt for train done.")
 
-    def printTable(List1, List2):
-        for i in range(len(List1[0])):
-            print("|", end=' ')
-            for j in range(len(List1)):
-                print(List1[j][i].rjust(int(List2[j])), end=' ')
+        def printTable(List1, List2):
+            for i in range(len(List1[0])):
                 print("|", end=' ')
-            print()
+                for j in range(len(List1)):
+                    print(List1[j][i].rjust(int(List2[j])), end=' ')
+                    print("|", end=' ')
+                print()
 
-    str_nums = [str(int(x)) for x in nums]
-    tableData = [
-        classes, str_nums
-    ]
-    colWidths = [0]*len(tableData)
-    len1 = 0
-    for i in range(len(tableData)):
-        for j in range(len(tableData[i])):
-            if len(tableData[i][j]) > colWidths[i]:
-                colWidths[i] = len(tableData[i][j])
-    printTable(tableData, colWidths)
-
-    if np.sum(nums) == 0:
-        print("在数据集中并未获得任何目标，请注意修改classes_path对应自己的数据集，并且保证标签名字正确，否则训练将会没有任何效果！")
-        print("在数据集中并未获得任何目标，请注意修改classes_path对应自己的数据集，并且保证标签名字正确，否则训练将会没有任何效果！")
-        print("在数据集中并未获得任何目标，请注意修改classes_path对应自己的数据集，并且保证标签名字正确，否则训练将会没有任何效果！")
-        print("（重要的事情说三遍）。")
+        str_nums = [str(int(x)) for x in nums]
+        tableData = [
+            classes, str_nums
+        ]
+        colWidths = [0]*len(tableData)
+        len1 = 0
+        for i in range(len(tableData)):
+            for j in range(len(tableData[i])):
+                if len(tableData[i][j]) > colWidths[i]:
+                    colWidths[i] = len(tableData[i][j])
+        printTable(tableData, colWidths)
+
+        if photo_nums[0] <= 500:
+            print("训练集数量小于500，属于较小的数据量，请注意设置较大的训练世代（Epoch）以满足足够的梯度下降次数（Step）。")
+
+        if np.sum(nums) == 0:
+            print("在数据集中并未获得任何目标，请注意修改classes_path对应自己的数据集，并且保证标签名字正确，否则训练将会没有任何效果！")
+            print("在数据集中并未获得任何目标，请注意修改classes_path对应自己的数据集，并且保证标签名字正确，否则训练将会没有任何效果！")
+            print("在数据集中并未获得任何目标，请注意修改classes_path对应自己的数据集，并且保证标签名字正确，否则训练将会没有任何效果！")
+            print("（重要的事情说三遍）。")