Add -c option to add custom class file (#13)

#12 Support custom class files.

README.md

@@ -9,8 +9,12 @@ sudo pip install -r requirements.txt
 ### Usage
 
 ```bash
-python xmltotxt.py -xml xml -out out
+python xmltotxt.py -c cls.txt -xml xml -out out
 ```
+#### Mandatory arguments
+-xml 
+#### Optional arguments
+-c, -out
 
 ### Example
 

reader.py

@@ -16,7 +16,7 @@ def get_xml_files(self):
         return xml_filenames
 
     @staticmethod
-    def get_classes(filename="classes.txt"):
+    def get_classes(filename):
         with open(os.path.join(os.path.dirname(os.path.realpath('__file__')), filename), "r", encoding="utf8") as f:
             lines = f.readlines()
             return {value: key for (key, value) in enumerate(list(map(lambda x: x.strip(), lines)))}

transformer.py

@@ -5,14 +5,15 @@
 
 
 class Transformer(object):
-    def __init__(self, xml_dir, out_dir):
+    def __init__(self, xml_dir, out_dir, class_file):
         self.xml_dir = xml_dir
         self.out_dir = out_dir
+        self.class_file = class_file
 
     def transform(self):
         reader = Reader(xml_dir=self.xml_dir)
         xml_files = reader.get_xml_files()
-        classes = reader.get_classes()
+        classes = reader.get_classes(self.class_file)
         object_mapper = ObjectMapper()
         annotations = object_mapper.bind_files(xml_files, xml_dir=self.xml_dir)
         self.write_to_txt(annotations, classes)

xmltotxt.py

@@ -10,6 +10,7 @@ def main():
     parser = argparse.ArgumentParser(description="Formatter from ImageNet xml to Darknet text format")
     parser.add_argument("-xml", help="Relative location of xml files directory", required=True)
     parser.add_argument("-out", help="Relative location of output txt files directory", default="out")
+    parser.add_argument("-c", help="Relative path to classes file", default="classes.txt")
     args = parser.parse_args()
 
     xml_dir = os.path.join(os.path.dirname(os.path.realpath('__file__')), args.xml)
@@ -22,10 +23,20 @@ def main():
         os.makedirs(out_dir)
 
     if not os.access(out_dir, os.W_OK):
-        print("%s folder is not writeable.")
+        print("%s folder is not writeable." % out_dir)
         sys.exit()
+
+    class_file = os.path.join(os.path.dirname(os.path.realpath('__file__')), args.c)
 
-    transformer = Transformer(xml_dir=xml_dir, out_dir=out_dir)
+    if not os.access(class_file, os.F_OK):
+        print("%s file is missing." % class_file)
+        sys.exit()
+
+    if not os.access(class_file, os.R_OK):
+        print("%s file is not readable." % class_file)
+        sys.exit()
+
+    transformer = Transformer(xml_dir=xml_dir, out_dir=out_dir, class_file=class_file)
     transformer.transform()