#!/usr/bin/env python # -*- coding: utf-8 -*- """ [사용 안내] - 이 코드는 /DATA2/ltb/si/it1_worker_helmet 아래에 있는 ① 원본 어노테이션 파일 (예: it1_250326_1819_4class_modify.json) ② albu 변형 어노테이션 파일 (예: albu_aug/it1_250326_1819_4class_albu_aug.json) 을 하나의 통합 데이터셋으로 병합합니다. - 통합 후, 각 이미지(절대경로로 기록됨)를 offline resize하여 SCALE_FACTORS (예: 1.0, 2.0, 0.5) 배수로 확대/축소한 이미지를 새로 저장하고, 그에 따른 bbox 좌표도 함께 업데이트한 새로운 COCO JSON 파일을 생성합니다. - 출력 이미지는 OUT_IMG_DIR 아래에 원본 폴더 구조를 유지하여 저장되며, JSON 파일의 file_name 필드에는 절대경로가 기록되어 원본 파일 복사 없이 용량을 절약할 수 있습니다. - 코드 내 디버깅 출력(print)으로 주요 경로, 파라미터, 처리 개수 등을 확인할 수 있습니다. [전체 구성] 1. [Cell 1] 전역 설정 및 라이브러리 임포트 - 데이터셋 루트, 입력 JSON 파일 목록, 출력 경로, 스케일 팩터 등 전역 변수 설정 2. [Cell 2] 함수 정의 - load_coco(), save_coco() 등 기본 입출력 함수 - merge_coco_datasets(): 여러 개의 COCO JSON을 클래스 이름 기준으로 병합 - (offline resize용) 각 이미지별 bbox 스케일 조정 함수 등 3. [Cell 3] 메인 실행 루프 - 통합 데이터셋(merged_coco)을 생성한 후, 각 이미지에 대해 offline resize를 적용하여 새로운 이미지와 어노테이션을 생성 및 저장하고 최종 JSON 파일을 출력 (!!!) 하단의 전역 파라미터만 환경에 맞게 수정하면 전체 코드가 자동으로 반영됩니다. """ # ============================================================================= # [Cell 1] 전역 설정 및 라이브러리 임포트 # ============================================================================= import os import json import cv2 import math import numpy as np from tqdm import tqdm from collections import defaultdict # (!!!) 수정 가능한 글로벌 파라미터 # 데이터셋 루트 (원본 이미지와 알부 변형 이미지 모두 포함) DATASET_ROOT = "/DATA2/ltb/si/it1_worker_helmet" # 입력 JSON 파일 목록 (절대경로 기반) INPUT_JSON_PATHS = [ os.path.join(DATASET_ROOT, "it1_250326_1819_4class_modify.json"), # 원본 어노테이션 os.path.join(DATASET_ROOT, "albu_aug", "it1_250326_1819_4class_albu_aug.json") # albu 변형 어노테이션 ] # 출력 결과 저장 경로 OUT_ROOT = "/DATA2/ltb/mmyolo/it1_worker_helmet_resize" # 출력 폴더 (기존 파일 덮어쓰지 않도록 exp, exp2 등으로 변경 가능) OUT_IMG_DIR = os.path.join(OUT_ROOT, "images") # 변환된 이미지 저장 폴더 OUT_JSON_PATH = os.path.join(OUT_ROOT, "merged_resized.json") # 최종 COCO JSON 파일 # offline resize 시 적용할 스케일 팩터 (1.0은 원본, 2.0은 2배 확대, 0.5는 50% 축소) SCALE_FACTORS = [1.0, 2.0, 0.5, 5.0] # 디버깅 출력: 설정값 확인 print("✅ DATASET_ROOT:", DATASET_ROOT) print("✅ INPUT_JSON_PATHS:", INPUT_JSON_PATHS) print("✅ OUT_ROOT:", OUT_ROOT) print("✅ OUT_IMG_DIR:", OUT_IMG_DIR) print("✅ OUT_JSON_PATH:", OUT_JSON_PATH) print("✅ SCALE_FACTORS:", SCALE_FACTORS) # ============================================================================= # [Cell 2] 함수 정의 # ============================================================================= def load_coco(json_path): """COCO 형식 JSON 파일을 읽어 반환합니다.""" with open(json_path, 'r', encoding='utf-8') as f: return json.load(f) def save_coco(data, out_path): """COCO 형식 데이터를 JSON 파일로 저장합니다. (필요한 폴더가 없으면 생성)""" os.makedirs(os.path.dirname(out_path), exist_ok=True) with open(out_path, 'w', encoding='utf-8') as f: json.dump(data, f, indent=2, ensure_ascii=False) print("✅ Saved COCO JSON:", out_path) def merge_coco_datasets(json_paths): """ 여러 개의 COCO JSON 파일을 병합하여 하나의 통합 데이터셋을 생성합니다. - 클래스(카테고리)는 클래스 이름을 기준으로 병합하며, 새로운 id를 부여합니다. - 각 이미지와 어노테이션의 id는 새롭게 순차적으로 재할당됩니다. 반환: merged_coco: {"images": [...], "annotations": [...], "categories": [...]} """ merged_images = [] merged_annotations = [] merged_categories = {} # key: category name, value: category dict new_cat_id = 1 new_img_id = 1 new_ann_id = 1 for json_path in json_paths: data = load_coco(json_path) # 각 JSON 파일의 카테고리 정보를 확인하여 병합 (클래스 이름 기준) # 각 JSON에는 "categories"가 있으며, id와 name, supercategory 정보가 있음. # 각 파일마다 old_cat_dict: 기존 id -> category name mapping 생성 old_cat_dict = {} for cat in data.get("categories", []): cat_name = cat["name"] old_cat_dict[cat["id"]] = cat_name if cat_name not in merged_categories: # 새로 발견한 클래스이면, 새로운 id를 부여하여 병합 merged_categories[cat_name] = { "id": new_cat_id, "name": cat_name, "supercategory": cat.get("supercategory", "none") } new_cat_id += 1 # 어노테이션 매핑: image_id -> list of annotations imgid_to_anns = defaultdict(list) for ann in data.get("annotations", []): imgid_to_anns[ann["image_id"]].append(ann) # 이미지 및 어노테이션 처리: 각 이미지의 id를 새롭게 부여하고, # 어노테이션의 image_id와 category_id도 새롭게 매핑 for img in data.get("images", []): # 새 이미지 항목 생성 new_img = { "id": new_img_id, "file_name": img["file_name"], # 파일명은 절대경로 그대로 사용 "width": img["width"], "height": img["height"] } merged_images.append(new_img) # 해당 이미지에 대한 어노테이션 처리 anns = imgid_to_anns.get(img["id"], []) for ann in anns: old_cat_id = ann["category_id"] cat_name = old_cat_dict.get(old_cat_id, None) if cat_name is None: continue # 해당 카테고리가 없으면 건너뜀 new_ann = ann.copy() new_ann["id"] = new_ann_id new_ann["image_id"] = new_img_id # 병합된 카테고리 id로 업데이트 new_ann["category_id"] = merged_categories[cat_name]["id"] merged_annotations.append(new_ann) new_ann_id += 1 new_img_id += 1 # 최종 merged categories: dictionary의 value들을 리스트로 변환 merged_categories_list = list(merged_categories.values()) print(f"✅ 병합 완료: 이미지 {len(merged_images)}개, 어노테이션 {len(merged_annotations)}개, 클래스 {len(merged_categories_list)}개") return { "images": merged_images, "annotations": merged_annotations, "categories": merged_categories_list, "info": {}, # 필요 시 info 항목 추가 "licenses": [] # 필요 시 licenses 항목 추가 } def clip_bbox(bbox, img_w, img_h): """ bbox: [x, y, w, h] 형식. 이미지 범위 내로 bbox 좌표를 자릅니다. """ x, y, w, h = bbox x = max(0, x) y = max(0, y) w = min(w, img_w - x) h = min(h, img_h - y) return [x, y, w, h] # ============================================================================= # [Cell 3] 메인 실행 루프: 통합 데이터셋 병합 및 Offline Resize Dataset 생성 # ============================================================================= def main(): # 1) 출력 폴더 생성 (기존 파일 덮어쓰지 않도록 필요 시 exp, exp2 등으로 이름 변경) os.makedirs(OUT_IMG_DIR, exist_ok=True) print("✅ Output 이미지 폴더 생성 완료:", os.path.abspath(OUT_IMG_DIR)) # 2) 입력 JSON 파일들을 병합하여 하나의 통합 데이터셋 생성 merged_coco = merge_coco_datasets(INPUT_JSON_PATHS) merged_images = merged_coco.get("images", []) merged_annotations = merged_coco.get("annotations", []) merged_categories = merged_coco.get("categories", []) # image_id 별 annotation 매핑 (merged dataset) imgid_to_anns = defaultdict(list) for ann in merged_annotations: imgid_to_anns[ann["image_id"]].append(ann) # 새로운 COCO 구성 요소 (offline resize 결과) new_images = [] new_annotations = [] # 카테고리는 병합된 것을 그대로 사용 new_categories = merged_categories # offline resize 시 새로운 id 부여를 위한 counter new_img_id_counter = 1 new_ann_id_counter = 1 # 3) 각 이미지에 대해 offline resize 수행 (여러 스케일 적용) for iminfo in tqdm(merged_images, desc="Offline Resize Dataset"): orig_img_id = iminfo["id"] file_name = iminfo["file_name"] # file_name은 절대경로 orig_w = iminfo["width"] orig_h = iminfo["height"] if not os.path.exists(file_name): print("❗ 파일 없음:", file_name) continue img = cv2.imread(file_name) if img is None: print("❗ 이미지 로드 실패:", file_name) continue # 이미지 크기가 JSON 정보와 일치하는지 확인 if img.shape[1] != orig_w or img.shape[0] != orig_h: print("❗ 이미지 크기 불일치:", file_name) continue # 해당 이미지에 속한 annotation 정보 anns = imgid_to_anns.get(orig_img_id, []) # 원본 bbox 리스트와 카테고리 id 리스트 (COCO 형식: [x, y, w, h]) bboxes_xywh = [] cat_ids = [] for ann in anns: bboxes_xywh.append(ann["bbox"]) cat_ids.append(ann["category_id"]) # 각 스케일 팩터마다 offline resize 수행 for scale in SCALE_FACTORS: newW = int(round(orig_w * scale)) newH = int(round(orig_h * scale)) if newW < 2 or newH < 2: print(f"❗ 스케일 {scale}: 너무 작은 사이즈 ({newW}x{newH}), 스킵") continue # ① 이미지 리사이즈 (cv2.resize는 (width, height) 순서) resized_img = cv2.resize(img, (newW, newH)) # ② bbox 좌표 스케일 적용 new_bboxes = [] for bbox in bboxes_xywh: x0, y0, w0, h0 = bbox nx = x0 * scale ny = y0 * scale nw = w0 * scale nh = h0 * scale new_bboxes.append([nx, ny, nw, nh]) # ③ 저장할 파일명 생성 # 원본 file_name에서 폴더와 파일명을 분리한 후 스케일 태그 추가 # 예: "/DATA2/ltb/si/it1_worker_helmet/images/2024-12-04 .../00251.jpg" => # "00251_2p0.jpg" (2배 확대의 경우) subfolder, basename = os.path.split(file_name) name_, ext_ = os.path.splitext(basename) scale_tag = f"{scale}".replace(".", "p") new_basename = f"{name_}_{scale_tag}{ext_}" # 출력 폴더에서 원본과 동일한 폴더 구조 유지 # 단, 파일이 여러 데이터셋에서 온 경우에도 그대로 저장됨 # (예: 원본은 .../images/..., albu 변형은 .../albu_aug/images/...) relative_folder = os.path.relpath(subfolder, DATASET_ROOT) out_subdir = os.path.join(OUT_IMG_DIR, relative_folder) os.makedirs(out_subdir, exist_ok=True) out_img_path = os.path.join(out_subdir, new_basename) # 이미지 저장 (cv2.imwrite: 저장 성공 시 True 반환) cv2.imwrite(out_img_path, resized_img) abs_out_img_path = os.path.abspath(out_img_path) # ④ COCO image 항목 생성 (새로운 image id 할당) new_img_entry = { "id": new_img_id_counter, "file_name": abs_out_img_path, # 절대경로 기록 "width": newW, "height": newH } new_images.append(new_img_entry) current_img_id = new_img_id_counter new_img_id_counter += 1 # ⑤ 각 bbox에 대해 annotation 항목 생성 for i, bbox in enumerate(new_bboxes): x_, y_, w_, h_ = bbox # 너무 작은 bbox는 건너뛰기 if w_ < 1 or h_ < 1: continue ann_entry = { "id": new_ann_id_counter, "image_id": current_img_id, "category_id": cat_ids[i], "bbox": [float(x_), float(y_), float(w_), float(h_)], "area": float(w_ * h_), "iscrowd": 0, "segmentation": [] # 빈 segmentation 리스트 } new_annotations.append(ann_entry) new_ann_id_counter += 1 # 4) 최종 COCO JSON 데이터 구성 (병합된 카테고리와 offline resize 결과 이미지/어노테이션) out_coco = { "info": merged_coco.get("info", {}), "licenses": merged_coco.get("licenses", []), "categories": new_categories, "images": new_images, "annotations": new_annotations } # 5) 최종 JSON 저장 save_coco(out_coco, OUT_JSON_PATH) print("✅ Offline Resize 및 병합 작업 완료!") print("✅ 결과 이미지 폴더:", os.path.abspath(OUT_IMG_DIR)) if __name__ == "__main__": main()