DICOM series에 object detection결과 그리기

워낙에 좋은 코드들이 많지만, 내가 사용하는 방법.

사실 결과에 따라서 항상 새로 코드를 짜는 편인데, 

이 방법은 ROI의 위치 확인용도로 사용한다.

 

딥러닝 결괏 값은 text파일에 (x1, y1, z1, x2, y2, z2) 이 순서로 좌표값을 저장하였다.

 

 

일반 jpg파일로 여러 장 저장할 수 있지만, case도 많고 1 case에 몇백 장씩 있으면 데이터 관리가 힘들다.

그래서 1개의 파일에 결과를 보기위해서 파일 포맷을 찾다가 nifti는 color값을 넣을 수가 없어서

(color값도 넣고 볼수있을수도 있지만, 난 도저히 찾아봐도 되지가 않았다)

 

그래서. mha 파일 포맷으로 저장하여 ImageJ를 이용하여 열어서 확인하는 방식으로 구성하였음.

 

import numpy as np
import SimpleITK as sidk
import os


# 실험에 사용한 window값을 가져왔는데 이 코드에서는 winodws[0]번째만 사용하였음
windows = [(500, 200), (700, 400), (1200, 400)]
def adjust_window(image, window):
    width = window[0]
    level = window[1]
    upper = level + width / 2
    lower = level - width / 2
    copied_image = image.clip(lower, upper)
    copied_image = copied_image - lower
    return (copied_image / (upper - lower)).astype('float32')

def get_matrix_from_txt(txt_path):
    f = open(txt_path, 'r')
    txt_matrix = list()
    while True:
        line = f.readline()
        if not line: break
        line = line.strip()
        line = line.split(',')
        line = list(map(int, line))
        txt_matrix.append(line)
    del line
    f.close()
    return txt_matrix

# +-1은 1pixel만 그리면 너무 작게 그려질까 걱정이 되서, 추가한 pixel임
# opencv를 사용해서 rect를 그리는 방법도 있는데, 그렇게 하면
# volume에서 이미지를 가져와서 그리고 다시 변환해줘야하는 번거로움이 있어서 이렇게 짯었다.
def draw_rectangle(matrix, image_vol, color):
    output_vol = image_vol.copy()
    for mat in matrix:
        z = mat[2]
        while z < mat[5]:
            x = mat[0]
            while x < mat[3]:
                for i in range(0, 3):
                    try:
                        output_vol[i][z][x][mat[1] - 1] = color[i]
                        output_vol[i][z][x][mat[1]] = color[i]
                        output_vol[i][z][x][mat[1] + 1] = color[i]
                    except:
                        print(mat[1])
                    try:
                        output_vol[i][z][x][mat[4] - 1] = color[i]
                        output_vol[i][z][x][mat[4]] = color[i]
                        output_vol[i][z][x][mat[4] + 1] = color[i]
                    except:
                        print(mat[4])
                x += 1
            y = mat[1]
            while y < mat[4]:
                for i in range(0, 3):
                    try:
                        output_vol[i][z][mat[0] - 1][y] = color[i]
                        output_vol[i][z][mat[0]][y] = color[i]
                        output_vol[i][z][mat[0] + 1][y] = color[i]
                    except:
                        print(mat[0])
                    try:
                        output_vol[i][z][mat[3] - 1][y] = color[i]
                        output_vol[i][z][mat[3]][y] = color[i]
                        output_vol[i][z][mat[3] + 1][y] = color[i]
                    except:
                        print(mat[3])
                y += 1
            z += 1

    return output_vol


def draw_numpy_gt_and_deep_result(img_dir, save_dir, gt, deep):
    deep_list_file = os.listdir(root_folder_deep)
    deep_list = [file for file in deep_list_file if file.endswith(".txt")]
    deep_list.sort()

    gt_file = os.listdir(gt)
    gt_list = [file for file in gt_file if file.endswith(".txt")]
    gt_list.sort()
	
    #python에서는 file을 읽어온 후 꼭 정렬이 필요함.

    for i in deep_list:
        deep_path = root_folder_deep + i
        deep_name = i.replace(root_folder_deep, '')
        img_path = img_dir + i.replace('.txt', '.npy')
        if not os.path.isfile(img_path):
            continue
        
        # DICOM을 매번 불러오는 것은 itk사용도 해야하고 시간이 많이걸려서
        # 학습데이터 만들때, 원본 값을 그대로 저장한 후에 필요할 때 이것만 불러와서 사용한다.
        img_vol = np.load(img_path)
        img = adjust_window(img_vol, windows[0])
        
        # rectangle 종류를 구분하기 위해서 color 채널값으로 만들어줘야하므로, stack을 하여
        # 3 channel로 만들어주는 과정
        img_3channel = np.stack((img, img, img), axis=0)
        for j in json_list:
            json_path = json_gt + j
            human_path = human_gt + j
            gt_name = j.replace(json_gt, '')
            if os.path.isfile(human_path) and deep_name == gt_name:
                print(deep_name)
                deep_matrix = get_matrix_from_txt(deep_path)
                gt_matrix = get_matrix_from_txt(json_path)
                red = [1, 0, 0]
                green = [0, 1, 0]
                blue = [0, 0, 1]
                rect_img = draw_rectangle(deep_matrix, img_3channel, red)
                rect_img = draw_rectangle(gt_matrix, rect_img, green)
                
                # float type 그대로 저장하면 용량이 너무 많이 차지하므로,
                # 8bit로 바꿔줘서 저장한다.
                rect_img *= 255.0
                rect_img = rect_img.astype('uint8')
                
                # 무식하게 channel값을 계속 뒤로 미뤄주는데
                # 혹시나 실수하여 x,y 좌표값이 바뀌거나 하는 문제 때문에 이렇게 사용함
                # 꼭 이렇게 할필요는 없음 다만 mha로 저장할때는 
                # 맨 마지막 shape처럼 나오도록해야 함
                # rect_img.shape[3,z,x,y]
                rect_img = np.moveaxis(rect_img, 0, 1)
                # rect_img.shape[z,3,x,y]
                rect_img = np.moveaxis(rect_img, 1, 2)
                # rect_img.shape[z,x,3,y]
                rect_img = np.moveaxis(rect_img, 2, 3)
                # rect_img.shape[z,x,y,3]
                
                output_img = sitk.GetImageFromArray(rect_img)
                output_path = save_dir + i.replace('.txt', '.mha')
                
                # sitk를 이용하여 저장한다. itk도 사용가능하지만, 귀찮아서 sitk를 사용했음
                sitk.WriteImage(output_img, output_path)
                break

 

 

내 코드를 보니 이 그림이 떠올랐음.....