UltralyticsAssistant 5 months ago
parent acec22b1e4
commit b1263cf4d5
  1. 14
      tests/test_cli.py
  2. 12
      ultralytics/cfg/__init__.py

@ -117,9 +117,7 @@ def run_yolo_info(args) -> tuple[str, str]:
def test_yolo_info_noargs():
"""Test yolo info command without arguments, same as `yolo checks`."""
stdout, stderr = run_yolo_info("")
checks_out = subprocess.run(
f"yolo checks", shell=True, capture_output=True, text=True, encoding="utf-8"
).stdout
checks_out = subprocess.run("yolo checks", shell=True, capture_output=True, text=True, encoding="utf-8").stdout
assert stdout == checks_out
assert stderr == ""
@ -135,7 +133,7 @@ def test_yolo_info_with_model():
def test_yolo_info_invalid_model():
"""Test yolo info command with an invalid model path."""
stdout, _ = run_yolo_info("model=not-a-model.pt")
assert f"Model 'not-a-model.pt' not found" in stdout
assert "Model 'not-a-model.pt' not found" in stdout
def test_yolo_info_with_show():
@ -158,12 +156,14 @@ def test_train_gpu(task, model, data):
@pytest.mark.slow
@pytest.mark.parametrize("model", {"yolo11n.pt", "yolov8n-seg.pt", "FastSAM-s.pt", "rtdetr-l.pt", "sam_b.pt", "yolo_nas_s.pt"})
def test_yolo_info_various_models(model:str):
@pytest.mark.parametrize(
"model", {"yolo11n.pt", "yolov8n-seg.pt", "FastSAM-s.pt", "rtdetr-l.pt", "sam_b.pt", "yolo_nas_s.pt"}
)
def test_yolo_info_various_models(model: str):
"""Test yolo info command with various models."""
model_path = WEIGHTS_DIR / model
if not model_path.exists():
run(f"yolo predict {model}") # download model before test
stdout, stderr = run_yolo_info(f"model={model_path.as_posix()}")
assert f"Model File: {model_path.as_posix().replace('/', os.sep)}" in stdout
assert stderr == ""
assert stderr == ""

@ -548,7 +548,7 @@ def handle_yolo_settings(args: List[str]) -> None:
def handle_yolo_info(args: List[str]) -> None:
"""Handles YOLO info command-line interface (CLI) commands."""
from ultralytics import FastSAM, NAS, RTDETR, SAM, YOLO
from ultralytics import NAS, RTDETR, SAM, YOLO, FastSAM
nl = "\n"
delim = "-" * 80
@ -559,7 +559,7 @@ def handle_yolo_info(args: List[str]) -> None:
"mobile_sam": SAM,
"sam_": SAM,
"sam2_": SAM,
"yolo_nas": NAS
"yolo_nas": NAS,
}
skip = {"model", "ema", "optimizer", "train_results"}
@ -571,14 +571,14 @@ def handle_yolo_info(args: List[str]) -> None:
"""Yields a formatted string representation of classes from a dictionary."""
for k, v in d.items():
yield f"\t{k}: {_obj_or_dict(v)}"
def _args_load(a: List[str]) -> tuple[dict, str]:
"""Loads arguments from a list of strings."""
f = a.pop(0) if "=" not in a[0] else ""
d:dict = dict(DEFAULT_CFG) | dict(parse_key_value_pair(e) for e in args)
d: dict = dict(DEFAULT_CFG) | dict(parse_key_value_pair(e) for e in args)
check_dict_alignment(dict(DEFAULT_CFG), d)
return d, f
try:
if args:
new, file = _args_load(args)
@ -592,7 +592,7 @@ def handle_yolo_info(args: List[str]) -> None:
f'{colorstr("blue", "bold", f"{nl}Ultralytics Model Info:")}'
f"{nl}{delim}{nl}"
f"Model File: {model.ckpt_path or file}"
)
)
# Display model metadata
f"{[LOGGER.info(f'{str(k).capitalize()}: {_obj_or_dict(v)}') for k,v in meta.items() if k.lower() not in skip]}"
# Display class-names & total count

Loading…
Cancel
Save