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"
    
    # 이미지가 있는 디렉토리 (사용자가 경로 변경시 여기 수정)
    # 여기서는 input_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:"helmet", 2:"no_helmet"
    # Person 제거
    category_map = {}
    for cat in original_categories:
        cid = cat['id']
        cname = cat['name']
        if cname == "Helmet":
            category_map[cid] = (1, "helmet")
        elif cname == "No-Helmet":
            category_map[cid] = (2, "no_helmet")
        elif cname == "Person":
            # person은 제외하므로 맵핑하지 않음
            pass
        else:
            # 혹시 모르는 추가 클래스가 있다면 제외
            # 필요에 따라 처리 가능
            pass
    
    # person 클래스(3번 id) 어노테이션 제거
    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)
        # person에 해당하는 ann은 스킵
    
    # 최종 카테고리 리스트 구성
    # 정렬을 위해 category_map에서 추출
    final_categories = [
        {"id": 1, "name": "helmet", "supercategory": ""},
        {"id": 2, "name": "no_helmet", "supercategory": ""}
    ]
    
    # 이미지 경로 수정:
    # 현재 file_name은 그대로 두고, 만약 경로를 변경하고 싶다면 아래와 같이 처리할 수 있음.
    # 여기서는 별도 수정 없이 file_name 그대로 유지.
    # 필요하다면:
    # for img in images:
    #     # 예: images 디렉토리 구조를 변경하고 싶다면 여기서 수정
    #     # 현재는 그대로 둠
    #     pass
    
    # 최종 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()