Update detect.py (Fix save-csv: Ensure header is written to CSV) (#13472)

This commit resolves an issue where the save-csv command did not write the CSV header. The code now correctly saves the header in the CSV file.

Signed-off-by: Ali Ghanbari <alighanbari446@gmail.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Ali Ghanbari 2025-01-02 00:16:59 +03:30 committed by GitHub
parent 915ce21ebc
commit f003c3df9e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -219,9 +219,10 @@ def run(
def write_to_csv(image_name, prediction, confidence):
"""Writes prediction data for an image to a CSV file, appending if the file exists."""
data = {"Image Name": image_name, "Prediction": prediction, "Confidence": confidence}
file_exists = os.path.isfile(csv_path)
with open(csv_path, mode="a", newline="") as f:
writer = csv.DictWriter(f, fieldnames=data.keys())
if not csv_path.is_file():
if not file_exists:
writer.writeheader()
writer.writerow(data)