Add new `--save-csv` argument to detect.py (#12042)

* Update detect.py

added support for saving result in csv,used for testing

Signed-off-by: Akash A Desai <62583018+akashAD98@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update detect.py

Signed-off-by: Akash A Desai <62583018+akashAD98@users.noreply.github.com>

* Update detect.py

changed save_in_csv to save_csv

Signed-off-by: Akash A Desai <62583018+akashAD98@users.noreply.github.com>

* Update detect.py

Signed-off-by: Akash A Desai <62583018+akashAD98@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Signed-off-by: Akash A Desai <62583018+akashAD98@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
pull/12079/head
Akash A Desai 2023-09-04 16:22:33 +05:30 committed by GitHub
parent 378ed74967
commit 8c30c583b0
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 23 additions and 0 deletions

View File

@ -29,6 +29,7 @@ Usage - formats:
""" """
import argparse import argparse
import csv
import os import os
import platform import platform
import sys import sys
@ -63,6 +64,7 @@ def run(
device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
view_img=False, # show results view_img=False, # show results
save_txt=False, # save results to *.txt save_txt=False, # save results to *.txt
save_csv=False, # save results in CSV format
save_conf=False, # save confidences in --save-txt labels save_conf=False, # save confidences in --save-txt labels
save_crop=False, # save cropped prediction boxes save_crop=False, # save cropped prediction boxes
nosave=False, # do not save images/videos nosave=False, # do not save images/videos
@ -135,6 +137,18 @@ def run(
# Second-stage classifier (optional) # Second-stage classifier (optional)
# pred = utils.general.apply_classifier(pred, classifier_model, im, im0s) # pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)
# Define the path for the CSV file
csv_path = save_dir / 'predictions.csv'
# Create or append to the CSV file
def write_to_csv(image_name, prediction, confidence):
data = {'Image Name': image_name, 'Prediction': prediction, 'Confidence': confidence}
with open(csv_path, mode='a', newline='') as f:
writer = csv.DictWriter(f, fieldnames=data.keys())
if not csv_path.is_file():
writer.writeheader()
writer.writerow(data)
# Process predictions # Process predictions
for i, det in enumerate(pred): # per image for i, det in enumerate(pred): # per image
seen += 1 seen += 1
@ -162,6 +176,14 @@ def run(
# Write results # Write results
for *xyxy, conf, cls in reversed(det): for *xyxy, conf, cls in reversed(det):
c = int(cls) # integer class
label = names[c] if hide_conf else f'{names[c]}'
confidence = float(conf)
confidence_str = f'{confidence:.2f}'
if save_csv:
write_to_csv(p.name, label, confidence_str)
if save_txt: # Write to file if save_txt: # Write to file
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
line = (cls, *xywh, conf) if save_conf else (cls, *xywh) # label format line = (cls, *xywh, conf) if save_conf else (cls, *xywh) # label format
@ -229,6 +251,7 @@ def parse_opt():
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
parser.add_argument('--view-img', action='store_true', help='show results') parser.add_argument('--view-img', action='store_true', help='show results')
parser.add_argument('--save-txt', action='store_true', help='save results to *.txt') parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
parser.add_argument('--save-csv', action='store_true', help='save results in CSV format')
parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels') parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
parser.add_argument('--save-crop', action='store_true', help='save cropped prediction boxes') parser.add_argument('--save-crop', action='store_true', help='save cropped prediction boxes')
parser.add_argument('--nosave', action='store_true', help='do not save images/videos') parser.add_argument('--nosave', action='store_true', help='do not save images/videos')