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
parent
378ed74967
commit
8c30c583b0
23
detect.py
23
detect.py
|
@ -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')
|
||||||
|
|
Loading…
Reference in New Issue