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>

detect.py

@@ -30,6 +30,8 @@
 
 import argparse
 import os
+import csv
+import pandas as pd
 import platform
 import sys
 from pathlib import Path
@@ -63,6 +65,7 @@ def run(
         device='',  # cuda device, i.e. 0 or 0,1,2,3 or cpu
         view_img=False,  # show results
         save_txt=False,  # save results to *.txt
+        save_in_csv=False,  # save results in CSV format
         save_conf=False,  # save confidences in --save-txt labels
         save_crop=False,  # save cropped prediction boxes
         nosave=False,  # do not save images/videos
@@ -135,6 +138,21 @@ def run(
         # Second-stage classifier (optional)
         # 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}
+            if not csv_path.is_file():
+                with open(csv_path, mode='w', newline='') as f:
+                    writer = csv.DictWriter(f, fieldnames=data.keys())
+                    writer.writeheader()
+            else:
+                with open(csv_path, mode='a', newline='') as f:
+                    writer = csv.DictWriter(f, fieldnames=data.keys())
+                    writer.writerow(data)
+
         # Process predictions
         for i, det in enumerate(pred):  # per image
             seen += 1
@@ -162,6 +180,14 @@ def run(
 
                 # Write results
                 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_in_csv:
+                        write_to_csv(p.name, label, confidence_str)
+
                     if save_txt:  # Write to file
                         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
@@ -229,6 +255,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('--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-in-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-crop', action='store_true', help='save cropped prediction boxes')
     parser.add_argument('--nosave', action='store_true', help='do not save images/videos')