modify pick_xml_img_by_xml.py

pick_xml_img_by_xml.py

@@ -6,7 +6,7 @@
 # This file should put on the same path with the folder
 # input: python pick_all_xml_img.py "/home/andy/data/img"
 # output:
-# 	./pickedLabel
+# 	./pickedXml
 # 	./pickedImg
 
 import os
@@ -17,61 +17,70 @@
 
 def parse_args():
     parser = argparse.ArgumentParser()
-    parser.add_argument('srcdir', help='file directory', type=str)
+    parser.add_argument('xmldir', help='xml file directory', type=str)
+    parser.add_argument('imgdir', help='image file directory', type=str)
 
     args = parser.parse_args()
     return args
 
 
 
-def pick(srcdir):
-    pickedLabel_dir = "./pickedLabel"
-    pickedImg_dir = "./pickedImg"
-    srcdir = os.path.abspath(srcdir)
-    if srcdir[-1] == "/":
-        srcdir = srcdir[:-1]
+def pick(xmldir, imgdir):
+    xmldir = os.path.abspath(xmldir)
+    if xmldir[-1] == "/":
+        xmldir = xmldir[:-1]
+    imgdir = os.path.abspath(imgdir)
+    if imgdir[-1] == "/":
+        imgdir = imgdir[:-1]
 
+    pickedXml_dir = "./pickedXml"
+    pickedImg_dir = "./pickedImg"
     # recreate folder
-    if os.path.exists(pickedLabel_dir):
-        shutil.rmtree(pickedLabel_dir)
-    os.makedirs(pickedLabel_dir) 
+    if os.path.exists(pickedXml_dir):
+        shutil.rmtree(pickedXml_dir)
+    os.makedirs(pickedXml_dir) 
 
     if os.path.exists(pickedImg_dir):
         shutil.rmtree(pickedImg_dir)
     os.makedirs(pickedImg_dir)
 
-    filelist = os.listdir(srcdir)
+    xmllist = os.listdir(xmldir)
 
-    for file in filelist:
+    for file in xmllist:
         print(file + "-->start!")
         fileInfor = file.split(".")
         if fileInfor[-1] != "xml":
             continue
 
-        src_img = srcdir + "/" + fileInfor[0] + ".jpg"
-        dst_img = pickedImg_dir + "/" + fileInfor[0] + ".jpg"
+        src_img = os.path.join(imgdir, fileInfor[0] + ".jpg")
+        dst_img = os.path.join(pickedImg_dir, fileInfor[0] + ".jpg")
         if not os.path.exists(src_img):
-            src_img = srcdir + "/" + fileInfor[0] + ".png"
-            dst_img = pickedImg_dir + "/" + fileInfor[0] + ".png"
+            src_img = os.path.join(imgdir, fileInfor[0] + ".png")
+            dst_img = os.path.join(pickedImg_dir, fileInfor[0] + ".png")
         if not os.path.exists(src_img):
+            print(os.path.basename(src_img), " is not exits!")
             continue
 
-        src_label = srcdir + "/" + fileInfor[0] + ".xml"
-        dst_label = pickedLabel_dir + "/" + fileInfor[0] + ".xml"
+        src_xml = os.path.join(xmldir, fileInfor[0] + ".xml")
+        dst_xml = os.path.join(pickedXml_dir, fileInfor[0] + ".xml")
 
-        shutil.copyfile(src_label, dst_label)
+        shutil.copyfile(src_xml, dst_xml)
         shutil.copyfile(src_img, dst_img)
-    print("Path of picked labels = ",os.path.abspath(pickedLabel_dir))
+    print("Path of picked xmls = ",os.path.abspath(pickedXml_dir))
     print("Path of picked images = ",os.path.abspath(pickedImg_dir))
 
 if __name__ == '__main__':
     print("Pick images and xml by xml")
     args = parse_args()
-    srcdir = args.srcdir
-
-    if not os.path.exists(srcdir):
-        print("Error !!! %s is not exists, please check the parameter"%srcdir)
+    xmldir = args.xmldir
+    imgdir = args.imgdir
+
+    if not os.path.exists(xmldir):
+        print("Error !!! %s is not exists, please check the parameter"%xmldir)
+        sys.exit(0)
+    if not os.path.exists(imgdir):
+        print("Error !!! %s is not exists, please check the parameter"%imgdir)
         sys.exit(0)
 
-    pick(srcdir)
+    pick(xmldir, imgdir)
     print("Done!")