diff --git a/main.py b/main.py index 49ad3a8..92ff1a3 100755 --- a/main.py +++ b/main.py @@ -1,43 +1,53 @@ -#!/bin/python3.10 -import sys +#!/usr/bin/env python3 +import argparse +import importlib from detector.paths import create_dirs -ARGS = """ -Usage: main.py [OPTION] -Options: - -h --help Displays this list - -e --extract Excracts and cropps traffic light images from given images in ./assets/exctraction/input/ to ./assets/exctraction/cropped/ - -t --train Trains model - -i --image Detecs traffic lights in images located in ./assets/detection/images_in/ - -v --video Detecs traffic lights in videos located in ./assets/detection/videos_in/ -""" +parser = argparse.ArgumentParser(description="Traffic light detection script.") +parser.add_argument( + "-e", + "--extract", + action="store_true", + help="extracts and crops traffic light images from given images in ./assets/extraction/input/ to ./assets/extraction/cropped/", +) +parser.add_argument( + "-t", + "--train", + action="store_true", + help="trains model.", +) +parser.add_argument( + "-i", + "--image", + action="store_true", + help="detects traffic lights in images located in ./assets/detection/images_in/", +) +parser.add_argument( + "-v", + "--video", + action="store_true", + help="detects traffic lights in videos located in ./assets/detection/videos_in/", +) -def main(argv) -> None: +def main(args) -> None: create_dirs() - for arg in argv: - if arg in ("-h", "--help"): - print(ARGS) - sys.exit() - elif arg in ("-e", "--extract"): - from detector.extract_traffic_lights import extract_traffic_lights - extract_traffic_lights() - elif arg in ("-t", "--train"): - from detector.train_traffic_light_color import train_traffic_light_color - train_traffic_light_color() - elif arg in ("-i", "--image"): - from detector.detect_traffic_light_color_image import ( - detect_traffic_light_color_image, - ) - detect_traffic_light_color_image() - elif arg in ("-v", "--video"): - from detector.detect_traffic_light_color_video import ( - detect_traffic_light_color_video, - ) - detect_traffic_light_color_video() + if args.extract: + module = importlib.import_module("detector.extract_traffic_lights") + module.extract_traffic_lights() + if args.train: + module = importlib.import_module("detector.train_traffic_light_color") + module.train_traffic_light_color() + if args.image: + module = importlib.import_module("detector.detect_traffic_light_color_image") + module.detect_traffic_light_color_image() + if args.video: + module = importlib.import_module("detector.detect_traffic_light_color_video") + module.detect_traffic_light_color_image() if __name__ == "__main__": - main(sys.argv[1:]) + args = parser.parse_args() + main(args)