From 41b5f45096e101b53fce50e7b641f600265ff32e Mon Sep 17 00:00:00 2001 From: Harsh Dhamecha <62664549+harshdhamecha@users.noreply.github.com> Date: Sun, 25 Aug 2024 03:10:38 +0530 Subject: [PATCH] Provide option for saving box coordinates in Pascal-VOC format. (#13272) * Provide option for saaving box coordinates in Pascal-VOC format. - Added save-format argument for saving box coordinates in Pascal-VOC format (Xmin, Ymin, Xmax, Ymax) when save_txt is True, 0 for YOLO and 1 for Pascal-VOC, By default - 0 * Auto-format by https://ultralytics.com/actions --------- Co-authored-by: UltralyticsAssistant --- detect.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/detect.py b/detect.py index 8a25ac235a3c..57d778740273 100644 --- a/detect.py +++ b/detect.py @@ -78,6 +78,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_format=0, # save boxes coordinates in YOLO format or Pascal-VOC format (0 for YOLO and 1 for Pascal-VOC) save_csv=False, # save results in CSV format save_conf=False, # save confidences in --save-txt labels save_crop=False, # save cropped prediction boxes @@ -260,8 +261,13 @@ def write_to_csv(image_name, prediction, confidence): 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 + if save_format == 0: + coords = ( + (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() + ) # normalized xywh + else: + coords = (torch.tensor(xyxy).view(1, 4) / gn).view(-1).tolist() # xyxy + line = (cls, *coords, conf) if save_conf else (cls, *coords) # label format with open(f"{txt_path}.txt", "a") as f: f.write(("%g " * len(line)).rstrip() % line + "\n") @@ -369,6 +375,12 @@ 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-format", + type=int, + default=0, + help="whether to save boxes coordinates in YOLO format or Pascal-VOC format when save-txt is True, 0 for YOLO and 1 for Pascal-VOC", + ) 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-crop", action="store_true", help="save cropped prediction boxes")