#!/usr/bin/env python # -*- coding: utf-8 -*- """ 여러 개의 데이터셋(예: 4개)을 하나의 데이터셋으로 합치는 스크립트입니다. - 각 데이터셋의 이미지 파일들을 하나의 폴더에 평면화하여 복사하고, - 각 데이터셋의 JSON 어노테이션 정보를 합쳐 하나의 JSON 파일로 생성합니다. 주요 기능: 1. 사용자가 설정한 여러 데이터셋의 정보를 기반으로 JSON 파일을 로드하여, 이미지 정보와 어노테이션 정보를 추출합니다. - JSON 파일이 COCO 형식("annotations" 키 존재)일 경우와 그렇지 않은 경우 모두 지원합니다. 2. 각 데이터셋의 이미지에 대해 새 파일명을 부여합니다. - 기존 file_name에서 "/"를 "_"로 변경하여 한 폴더에 평면화 저장이 가능하도록 함. - 추가적으로 데이터셋별 접두어(prefix)를 파일명 앞에 붙여 충돌을 방지할 수 있습니다. 3. 이미지와 어노테이션의 ID를 전체 데이터셋 기준으로 새롭게 순차 부여하여 중복 문제를 해결합니다. 4. 어노테이션의 클래스(카테고리) 정보는 각 데이터셋마다 다를 수 있으므로, 사용자가 최종 데이터셋에 포함할 클래스들을 지정하고, 이름 변경이나 병합, 삭제가 가능하도록 지원합니다. - final_category_mapping 변수에서 각 원본 클래스(소문자)를 최종 클래스 이름으로 매핑합니다. 값이 None인 경우 해당 클래스는 최종 데이터셋에 포함되지 않습니다. 예) {"worker": "worker", "head": "head", "helmet": "head", "background": None} => worker는 그대로, head는 head로, helmet은 head로 병합, background는 삭제됨. 5. 모든 이미지 파일들을 지정한 대상 폴더로 복사하며, 복사 진행 상황(퍼센트)과 전체 소요시간을 출력합니다. 6. 최종 합쳐진 JSON 파일을 저장하며, 임시로 사용한 키("orig_file_name", "dataset")는 제거합니다. 사용 방법: 1. 아래의 dataset_configs 리스트에 각 데이터셋의 정보를 입력합니다. 각 항목은 딕셔너리 형태로, 다음 키를 포함해야 합니다. - "name": 데이터셋 식별을 위한 이름 (예: "head", "fall", "original", "helmet") - "images_dir": 해당 데이터셋의 이미지가 저장된 폴더의 절대경로 - "json_path": 해당 데이터셋의 어노테이션 JSON 파일 경로 - "prefix": 파일명에 붙일 접두어 (없으면 빈 문자열, 예: "head_", "fall_") 2. combined_dataset_dir 변수에 최종 합친 데이터셋을 저장할 상위 폴더의 경로를 지정합니다. 3. final_category_mapping 변수에서 최종 데이터셋에 포함할 클래스(카테고리)를 지정합니다. - key: 원본 데이터셋에 나타난 클래스 이름(소문자) - value: 최종 데이터셋에서 사용할 클래스 이름 (변경 없이 그대로 사용하려면 동일한 이름, 병합하려면 동일한 값, 삭제하려면 None) 4. 스크립트를 실행하면, 지정한 모든 데이터셋의 이미지와 어노테이션이 하나로 합쳐집니다. 이미지 파일은 combined_dataset_dir/images 폴더에 복사되고, JSON 파일은 combined_dataset_dir/combined.json으로 저장됩니다. 예시 실행: python merge_multiple_datasets.py 작성자: [작성자 이름 또는 팀] 작성일: [작성일] """ import os import json import shutil import time def load_and_separate_json(json_path): """ JSON 파일을 로드한 후, 이미지 정보와 어노테이션 정보를 분리하는 함수. COCO 형식처럼 "annotations" 키가 존재하면 해당 키의 내용을 사용하며, 그렇지 않을 경우 "images" 리스트 내에서 "file_name" 유무로 이미지와 어노테이션 정보를 구분합니다. 매개변수: json_path (str): JSON 파일의 절대/상대 경로. 반환값: images (list): 이미지 정보 리스트. annotations (list): 어노테이션 정보 리스트. categories (list): 카테고리 정보 리스트 (존재 시). info (dict): info 필드 (존재 시). licenses (list): licenses 필드 (존재 시). """ with open(json_path, 'r', encoding='utf-8') as f: data = json.load(f) if "annotations" in data: images = data.get("images", []) annotations = data.get("annotations", []) else: images = [] annotations = [] for item in data["images"]: if "file_name" in item: images.append(item) else: annotations.append(item) categories = data.get("categories", []) info = data.get("info", {}) licenses = data.get("licenses", []) return images, annotations, categories, info, licenses def main(): # ================= 사용자 설정 변수 (필요 시 수정) ================= # 1. 데이터셋 구성 정보 # dataset_configs: 합칠 데이터셋들의 정보를 리스트로 구성 # 각 데이터셋은 아래 정보를 포함하는 딕셔너리로 구성: # "name": 데이터셋 이름 (예: "head", "fall", "original", "helmet") # "images_dir": 해당 데이터셋의 이미지 폴더 경로 (절대경로 권장) # "json_path": 해당 데이터셋의 어노테이션 JSON 파일 경로 # "prefix": 이미지 파일명에 붙일 접두어 (없으면 빈 문자열) dataset_configs = [ { "name": "dl1", "images_dir": "/ltb/media/ltb/90887173887158A4/Users/ltb/si/dl-e-c-project/images", "json_path": "/ltb/media/ltb/90887173887158A4/Users/ltb/si/dl-e-c-project/mm_coco_annotations.json", "prefix": "dl1_" }, { "name": "dl2", "images_dir": "/ltb/media/ltb/90887173887158A4/Users/ltb/si/dl-e-c-project_2/images", "json_path": "/ltb/media/ltb/90887173887158A4/Users/ltb/si/dl-e-c-project_2/mm_coco_annotations.json", "prefix": "dl2_" }, { "name": "dl3", "images_dir": "/ltb/media/ltb/90887173887158A4/Users/ltb/si/dl-e-c-project_3/images", "json_path": "/ltb/media/ltb/90887173887158A4/Users/ltb/si/dl-e-c-project_3/mm_coco_annotations.json", "prefix": "dl3_" }, { "name": "dl4", "images_dir": "/ltb/media/ltb/90887173887158A4/Users/ltb/si/dl-e-c-project_4/images", "json_path": "/ltb/media/ltb/90887173887158A4/Users/ltb/si/dl-e-c-project_4/mm_coco_annotations.json", "prefix": "dl4_" } ] # 2. 최종 합친 데이터셋 저장 경로 combined_dataset_dir = "/ltb/media/ltb/90887173887158A4/Users/ltb/si/dl_all_fusion" combined_images_dir = os.path.join(combined_dataset_dir, "images") combined_json_path = os.path.join(combined_dataset_dir, "combined.json") # 3. 최종 데이터셋에 포함할 클래스(카테고리) 설정 # final_category_mapping: key는 원본 데이터셋에 나타난 클래스명(소문자), # value는 최종 데이터셋에서 사용할 클래스명. # value가 None이면 해당 클래스는 최종 데이터셋에 포함되지 않습니다. # 예시: # - head 데이터셋: worker, head # - fall 데이터셋: worker, background # - original: worker, helmet, head # - helmet: worker, helmet, head # 사용자가 원하는 대로 이름을 변경하거나 병합할 수 있습니다. # 예: 만약 helmet과 head를 모두 head로 병합하고 싶다면 "helmet": "head"로 지정. final_category_mapping = { "worker_dl": "worker", "signalman_dl": "signalman", "hardhat_dl": "helmet", "worker_helmat_off_dl": "worker", "harness_dl": "harness", "mixer_truck_dl": "mixer_truck", "excavator_dl": "excavator" } # ============================================================================= # 폴더 생성 (이미지 저장용) os.makedirs(combined_images_dir, exist_ok=True) # 데이터셋별 정보를 빠르게 조회하기 위한 딕셔너리 (키: 데이터셋 이름) dataset_info = {cfg["name"]: cfg for cfg in dataset_configs} # ---------------- 글로벌 변수 초기화 ---------------- combined_images = [] # 최종 JSON의 "images" 항목 combined_annotations = [] # 최종 JSON의 "annotations" 항목 new_image_id = 1 # 전체 데이터셋 기준 이미지 ID 순차 부여 new_annotation_id = 1 # 전체 데이터셋 기준 어노테이션 ID 순차 부여 # 최종 데이터셋의 카테고리 정보를 구축할 글로벌 딕셔너리 # key: 최종 클래스 이름, value: 새 id global_final_categories = {} final_cat_counter = 1 # 각 데이터셋별 이미지 id 매핑 (원래 JSON의 id -> 새 id) dataset_image_id_mapping = {} # key: dataset name, value: dict(mapping) # 각 데이터셋별 카테고리 id 매핑 (원본 카테고리 id -> 최종 카테고리 id) dataset_category_id_mapping = {} # key: dataset name, value: dict(mapping) # ---------------- 각 데이터셋 처리 ---------------- for cfg in dataset_configs: dataset_name = cfg["name"] json_path = cfg["json_path"] prefix = cfg["prefix"] print(f">> '{dataset_name}' 데이터셋 JSON 파일 로드 중... ({json_path})") images, annotations, categories, info, licenses = load_and_separate_json(json_path) # 초기화: 현재 데이터셋에 대한 이미지 및 카테고리 ID 매핑 dataset_image_id_mapping[dataset_name] = {} dataset_category_id_mapping[dataset_name] = {} # 1. 각 데이터셋의 카테고리 정보 처리 # 원본 JSON의 categories 필드를 순회하며, 각 카테고리명을 소문자로 변환 후, # final_category_mapping에 있는지 확인. 만약 값이 None이면 최종에 포함하지 않음. for cat in categories: orig_cat_name = cat.get("name", "").lower() if orig_cat_name not in final_category_mapping: # 매핑에 없으면 무시 (필요시 사용자에게 경고 출력 가능) continue final_name = final_category_mapping[orig_cat_name] if final_name is None: # None이면 최종 데이터셋에 포함하지 않음. continue # global_final_categories에 해당 final_name이 없으면 새 id 할당 if final_name not in global_final_categories: global_final_categories[final_name] = final_cat_counter final_cat_counter += 1 # 현재 데이터셋의 원본 카테고리 id를 최종 id로 매핑 dataset_category_id_mapping[dataset_name][cat["id"]] = global_final_categories[final_name] # 2. 이미지 정보 처리 print(f">> '{dataset_name}' 데이터셋 이미지 정보 처리 중...") for img in images: orig_file_name = img["file_name"] new_file_name = prefix + orig_file_name.replace("/", "_") new_img = img.copy() new_img["id"] = new_image_id new_img["file_name"] = new_file_name new_img["orig_file_name"] = orig_file_name new_img["dataset"] = dataset_name combined_images.append(new_img) dataset_image_id_mapping[dataset_name][img["id"]] = new_image_id new_image_id += 1 # 3. 어노테이션 정보 처리 print(f">> '{dataset_name}' 데이터셋 어노테이션 정보 처리 중...") for ann in annotations: # 각 어노테이션의 category_id를 최종 id로 변경하기 위해, # 현재 데이터셋의 category id 매핑을 확인합니다. orig_cat_id = ann.get("category_id") if orig_cat_id not in dataset_category_id_mapping[dataset_name]: # 해당 어노테이션의 클래스가 최종 데이터셋에 포함되지 않는 경우 건너뜁니다. # (예: final_category_mapping에서 None으로 지정된 경우) # 필요시 경고 메시지 출력 가능 continue new_ann = ann.copy() # 이미지 id 변경: 현재 데이터셋 내 매핑 사용 old_img_id = ann["image_id"] if old_img_id in dataset_image_id_mapping[dataset_name]: new_ann["image_id"] = dataset_image_id_mapping[dataset_name][old_img_id] else: print(f"경고: '{dataset_name}' 데이터셋 어노테이션에서 이미지 id {old_img_id}를 찾을 수 없습니다.") continue # 최종 카테고리 id로 갱신 new_ann["category_id"] = dataset_category_id_mapping[dataset_name][orig_cat_id] new_ann["id"] = new_annotation_id new_annotation_id += 1 combined_annotations.append(new_ann) # ----------------------------------------------------- # ---------------- 최종 JSON 파일 구성 ---------------- # 최종 카테고리 목록을 global_final_categories 딕셔너리로부터 생성 (id 순 정렬) combined_categories = [{"id": cid, "name": name} for name, cid in sorted(global_final_categories.items(), key=lambda x: x[1])] combined_info = { "description": "여러 데이터셋을 합친 데이터셋 (최종 클래스 매핑 적용)", "date_created": time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) } combined_licenses = [] # 필요 시 라이선스 정보 추가 combined_json = { "info": combined_info, "licenses": combined_licenses, "images": combined_images, "annotations": combined_annotations, "categories": combined_categories } # ----------------------------------------------------- # ---------------- 이미지 파일 복사 ---------------- print(">> 이미지 파일 복사 시작...") start_time = time.time() total_images = len(combined_images) copied_count = 0 for img in combined_images: dataset_name = img["dataset"] orig_file_name = img["orig_file_name"] new_file_name = img["file_name"] source_images_dir = dataset_info[dataset_name]["images_dir"] source_path = os.path.join(source_images_dir, orig_file_name) dest_path = os.path.join(combined_images_dir, new_file_name) if os.path.exists(source_path): shutil.copy2(source_path, dest_path) else: print(f"\n경고: 소스 이미지 파일이 존재하지 않습니다: {source_path}") copied_count += 1 if copied_count % max(1, total_images // 100) == 0 or copied_count == total_images: progress = (copied_count / total_images) * 100 print(f"이미지 복사 진행 상황: {copied_count}/{total_images} ({progress:.2f}%)", end='\r') end_time = time.time() elapsed_time = end_time - start_time print(f"\n>> 이미지 복사 완료: 총 {copied_count}개 파일 복사, 소요 시간: {elapsed_time:.2f}초") # ----------------------------------------------------- # ---------------- JSON 파일 저장 전 임시 키 제거 ---------------- for img in combined_images: if "orig_file_name" in img: del img["orig_file_name"] if "dataset" in img: del img["dataset"] with open(combined_json_path, 'w', encoding='utf-8') as f: json.dump(combined_json, f, ensure_ascii=False, indent=4) print(f">> 합쳐진 JSON 파일 저장 완료: {combined_json_path}") if __name__ == '__main__': main()