Skip to content

Commit

Permalink
Merge branch 'master' into main
Browse files Browse the repository at this point in the history
  • Loading branch information
UltralyticsAssistant committed Aug 24, 2024
2 parents fc4942e + 41b5f45 commit 72f7a2f
Showing 1 changed file with 14 additions and 2 deletions.
16 changes: 14 additions & 2 deletions detect.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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")

Expand Down Expand Up @@ -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")
Expand Down

0 comments on commit 72f7a2f

Please sign in to comment.