test_model.py

"""
Quick test runner for the trained model.
Runs a full evaluation (on test split) and a single-sample inference using the model path from `config.py`.
Usage:
    python test_model.py        # runs evaluation and a sample inference
    python test_model.py --model path/to/model.pt  # evaluate a specific model
"""

import os
import argparse
from pathlib import Path

from config import Config
from evaluate_model import evaluate_model, test_on_image

OUTPUT_DIR = Path("./evaluation_results")
OUTPUT_DIR.mkdir(exist_ok=True)


def find_sample_image():
    test_img_dir = Path("./datasets/weapons/test/images")
    if not test_img_dir.exists():
        return None
    imgs = [p for p in test_img_dir.iterdir() if p.suffix.lower() in ('.jpg', '.jpeg', '.png')]
    return imgs[0] if imgs else None


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', type=str, help='Path to model to test')
    parser.add_argument('--conf', type=float, default=0.5, help='Confidence threshold for sample inference')
    args = parser.parse_args()

    model_path = args.model or Config.MODEL_PATH

    print(f"Using model: {model_path}")
    if not os.path.exists(model_path):
        print(f"Error: model not found at {model_path}")
        return

    # Evaluate (full test split)
    print("\n== Running evaluation on test split ==")
    metrics = evaluate_model(model_path)
    print("\nEvaluation metrics:")
    print(metrics)

    # Sample inference
    sample_img = find_sample_image()
    if sample_img:
        print(f"\n== Running sample inference on: {sample_img} ==")
        test_on_image(model_path, str(sample_img), conf_threshold=args.conf)
    else:
        print("No sample image found to run inference on.")


if __name__ == '__main__':
    main()