File size: 2,554 Bytes
6a5bb7e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | """用训练好的模型对未标注图片做推理，输出labelme JSON预标注"""
import json
import argparse
from pathlib import Path
from ultralytics import YOLO
IMG_W, IMG_H = 3200, 1800
def yolo_box_to_labelme_shape(cx_n, cy_n, w_n, h_n):
cx = cx_n * IMG_W
cy = cy_n * IMG_H
w = w_n * IMG_W
h = h_n * IMG_H
x1, y1 = cx - w / 2, cy - h / 2
x2, y2 = cx + w / 2, cy + h / 2
return {
"label": "goat",
"points": [[round(x1, 4), round(y1, 4)], [round(x2, 4), round(y2, 4)]],
"group_id": None,
"description": "",
"shape_type": "rectangle",
"flags": {},
"mask": None,
}
def make_labelme_json(img_path: Path, shapes: list) -> dict:
return {
"version": "5.5.0",
"flags": {},
"shapes": shapes,
"imagePath": img_path.name,
"imageData": None,
"imageHeight": IMG_H,
"imageWidth": IMG_W,
}
def main():
parser = argparse.ArgumentParser()
parser.add_argument("model_path", help="best.pt路径")
parser.add_argument("unlabeled_root", help="Unlabeled_images根目录")
parser.add_argument("out_root", help="Auto_labeled输出根目录")
parser.add_argument("--conf", type=float, default=0.25)
parser.add_argument("--iou", type=float, default=0.45)
parser.add_argument("--imgsz", type=int, default=1280)
args = parser.parse_args()
model = YOLO(args.model_path)
unlabeled_root = Path(args.unlabeled_root)
out_root = Path(args.out_root)
img_paths = sorted(unlabeled_root.rglob("*.jpg"))
print(f"Found {len(img_paths)} unlabeled images")
for img_path in img_paths:
rel = img_path.parent.relative_to(unlabeled_root)
period = rel.parts[-1]
cam = rel.parts[0]
out_dir = out_root / cam / f"{period}Label"
out_dir.mkdir(parents=True, exist_ok=True)
results = model.predict(
str(img_path),
conf=args.conf,
iou=args.iou,
imgsz=args.imgsz,
verbose=False,
)
shapes = []
for box in results[0].boxes:
cx_n, cy_n, w_n, h_n = box.xywhn[0].tolist()
shapes.append(yolo_box_to_labelme_shape(cx_n, cy_n, w_n, h_n))
out_json = out_dir / (img_path.stem + ".json")
out_json.write_text(
json.dumps(make_labelme_json(img_path, shapes), ensure_ascii=False, indent=2),
encoding="utf-8",
)
print(f"Done. Output: {out_root}")
if __name__ == "__main__":
main()
|