From 82e8daef7200d1f8884a9abe947fbcc9843288aa Mon Sep 17 00:00:00 2001 From: Burhan <62214284+Burhan-Q@users.noreply.github.com> Date: Thu, 14 Mar 2024 16:02:28 -0400 Subject: [PATCH] New test for labels and crops (#8861) Co-authored-by: Glenn Jocher --- tests/test_python.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/tests/test_python.py b/tests/test_python.py index c096a38c19..2a72b417cb 100644 --- a/tests/test_python.py +++ b/tests/test_python.py @@ -334,6 +334,28 @@ def test_results(): print(r, len(r), r.path) +def test_labels_and_crops(): + """Test output from prediction args for saving detection labels and crops.""" + imgs = [SOURCE, ASSETS / "zidane.jpg"] + results = YOLO(WEIGHTS_DIR / "yolov8n.pt")(imgs, imgsz=160, save_txt=True, save_crop=True) + save_path = Path(results[0].save_dir) + for r in results: + im_name = Path(r.path).stem + cls_idxs = r.boxes.cls.int().tolist() + # Check label path + labels = save_path / f"labels/{im_name}.txt" + assert labels.exists() + # Check detections match label count + assert len(r.boxes.data) == len([l for l in labels.read_text().splitlines() if l]) + # Check crops path and files + crop_dirs = [p for p in (save_path / "crops").iterdir()] + crop_files = [f for p in crop_dirs for f in p.glob("*")] + # Crop directories match detections + assert all([r.names.get(c) in [d.name for d in crop_dirs] for c in cls_idxs]) + # Same number of crops as detections + assert len([f for f in crop_files if im_name in f.name]) == len(r.boxes.data) + + @pytest.mark.skipif(not ONLINE, reason="environment is offline") def test_data_utils(): """Test utility functions in ultralytics/data/utils.py."""