Simplify `auto_annotate` function (#19400)

Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
pull/19418/head
Burhan 2 weeks ago committed by GitHub
parent b0e8938172
commit 58936e9661
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
  1. 12
      ultralytics/data/annotator.py

@ -58,15 +58,13 @@ def auto_annotate(
for result in det_results:
class_ids = result.boxes.cls.int().tolist() # noqa
if len(class_ids):
if class_ids:
boxes = result.boxes.xyxy # Boxes object for bbox outputs
sam_results = sam_model(result.orig_img, bboxes=boxes, verbose=False, save=False, device=device)
segments = sam_results[0].masks.xyn # noqa
with open(f"{Path(output_dir) / Path(result.path).stem}.txt", "w") as f:
for i in range(len(segments)):
s = segments[i]
if len(s) == 0:
continue
segment = map(str, segments[i].reshape(-1).tolist())
f.write(f"{class_ids[i]} " + " ".join(segment) + "\n")
for i, s in enumerate(segments):
if s.any():
segment = map(str, s.reshape(-1).tolist())
f.write(f"{class_ids[i]} " + " ".join(segment) + "\n")

Loading…
Cancel
Save