first commit

small_object/aitod_getclass.py

@@ -0,0 +1,57 @@
+"""
+author:Hatcher
+email:hcwang@semi.ac.cn
+use:将aitod小目标检测数据集按照类别分开
+input：
+voc
+    train_images
+    train_annos
+    ....
+    test&val
+output:
+voc_class
+    classname1
+        images
+        annos
+    classname2
+        iamges
+        annos
+        .....
+
+"""
+import os
+from tqdm import tqdm
+import  xml.dom.minidom
+import shutil
+source_path_root = "/data4/Hatcher/smallobject/voc/test_annos/"
+det_path_root = "/data4/Hatcher/smallobject/voc_class/"
+class_names  = ["airplane"]#os.listdir(det_path_root)
+for class_name in class_names:
+    for source_path_anno in tqdm(os.listdir(source_path_root)):
+        source_path = os.path.join(source_path_root,source_path_anno)
+        dom = xml.dom.minidom.parse(source_path)
+        root = dom.documentElement
+        objects=dom.getElementsByTagName("object")
+        for object in objects:
+            bndbox = object.getElementsByTagName('bndbox')[0]
+            xmin = bndbox.getElementsByTagName('xmin')[0]
+            ymin = bndbox.getElementsByTagName('ymin')[0]
+            xmax = bndbox.getElementsByTagName('xmax')[0]
+            ymax = bndbox.getElementsByTagName('ymax')[0]
+            #print(xmin.childNodes[0].data)
+            #print(type(xmin.childNodes[0].data))
+            xmin_data=int(xmin.childNodes[0].data)
+            ymin_data=int(ymin.childNodes[0].data)
+            xmax_data=int(xmax.childNodes[0].data)
+            ymax_data=int(ymax.childNodes[0].data)
+            label_name=object.getElementsByTagName('name')[0].childNodes[0].data
+            if not (label_name==class_name):
+                 dom.documentElement.removeChild(object)
+        objects=dom.getElementsByTagName("object")
+        if not (len(objects)==0):
+            #存储xml，复制图像
+            xml_path = os.path.join(det_path_root,class_name,"annos",source_path_anno)
+            fp = open(xml_path, 'w', encoding='utf-8')
+            dom.writexml(fp, indent='', addindent='\t', newl='\n', encoding='utf-8')
+            fp.close()
+            shutil.copy(source_path.replace("test_annos","test_images").replace(".xml",".png"),os.path.join(det_path_root,class_name,"images",source_path_anno.replace(".xml",".png")))