Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add detect.py --hide-conf --hide-labels --line-thickness options #2658

Merged
merged 9 commits into from
Apr 23, 2021
8 changes: 6 additions & 2 deletions detect.py
Original file line number Diff line number Diff line change
Expand Up @@ -101,6 +101,7 @@ def detect(save_img=False):
s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string

# Write results
label = None
for *xyxy, conf, cls in reversed(det):
if save_txt: # Write to file
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
Expand All @@ -109,8 +110,9 @@ def detect(save_img=False):
f.write(('%g ' * len(line)).rstrip() % line + '\n')

if save_img or view_img: # Add bbox to image
label = f'{names[int(cls)]} {conf:.2f}'
plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)
if not opt.hide_labels:
label = f'{names[int(cls)]} {conf:.2f}'
plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=opt.line_thickness)

# Print time (inference + NMS)
print(f'{s}Done. ({t2 - t1:.3f}s)')
Expand Down Expand Up @@ -165,6 +167,8 @@ def detect(save_img=False):
parser.add_argument('--project', default='runs/detect', help='save results to project/name')
parser.add_argument('--name', default='exp', help='save results to project/name')
parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
parser.add_argument('--line-thickness', default=3, type=int, help='thickness of the bounding boxes')
parser.add_argument('--hide-labels', default=False, action='store_true', help='hide class labels on bounding boxes')
opt = parser.parse_args()
print(opt)
check_requirements(exclude=('pycocotools', 'thop'))
Expand Down