70 lines
2.2 KiB
Python
70 lines
2.2 KiB
Python
|
import json
|
||
|
|
from pathlib import Path
|
||
|
|
from pycocotools.coco import COCO
|
||
|
|
|
||
|
|
def generate_coco_labels(annotations_file, images_dir, labels_dir):
|
||
|
|
"""从COCO JSON标注生成YOLO格式标签"""
|
||
|
|
annotations_file = Path(annotations_file)
|
||
|
|
images_dir = Path(images_dir)
|
||
|
|
labels_dir = Path(labels_dir)
|
||
|
|
|
||
|
|
labels_dir.mkdir(parents=True, exist_ok=True)
|
||
|
|
|
||
|
|
print(f"Loading annotations from {annotations_file}...")
|
||
|
|
coco = COCO(str(annotations_file))
|
||
|
|
|
||
|
|
img_ids = sorted(coco.getImgIds())
|
||
|
|
|
||
|
|
total_boxes = 0
|
||
|
|
|
||
|
|
for img_id in img_ids:
|
||
|
|
img_info = coco.loadImgs(img_id)[0]
|
||
|
|
img_name = img_info['file_name']
|
||
|
|
img_path = images_dir / img_name
|
||
|
|
|
||
|
|
if not img_path.exists():
|
||
|
|
continue
|
||
|
|
|
||
|
|
ann_ids = coco.getAnnIds(imgIds=img_id)
|
||
|
|
anns = coco.loadAnns(ann_ids)
|
||
|
|
|
||
|
|
label_path = labels_dir / (img_path.stem + '.txt')
|
||
|
|
|
||
|
|
with open(label_path, 'w') as f:
|
||
|
|
for ann in anns:
|
||
|
|
if ann['category_id'] == 0:
|
||
|
|
continue
|
||
|
|
|
||
|
|
cat_id = ann['category_id']
|
||
|
|
bbox = ann['bbox']
|
||
|
|
x, y, w, h = bbox
|
||
|
|
|
||
|
|
x_center = x + w / 2
|
||
|
|
y_center = y + h / 2
|
||
|
|
|
||
|
|
img_w = img_info['width']
|
||
|
|
img_h = img_info['height']
|
||
|
|
|
||
|
|
x_center /= img_w
|
||
|
|
y_center /= img_h
|
||
|
|
w /= img_w
|
||
|
|
h /= img_h
|
||
|
|
|
||
|
|
f.write(f"{cat_id} {x_center:.6f} {y_center:.6f} {w:.6f} {h:.6f}\n")
|
||
|
|
total_boxes += 1
|
||
|
|
|
||
|
|
print(f"Generated {total_boxes} bounding boxes in {labels_dir}")
|
||
|
|
print(f"Labels saved to: {labels_dir}")
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
base_dir = Path(r"C:\Users\16337\PycharmProjects\Security\datasets\coco")
|
||
|
|
|
||
|
|
generate_coco_labels(
|
||
|
|
annotations_file=base_dir / "annotations" / "instances_val2017.json",
|
||
|
|
images_dir=base_dir / "images" / "val2017",
|
||
|
|
labels_dir=base_dir / "labels" / "val2017"
|
||
|
|
)
|
||
|
|
|
||
|
|
print("\nNow run validation with batch=1:")
|
||
|
|
print('yolo val model=yolo11n_int8_b1_8.engine data=coco.yaml imgsz=640 rect=False batch=1')
|