Simplify `auto_annotate` function (#19400)

Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
pull/19418/head
Burhan 1 month 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: for result in det_results:
class_ids = result.boxes.cls.int().tolist() # noqa class_ids = result.boxes.cls.int().tolist() # noqa
if len(class_ids): if class_ids:
boxes = result.boxes.xyxy # Boxes object for bbox outputs boxes = result.boxes.xyxy # Boxes object for bbox outputs
sam_results = sam_model(result.orig_img, bboxes=boxes, verbose=False, save=False, device=device) sam_results = sam_model(result.orig_img, bboxes=boxes, verbose=False, save=False, device=device)
segments = sam_results[0].masks.xyn # noqa segments = sam_results[0].masks.xyn # noqa
with open(f"{Path(output_dir) / Path(result.path).stem}.txt", "w") as f: with open(f"{Path(output_dir) / Path(result.path).stem}.txt", "w") as f:
for i in range(len(segments)): for i, s in enumerate(segments):
s = segments[i] if s.any():
if len(s) == 0: segment = map(str, s.reshape(-1).tolist())
continue f.write(f"{class_ids[i]} " + " ".join(segment) + "\n")
segment = map(str, segments[i].reshape(-1).tolist())
f.write(f"{class_ids[i]} " + " ".join(segment) + "\n")

Loading…
Cancel
Save