kristoferssolo/Traffic-Light-Detector
1
0
Fork 0
mirror of https://github.com/kristoferssolo/Traffic-Light-Detector.git synced 2025-10-21 20:00:36 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Updated main

Browse Source
This commit is contained in:
Kristofers Solo 2022-12-17 16:12:18 +02:00
parent 7af7bd8748
commit a8220d9a6c
1 changed files with 14 additions and 35 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

49
main.py
View File

@ -1,59 +1,38 @@
#!/usr/bin/env python3 #!/usr/bin/env python3
import argparse import argparse
import importlib
from detector.paths import BASE_PATH, create_dirs, LOGS_PATH
from loguru import logger from loguru import logger
log_level = "DEBUG" if BASE_PATH.joinpath("debug").exists() else "INFO" from paths import create_dirs, IMAGES_IN_PATH
# Set up logging
logger.add(LOGS_PATH.joinpath("detection.log"), format="{time} | {level} | {message}", level=log_level, rotation="1 MB", compression="zip")
parser = argparse.ArgumentParser(description="Traffic light detection script.") parser = argparse.ArgumentParser(description="Traffic light detection script.")
parser.add_argument( parser.add_argument(
"-e", "-w",
"--extract", "--webcam",
action="store_true", action="store_true",
help="extracts and crops traffic light images from given images in ./assets/extraction/input/ to ./assets/extraction/cropped/", help="reads webcam inputs to determine traffic light color",
)
parser.add_argument(
"-t",
"--train",
action="store_true",
help="trains model.",
) )
parser.add_argument( parser.add_argument(
"-i", "-i",
"--image", "--image",
action="store_true", action="store_true",
help="detects traffic lights in images located in ./assets/detection/images_in/", help="detects traffic lights in images located in ./assets/images_in/",
) )
# parser.add_argument(
# "-v",
# "--video",
# action="store_true",
# help="detects traffic lights in videos located in ./assets/detection/videos_in/",
# )
@logger.catch @logger.catch
def main(args) -> None: def main(args) -> None:
create_dirs() create_dirs()
if args.extract: if args.webcam:
module = importlib.import_module("detector.extract_traffic_lights") from TrafficLightDetector.traffic_light_webcam import TrafficLightDetectorWebcam
module.extract_traffic_lights() camera = TrafficLightDetectorWebcam()
if args.train: camera.enable()
module = importlib.import_module("detector.train_traffic_light_color")
module.train_traffic_light_color()
if args.image: if args.image:
module = importlib.import_module("detector.detect_traffic_light_color_image") from TrafficLightDetector.traffic_light_images import TrafficLightDetectorImages
module.detect_traffic_light_color_image() for path in IMAGES_IN_PATH.iterdir():
# if args.video: image = TrafficLightDetectorImages(path)
# module = importlib.import_module("detector.detect_traffic_light_color_video") image.draw()
# module.detect_traffic_light_color_video()
if __name__ == "__main__": if __name__ == "__main__":