add : dcm / nii.gz 뷰어 / cvat 컨버터

DCM_NII_to_CVAT.py

@@ -0,0 +1,171 @@
+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()