140 lines
3.7 KiB
Python
140 lines
3.7 KiB
Python
import argparse
|
|
import os
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
# ===== Import Project Modules =====
|
|
import evaluate
|
|
import setup
|
|
import train
|
|
import visualize
|
|
|
|
|
|
# ===== Display Banner Function =====
|
|
def display_banner():
|
|
"""Display the project banner in the console."""
|
|
print("\n" + "=" * 60)
|
|
print(" YOLO-TinyPerson: Tiny Person Detection with YOLO Models")
|
|
print("=" * 60)
|
|
print(" A framework for training and evaluating YOLO models")
|
|
print(" on the TinyPerson dataset for small object detection")
|
|
print("=" * 60 + "\n")
|
|
|
|
|
|
# ===== Command Functions =====
|
|
def setup_environment(args):
|
|
"""Set up the environment, datasets, and models."""
|
|
print("\n[*] Setting up environment...")
|
|
setup.main()
|
|
|
|
|
|
def train_models(args):
|
|
"""Train YOLO models on the TinyPerson dataset."""
|
|
print("\n[*] Training models...")
|
|
train.main()
|
|
|
|
|
|
def evaluate_models(args):
|
|
"""Evaluate trained YOLO models on the TinyPerson dataset."""
|
|
print("\n[*] Evaluating models...")
|
|
evaluate.main()
|
|
|
|
|
|
def visualize_results(args):
|
|
"""Create visualizations of training and evaluation results."""
|
|
print("\n[*] Generating visualizations...")
|
|
visualize.main()
|
|
|
|
|
|
def run_all(args):
|
|
"""Run the complete pipeline: setup, train, evaluate, and visualize."""
|
|
print("\n[*] Running complete pipeline...")
|
|
|
|
start_time = time.time()
|
|
|
|
print("\n[1/4] Setting up environment")
|
|
setup.main()
|
|
|
|
print("\n[2/4] Training models")
|
|
train.main()
|
|
|
|
print("\n[3/4] Evaluating models")
|
|
evaluate.main()
|
|
|
|
print("\n[4/4] Generating visualizations")
|
|
visualize.main()
|
|
|
|
# ----- Calculate Total Execution Time -----
|
|
total_duration = time.time() - start_time
|
|
hours, remainder = divmod(total_duration, 3600)
|
|
minutes, seconds = divmod(remainder, 60)
|
|
|
|
print("\n" + "=" * 60)
|
|
print(
|
|
f"[✓] Complete pipeline finished in {int(hours)}h {int(minutes)}m {int(seconds)}s"
|
|
)
|
|
print("=" * 60)
|
|
|
|
|
|
# ===== Command Line Interface =====
|
|
def create_parser():
|
|
"""Create the command line argument parser."""
|
|
parser = argparse.ArgumentParser(
|
|
description="YOLO-TinyPerson: Tiny Person Detection with YOLO Models",
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
)
|
|
|
|
subparsers = parser.add_subparsers(dest="command", help="Command to execute")
|
|
|
|
# ----- Setup Command -----
|
|
setup_parser = subparsers.add_parser(
|
|
"setup", help="Set up environment, download dataset and models"
|
|
)
|
|
setup_parser.set_defaults(func=setup_environment)
|
|
|
|
# ----- Train Command -----
|
|
train_parser = subparsers.add_parser(
|
|
"train", help="Train YOLO models on TinyPerson dataset"
|
|
)
|
|
train_parser.set_defaults(func=train_models)
|
|
|
|
# ----- Evaluate Command -----
|
|
eval_parser = subparsers.add_parser("evaluate", help="Evaluate trained models")
|
|
eval_parser.set_defaults(func=evaluate_models)
|
|
|
|
# ----- Visualize Command -----
|
|
viz_parser = subparsers.add_parser(
|
|
"visualize", help="Generate visualizations of results"
|
|
)
|
|
viz_parser.set_defaults(func=visualize_results)
|
|
|
|
# ----- All Command -----
|
|
all_parser = subparsers.add_parser("all", help="Run complete pipeline")
|
|
all_parser.set_defaults(func=run_all)
|
|
|
|
return parser
|
|
|
|
|
|
# ===== Main Function =====
|
|
def main():
|
|
"""Main entry point for the application."""
|
|
display_banner()
|
|
|
|
parser = create_parser()
|
|
args = parser.parse_args()
|
|
|
|
# ----- Check for Command (use 'all' as default) -----
|
|
if not args.command:
|
|
print("[i] No command specified. Running complete pipeline by default.\n")
|
|
run_all(args)
|
|
elif hasattr(args, "func"):
|
|
args.func(args)
|
|
else:
|
|
parser.print_help()
|
|
|
|
|
|
# ===== Script Entry Point =====
|
|
if __name__ == "__main__":
|
|
main()
|