| |
| """ |
| Water Surface Segmentation Evaluation Script |
| Evaluate the trained model on a validation dataset. |
| """ |
|
|
| import argparse |
| import os |
| import sys |
| from pathlib import Path |
| from ultralytics import YOLO |
|
|
|
|
| def parse_arguments() -> argparse.Namespace: |
| """Parse command line arguments.""" |
| parser = argparse.ArgumentParser( |
| description="Evaluate water surface segmentation model", |
| formatter_class=argparse.ArgumentDefaultsHelpFormatter |
| ) |
|
|
| parser.add_argument( |
| "--data", |
| type=str, |
| required=True, |
| help="Path to validation dataset or data.yaml file" |
| ) |
|
|
| parser.add_argument( |
| "--weights", |
| type=str, |
| default="model/nwsd-v2.pt", |
| help="Path to model weights file" |
| ) |
|
|
| parser.add_argument( |
| "--img", |
| type=int, |
| default=640, |
| help="Image size for evaluation" |
| ) |
|
|
| parser.add_argument( |
| "--batch", |
| type=int, |
| default=16, |
| help="Batch size for evaluation" |
| ) |
|
|
| parser.add_argument( |
| "--conf", |
| type=float, |
| default=0.25, |
| help="Confidence threshold" |
| ) |
|
|
| parser.add_argument( |
| "--iou", |
| type=float, |
| default=0.45, |
| help="IoU threshold for NMS" |
| ) |
|
|
| parser.add_argument( |
| "--device", |
| type=str, |
| default="", |
| help="Device to use for evaluation (cpu, cuda, mps)" |
| ) |
|
|
| parser.add_argument( |
| "--project", |
| type=str, |
| default="runs/segment", |
| help="Project directory for results" |
| ) |
|
|
| parser.add_argument( |
| "--name", |
| type=str, |
| default="nwsd_eval", |
| help="Experiment name" |
| ) |
|
|
| parser.add_argument( |
| "--save-json", |
| action="store_true", |
| help="Save results in JSON format" |
| ) |
|
|
| parser.add_argument( |
| "--save-txt", |
| action="store_true", |
| help="Save results in TXT format" |
| ) |
|
|
| parser.add_argument( |
| "--plots", |
| action="store_true", |
| help="Generate evaluation plots" |
| ) |
|
|
| return parser.parse_args() |
|
|
|
|
| def validate_inputs(args: argparse.Namespace) -> None: |
| """Validate input arguments.""" |
| if not os.path.exists(args.data): |
| raise FileNotFoundError(f"Data path not found: {args.data}") |
|
|
| if not os.path.exists(args.weights): |
| raise FileNotFoundError(f"Model weights not found: {args.weights}") |
|
|
|
|
| def main(): |
| """Main evaluation function.""" |
| args = parse_arguments() |
|
|
| try: |
| validate_inputs(args) |
|
|
| print(f"Loading model: {args.weights}") |
| model = YOLO(args.weights) |
|
|
| eval_params = { |
| 'data': args.data, |
| 'imgsz': args.img, |
| 'batch': args.batch, |
| 'conf': args.conf, |
| 'iou': args.iou, |
| 'device': args.device, |
| 'project': args.project, |
| 'name': args.name, |
| 'save_json': args.save_json, |
| 'save_txt': args.save_txt, |
| 'plots': args.plots, |
| 'verbose': True, |
| } |
|
|
| print("Starting evaluation with parameters:") |
| for key, value in eval_params.items(): |
| print(f" {key}: {value}") |
|
|
| results = model.val(**eval_params) |
|
|
| print("\n" + "="*50) |
| print("EVALUATION RESULTS SUMMARY") |
| print("="*50) |
|
|
| if hasattr(results, 'box') and results.box is not None: |
| print(f"mAP50: {results.box.map50:.4f}") |
| print(f"mAP50-95: {results.box.map:.4f}") |
|
|
| if hasattr(results, 'seg') and results.seg is not None: |
| print(f"Segmentation mAP50: {results.seg.map50:.4f}") |
| print(f"Segmentation mAP50-95: {results.seg.map:.4f}") |
|
|
| print("\nEvaluation completed successfully!") |
|
|
| except Exception as e: |
| print(f"Error: {str(e)}", file=sys.stderr) |
| sys.exit(1) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
