diff --git a/README.md b/README.md index 277ed06..d1da67d 100644 --- a/README.md +++ b/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 diff --git a/reader.py b/reader.py index 96f848a..f072b2c 100644 --- a/reader.py +++ b/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)))} diff --git a/transformer.py b/transformer.py index a9f1cea..9ab13b6 100644 --- a/transformer.py +++ b/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) diff --git a/xmltotxt.py b/xmltotxt.py index 275ac31..ca7f480 100644 --- a/xmltotxt.py +++ b/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()