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')