import json
import os

def main():
    # 원본 JSON 경로 (사용자가 변경 가능)
    input_json = "/DATA2/ltb/mmyolo/roboflow_helmet/train/_annotations.coco.json"
    # 결과 JSON 경로
    output_json = "/DATA2/ltb/mmyolo/roboflow_helmet/robo_no_helmet.json"
    
    # 이미지가 있는 디렉토리 (사용자가 경로 변경시 여기 수정)
    images_dir = "/mmyolo/roboflow_helmet/train"
    
    with open(input_json, 'r') as f:
        data = json.load(f)
    
    # 원본 카테고리 정보를 가져옴
    original_categories = data.get('categories', [])
    # 원본 이미지, 어노테이션
    images = data.get('images', [])
    annotations = data.get('annotations', [])
    
    # 카테고리 매핑 (원본 → 최종)
    # 원본: 1:"Helmet", 2:"No-Helmet", 3:"Person"
    # 최종: 1:"no_helmet", 2:"background"
    category_map = {}
    for cat in original_categories:
        cid = cat['id']
        cname = cat['name']
        if cname == "No-Helmet":
            category_map[cid] = (1, "no_helmet")
        else:
            # 'Helmet', 'Person' 및 기타 모든 클래스는 제외
            pass
    
    # 어노테이션 필터링: 'No-Helmet'에 해당하는 것만 남김
    filtered_annotations = []
    for ann in annotations:
        cid = ann['category_id']
        if cid in category_map:
            # 매핑된 cid로 변경
            new_cid, new_name = category_map[cid]
            ann['category_id'] = new_cid
            filtered_annotations.append(ann)
        # 'No-Helmet' 이외의 어노테이션은 스킵
    
    # 최종 카테고리 리스트 구성
    final_categories = [
        {"id": 1, "name": "no_helmet", "supercategory": ""},
        {"id": 2, "name": "background", "supercategory": ""}
    ]
    
    # 최종 COCO 구조
    final_coco = {
        "images": images,
        "annotations": filtered_annotations,
        "categories": final_categories
    }
    
    # 저장
    with open(output_json, 'w') as f:
        json.dump(final_coco, f, indent=4)
    
    print("변환 완료!")
    print(f"입력 파일: {input_json}")
    print(f"결과 파일: {output_json}")
    print(f"총 이미지 수: {len(images)}")
    print(f"최종 어노테이션 수: {len(filtered_annotations)}")
    print("카테고리:", [c['name'] for c in final_categories])

if __name__ == "__main__":
    main()