78 lines
2.5 KiB
Python
78 lines
2.5 KiB
Python
"""
|
|
Pre-label images using an existing YOLO model.
|
|
Generates .txt labels for unlabeled images so the annotator can load them for review.
|
|
|
|
Usage:
|
|
python prelabel.py [image_dir] [--model boss-v1] [--conf 0.20]
|
|
"""
|
|
import argparse
|
|
import glob
|
|
import os
|
|
|
|
def run_prelabel(args):
|
|
"""Run pre-labeling. Called from main() or manage.py."""
|
|
img_dir = os.path.abspath(args.img_dir)
|
|
model_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "models", f"{args.model}.pt")
|
|
|
|
if not os.path.exists(model_path):
|
|
print(f"Model not found: {model_path}")
|
|
return
|
|
|
|
from ultralytics import YOLO
|
|
model = YOLO(model_path)
|
|
|
|
extensions = ("*.jpg", "*.jpeg", "*.png")
|
|
files = []
|
|
for ext in extensions:
|
|
files.extend(glob.glob(os.path.join(img_dir, ext)))
|
|
files.sort()
|
|
|
|
# Only process unlabeled images
|
|
unlabeled = []
|
|
for f in files:
|
|
label_path = os.path.splitext(f)[0] + ".txt"
|
|
if not os.path.exists(label_path):
|
|
unlabeled.append(f)
|
|
|
|
print(f"Found {len(files)} images, {len(files) - len(unlabeled)} already labeled, {len(unlabeled)} to pre-label")
|
|
|
|
if not unlabeled:
|
|
print("All images already have labels!")
|
|
return
|
|
|
|
labeled = 0
|
|
skipped = 0
|
|
for filepath in unlabeled:
|
|
results = model(filepath, conf=args.conf, verbose=False)
|
|
boxes = results[0].boxes
|
|
|
|
if len(boxes) == 0:
|
|
skipped += 1
|
|
continue
|
|
|
|
label_path = os.path.splitext(filepath)[0] + ".txt"
|
|
with open(label_path, "w") as f:
|
|
for box in boxes:
|
|
cls = int(box.cls[0])
|
|
xywhn = box.xywhn[0] # normalized center x, y, w, h
|
|
cx, cy, w, h = xywhn.tolist()
|
|
f.write(f"{cls} {cx:.6f} {cy:.6f} {w:.6f} {h:.6f}\n")
|
|
|
|
labeled += 1
|
|
fname = os.path.basename(filepath)
|
|
conf = boxes.conf[0].item()
|
|
print(f" {fname}: {len(boxes)} box(es), best conf={conf:.2f}")
|
|
|
|
print(f"\nPre-labeled {labeled} images, skipped {skipped} (no detections)")
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Pre-label images with YOLO model")
|
|
parser.add_argument("img_dir", nargs="?", default="../../training-data/kulemak/raw")
|
|
parser.add_argument("--model", default="boss-kulemak", help="Model name in models/")
|
|
parser.add_argument("--conf", type=float, default=0.20, help="Confidence threshold")
|
|
args = parser.parse_args()
|
|
run_prelabel(args)
|
|
|
|
if __name__ == "__main__":
|
|
main()
|