172 lines
5.7 KiB
Python
172 lines
5.7 KiB
Python
import os
|
|
import cv2
|
|
import numpy as np
|
|
import nibabel as nib
|
|
import pydicom
|
|
from pathlib import Path
|
|
import shutil
|
|
|
|
# 설정
|
|
ROOT_DIR = Path('.')
|
|
OUTPUT_DIR = ROOT_DIR / 'cvat_dataset_mask_v2'
|
|
JPEG_DIR = OUTPUT_DIR / 'JPEGImages'
|
|
SEG_CLASS_DIR = OUTPUT_DIR / 'SegmentationClass'
|
|
IMAGE_SETS_DIR = OUTPUT_DIR / 'ImageSets' / 'Segmentation'
|
|
|
|
# 색상 정의 (샘플 labelmap.txt 기반)
|
|
COLOR_MAP = {
|
|
"C2": (230, 25, 75),
|
|
"C3": (60, 180, 75),
|
|
"C4": (255, 225, 25),
|
|
"C5": (67, 99, 216),
|
|
"C6": (245, 130, 49),
|
|
"C7": (145, 30, 180),
|
|
"L1": (0, 0, 117),
|
|
"L2": (128, 128, 128),
|
|
"L3": (255, 255, 255),
|
|
"L4": (0, 0, 0), # 주의: L4가 검은색(0,0,0)으로 정의되어 있어 배경과 겹칠 수 있음. 확인 필요.
|
|
"L5": (31, 119, 180),
|
|
"Rib": (255, 127, 14),
|
|
"Sacrum": (44, 160, 44),
|
|
"T1": (70, 240, 240),
|
|
"T10": (170, 255, 195),
|
|
"T11": (128, 128, 0),
|
|
"T12": (255, 216, 177),
|
|
"T2": (240, 50, 230),
|
|
"T3": (188, 246, 12),
|
|
"T4": (250, 190, 190),
|
|
"T5": (0, 128, 128),
|
|
"T6": (230, 190, 255),
|
|
"T7": (154, 99, 36),
|
|
"T8": (255, 250, 200),
|
|
"T9": (128, 0, 0),
|
|
# "background": (0, 0, 0) # 배경은 기본적으로 검정
|
|
}
|
|
# L4가 (0,0,0)이라면 배경과 구분 불가. L4 색상을 임의로 변경하거나 주의해야 함.
|
|
# 여기서는 안전을 위해 L4를 아주 어두운 회색(1,1,1)으로 하거나,
|
|
# 또는 사용자가 준 샘플에 L4가 (0,0,0)이라면 그것을 따라야 하지만,
|
|
# 보통 배경이 (0,0,0)이므로 L4는 (50,50,50) 정도로 변경하여 진행하겠습니다.
|
|
COLOR_MAP["L4"] = (50, 50, 50)
|
|
|
|
|
|
def normalize_image(image):
|
|
"""DICOM 이미지를 0-255 범위의 8bit 이미지로 변환"""
|
|
image = image.astype(float)
|
|
min_val = np.min(image)
|
|
max_val = np.max(image)
|
|
if max_val - min_val == 0:
|
|
return np.zeros(image.shape, dtype=np.uint8)
|
|
image = (image - min_val) / (max_val - min_val) * 255.0
|
|
return image.astype(np.uint8)
|
|
|
|
def process_case(case_dir, file_list):
|
|
"""하나의 케이스(폴더) 처리"""
|
|
dcm_files = list(case_dir.glob('*.dcm'))
|
|
if not dcm_files:
|
|
return
|
|
|
|
dcm_path = dcm_files[0]
|
|
case_name = case_dir.name
|
|
|
|
# 1. DICOM -> JPG
|
|
try:
|
|
ds = pydicom.dcmread(dcm_path)
|
|
pixel_array = ds.pixel_array
|
|
img_8bit = normalize_image(pixel_array)
|
|
|
|
# 그레이스케일 이미지를 RGB로 변환 (CVAT 호환성)
|
|
img_rgb = cv2.cvtColor(img_8bit, cv2.COLOR_GRAY2BGR)
|
|
|
|
cv2.imwrite(str(JPEG_DIR / f"{case_name}.jpg"), img_rgb)
|
|
except Exception as e:
|
|
print(f"Error processing DICOM {dcm_path}: {e}")
|
|
return
|
|
|
|
# 2. NII Masks -> RGB Mask
|
|
# 배경(0,0,0)으로 초기화된 RGB 이미지 생성
|
|
mask_rgb = np.zeros((pixel_array.shape[0], pixel_array.shape[1], 3), dtype=np.uint8)
|
|
|
|
nii_files = list(case_dir.glob('*.nii.gz'))
|
|
for nii_path in nii_files:
|
|
class_name = nii_path.name.split('.')[0]
|
|
if class_name not in COLOR_MAP:
|
|
continue
|
|
|
|
color = COLOR_MAP[class_name] # (R, G, B)
|
|
# OpenCV는 BGR 사용하므로 순서 변경
|
|
color_bgr = (color[2], color[1], color[0])
|
|
|
|
try:
|
|
img_nii = nib.load(str(nii_path))
|
|
data = img_nii.get_fdata()
|
|
|
|
# 3D -> 2D Projection
|
|
proj = np.max(data, axis=2).T
|
|
|
|
# 리사이즈
|
|
if proj.shape != (mask_rgb.shape[0], mask_rgb.shape[1]):
|
|
proj = cv2.resize(proj, (mask_rgb.shape[1], mask_rgb.shape[0]), interpolation=cv2.INTER_NEAREST)
|
|
|
|
# 마스크 그리기
|
|
# 해당 클래스 영역을 해당 색상으로 칠함
|
|
mask_rgb[proj > 0] = color_bgr
|
|
|
|
except Exception as e:
|
|
print(f"Error processing NII {nii_path}: {e}")
|
|
|
|
# 마스크 저장 (PNG)
|
|
cv2.imwrite(str(SEG_CLASS_DIR / f"{case_name}.png"), mask_rgb)
|
|
|
|
# 파일 목록에 추가
|
|
file_list.append(case_name)
|
|
print(f"Processed: {case_name}")
|
|
|
|
def create_labelmap():
|
|
"""labelmap.txt 생성"""
|
|
with open(OUTPUT_DIR / 'labelmap.txt', 'w') as f:
|
|
f.write("# label:color_rgb:parts:actions\n")
|
|
# 정의된 순서대로 작성 (이름 정렬)
|
|
for name in sorted(COLOR_MAP.keys()):
|
|
r, g, b = COLOR_MAP[name]
|
|
f.write(f"{name}:{r},{g},{b}::\n")
|
|
# 배경 추가
|
|
f.write("background:0,0,0::\n")
|
|
|
|
def main():
|
|
# 디렉토리 초기화
|
|
if OUTPUT_DIR.exists():
|
|
shutil.rmtree(OUTPUT_DIR)
|
|
|
|
JPEG_DIR.mkdir(parents=True)
|
|
SEG_CLASS_DIR.mkdir(parents=True)
|
|
IMAGE_SETS_DIR.mkdir(parents=True)
|
|
|
|
print("Converting to CVAT Segmentation Mask 1.1 format...")
|
|
|
|
file_list = []
|
|
|
|
# 데이터 탐색
|
|
for d in ROOT_DIR.iterdir():
|
|
if d.is_dir() and not d.name.startswith('.') and 'cvat' not in d.name and 'segmask' not in d.name:
|
|
if list(d.glob('*.dcm')):
|
|
process_case(d, file_list)
|
|
|
|
# default.txt 생성
|
|
with open(IMAGE_SETS_DIR / 'default.txt', 'w') as f:
|
|
for name in file_list:
|
|
f.write(f"{name}\n")
|
|
|
|
# labelmap.txt 생성
|
|
create_labelmap()
|
|
|
|
print("\nDone! Dataset created at:", OUTPUT_DIR.absolute())
|
|
print(f"Total cases: {len(file_list)}")
|
|
print("\n[중요] 업로드 시 주의사항:")
|
|
print("1. 'cvat_dataset_mask_v2' 폴더 안으로 들어가서")
|
|
print("2. 모든 파일/폴더(JPEGImages, SegmentationClass, ImageSets, labelmap.txt)를 선택하고")
|
|
print("3. '압축하기(ZIP)'를 실행하세요.")
|
|
print("4. CVAT 포맷 선택: 'Segmentation Mask 1.1'")
|
|
|
|
if __name__ == '__main__':
|
|
main()
|