#!/usr/bin/env python # -*- coding: utf-8 -*- """ 여러 개의 데이터셋(예: 4개)을 하나의 데이터셋으로 합치는 스크립트입니다. - 각 데이터셋의 이미지 파일들을 하나의 폴더에 평면화하여 복사하고, - 각 데이터셋의 JSON 어노테이션 정보를 합쳐 하나의 JSON 파일로 생성합니다. 주요 기능: 1. 사용자가 설정한 여러 데이터셋의 정보를 기반으로 JSON 파일을 로드하여, 이미지 정보와 어노테이션 정보를 추출합니다. - JSON 파일이 COCO 형식("annotations" 키 존재)일 경우와 그렇지 않은 경우 모두 지원합니다. 2. 각 데이터셋의 이미지에 대해 새 파일명을 부여합니다. - 기존에는 file_name의 "/"를 "_"로 변경했으나, 이번 수정에서는 오직 데이터셋별 접두어(prefix)만 앞에 붙여 충돌을 방지합니다. 3. 이미지와 어노테이션의 ID를 전체 데이터셋 기준으로 새롭게 순차 부여하여 중복 문제를 해결합니다. 4. 어노테이션의 클래스(카테고리) 정보는 각 데이터셋마다 다를 수 있으므로, 사용자가 최종 데이터셋에 포함할 클래스들을 지정하고, 이름 변경이나 병합, 삭제가 가능하도록 지원합니다. 5. 모든 이미지 파일들을 지정한 대상 폴더로 복사하며, 복사 진행 상황(퍼센트)과 전체 소요시간을 출력합니다. 6. 최종 합쳐진 JSON 파일을 저장하며, 이미지의 file_name은 복사된 이미지의 절대경로로 업데이트됩니다. 사용 방법: 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": "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_combined" 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 = { "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 순차 부여 # 최종 데이터셋의 카테고리 정보를 구축할 글로벌 딕셔너리 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. 각 데이터셋의 카테고리 정보 처리 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: continue if final_name not in global_final_categories: global_final_categories[final_name] = final_cat_counter final_cat_counter += 1 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"] # [수정] 파일명 수정: 기존 경로 전체를 변환하지 않고, basename만 추출하여 prefix만 추가 base_file_name = os.path.basename(orig_file_name) new_file_name = prefix + base_file_name new_img = img.copy() new_img["id"] = new_image_id # 초기에는 파일명은 단순 파일명으로 설정 (추후 JSON 저장 시 절대경로로 업데이트) 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: orig_cat_id = ann.get("category_id") if orig_cat_id not in dataset_category_id_mapping[dataset_name]: continue new_ann = ann.copy() 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 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 파일 구성 ---------------- 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 파일 저장 전 임시 키 제거 및 절대경로 업데이트 ---------------- # [수정] 각 이미지의 file_name을 combined_images_dir를 포함한 절대경로로 업데이트 for img in combined_images: img["file_name"] = os.path.join(combined_images_dir, img["file_name"]) 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()