|
|
|
|
|
""" |
|
|
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() |
|
|
|
