evaluation/RHLV_quantification.py

import json
import numpy as np
import os
import nibabel as nib
import matplotlib.pyplot as plt
import cv2
import pandas as pd
from sklearn.model_selection import ParameterGrid

def rotate_image_to_horizontal(binary_image):
    """
    Rotates the image to make the major axis of the object horizontal.
    """
    # 寻找轮廓
    binary_image = binary_image.astype(np.uint8)
    contours, _ = cv2.findContours(binary_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # 假设最大的轮廓是椎体
    contour = max(contours, key=cv2.contourArea)

    # 获取轮廓的最小外接矩形
    rect = cv2.minAreaRect(contour)
    box = cv2.boxPoints(rect)
    box = np.int0(box)

    # 计算旋转角度
    angle = rect[2]
    if angle < -45:
        angle += 90
    if angle > 45:
        angle-=90

    # 旋转图像
    (h, w) = binary_image.shape[:2]
    center = (w // 2, h // 2)
    M = cv2.getRotationMatrix2D(center, angle, 1.0)
    rotated_image = cv2.warpAffine(binary_image, M, (w, h), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_CONSTANT)

    return rotated_image

def calculate_heights(segmentation_fake,segmentation_label,height_threshold):
    all_heights_fake = []
    all_heights_label = []
    pre_heights_fake = []
    pre_heights_label = []
    mid_heights_fake = []
    mid_heights_label = []
    post_heights_fake = []
    post_heights_label = []
    # 遍历z轴上的每个层
    for z in range(segmentation_label.shape[2]):
        if np.any(segmentation_label[:, :, z]) and np.any(segmentation_fake[:, :, z]):
            segmentation_label_slice = segmentation_label[:, :, z]
            segmentation_fake_slice = segmentation_fake[:, :, z]
            
            loc = np.where(segmentation_fake_slice)[1]
            y_min = int(np.min(loc))
            y_max = int(np.max(loc))
            y_range = y_max-y_min
            one_third_y = int(y_min + y_range/3)
            two_third_y = int(y_min + 2*y_range/3)
            center_height_fake = np.count_nonzero(segmentation_fake_slice[:, int(np.mean(loc))])
            all_height_fake = np.count_nonzero(segmentation_fake_slice, axis=0)
            pre_height_fake = np.count_nonzero(segmentation_fake_slice[:,:one_third_y], axis=0)
            mid_height_fake = np.count_nonzero(segmentation_fake_slice[:,one_third_y:two_third_y], axis=0)
            post_height_fake = np.count_nonzero(segmentation_fake_slice[:, two_third_y:], axis=0)
            
            loc = np.where(segmentation_label[:, :, z])[1]
            center_height_label = np.count_nonzero(segmentation_label_slice[:, int(np.mean(loc))])
            all_height_label = np.count_nonzero(segmentation_label_slice, axis=0)
            pre_height_label = np.count_nonzero(segmentation_label_slice[:, :one_third_y], axis=0)
            mid_height_label = np.count_nonzero(segmentation_label_slice[:, one_third_y:two_third_y], axis=0)
            post_height_label = np.count_nonzero(segmentation_label_slice[:, two_third_y:], axis=0)
            

            all_scale_ratio = 1
            pre_scale_ratio = 1
            mid_scale_ratio = 1
            post_scale_ratio = 1
            if all_height_label.size > 0 and all_height_fake.size > 0:
                if all_height_label.max()>all_height_fake.max():
                    all_scale_ratio = all_height_label.max()/(all_height_fake.max()+1e-6)
            if pre_height_label.size > 0 and pre_height_fake.size > 0:
                if pre_height_label.max()>pre_height_fake.max():
                    pre_scale_ratio = pre_height_label.max()/(pre_height_fake.max()+1e-6)
            if mid_height_label.size > 0 and mid_height_fake.size > 0:
                if mid_height_label.max()>mid_height_fake.max():
                    mid_scale_ratio = mid_height_label.max()/(mid_height_fake.max()+1e-6)
            if post_height_label.size > 0 and post_height_fake.size > 0:
                if post_height_label.max()>post_height_fake.max():
                    post_scale_ratio = post_height_label.max()/(post_height_fake.max()+1e-6)
                
            all_height_fake = all_height_fake*all_scale_ratio
            center_height_fake = center_height_fake*all_scale_ratio
            pre_height_fake = pre_height_fake*pre_scale_ratio
            mid_height_fake = mid_height_fake*mid_scale_ratio
            post_height_fake = post_height_fake*post_scale_ratio
                
            all_heights_fake.extend(all_height_fake[all_height_fake > (center_height_fake*height_threshold)])  # 仅添加非零高度
            all_heights_label.extend(all_height_label[all_height_label > (center_height_label*height_threshold)])  # 仅添加非零高度
            pre_heights_fake.extend(pre_height_fake[pre_height_fake > (center_height_fake*height_threshold)])  # 仅添加非零高度
            pre_heights_label.extend(pre_height_label[pre_height_label > (center_height_label*height_threshold)])  # 仅添加非零高度
            mid_heights_fake.extend(mid_height_fake[mid_height_fake > (center_height_fake*height_threshold)])  # 仅添加非零高度
            mid_heights_label.extend(mid_height_label[mid_height_label > (center_height_label*height_threshold)])  # 仅添加非零高度
            post_heights_fake.extend(post_height_fake[post_height_fake > (center_height_fake*height_threshold)])  # 仅添加非零高度
            post_heights_label.extend(post_height_label[post_height_label > (center_height_label*height_threshold)])  # 仅添加非零高度
    
    # 将heights转换为numpy数组以便使用numpy的功能
    all_heights_fake = np.array(all_heights_fake)
    all_heights_label = np.array(all_heights_label)
    pre_heights_fake = np.array(pre_heights_fake)
    pre_heights_label = np.array(pre_heights_label)
    mid_heights_fake = np.array(mid_heights_fake)
    mid_heights_label = np.array(mid_heights_label)
    post_heights_fake = np.array(post_heights_fake)
    post_heights_label = np.array(post_heights_label)

    return all_heights_fake, all_heights_label,pre_heights_fake, pre_heights_label,mid_heights_fake, mid_heights_label,post_heights_fake, post_heights_label


def calculate_rhlv(segmentation_fake, segmentation_label, center_z, length,vertebra,height_threshold):
    """
    Calculate the Relative Height Loss Value (RHLV) between fake and label segmentations.
    """
    seg_fake_filtered = segmentation_fake[:, :, center_z-length:center_z+length]
    seg_label_filtered = segmentation_label[:, :, center_z-length:center_z+length]

    all_heights_fake, all_heights_label,pre_heights_fake, pre_heights_label,mid_heights_fake, mid_heights_label,post_heights_fake, post_heights_label\
        = calculate_heights(seg_fake_filtered, seg_label_filtered,height_threshold)
    all_height_fake = np.mean(all_heights_fake) if all_heights_fake.size > 0 else 0
    all_height_label = np.mean(all_heights_label) if all_heights_label.size > 0 else 0
    pre_height_fake = np.mean(pre_heights_fake) if pre_heights_fake.size > 0 else 0
    pre_height_label = np.mean(pre_heights_label) if pre_heights_label.size > 0 else 0
    mid_height_fake = np.mean(mid_heights_fake) if mid_heights_fake.size > 0 else 0
    mid_height_label = np.mean(mid_heights_label) if mid_heights_label.size > 0 else 0
    post_height_fake = np.mean(post_heights_fake) if post_heights_fake.size > 0 else 0
    post_height_label = np.mean(post_heights_label) if post_heights_label.size > 0 else 0
    
    all_rhlv = (all_height_fake - all_height_label) / (all_height_fake +1e-6)
    pre_rhlv = (pre_height_fake - pre_height_label) / (pre_height_fake +1e-6)
    mid_rhlv = (mid_height_fake - mid_height_label) / (mid_height_fake +1e-6)
    post_rhlv = (post_height_fake - post_height_label) / (post_height_fake +1e-6)
    min_height = min(pre_height_label,mid_height_label,post_height_label)
    max_height = max(pre_height_label,mid_height_label,post_height_label)
    relative_height_label = min_height/(max_height+1e-6)

    return all_rhlv,pre_rhlv,mid_rhlv,post_rhlv,relative_height_label

def process_datasets_to_excel(dataset_info, label_folder, fake_folder, output_file,length_divisor=5, height_threshold=0.64):
    results = []
    for dataset_type, data in dataset_info.items():
        for vertebra, label in data.items():
            label_path = os.path.join(label_folder, vertebra + '.nii.gz')
            fake_path = os.path.join(fake_folder, vertebra + '.nii.gz')

            if not os.path.exists(label_path) or not os.path.exists(fake_path):
                continue

            segmentation_label_temp = nib.load(label_path).get_fdata()
            segmentation_label = np.zeros_like(segmentation_label_temp)
            
            segmentation_fake_temp = nib.load(fake_path).get_fdata()
            segmentation_fake = np.zeros_like(segmentation_fake_temp)
            
            label_index = int(vertebra.split('_')[-1])
            segmentation_label[segmentation_label_temp == label_index] = 1
            segmentation_fake[segmentation_fake_temp == label_index] = 1
            
            loc = np.where(segmentation_label)[2]
            if loc.size == 0:
                continue  # Skip if no label index found

            min_z = np.min(loc)
            max_z = np.max(loc)
            center_z = int(np.mean(loc))
            length = (max_z - min_z) // length_divisor  # Divisor adjusted based on your setup
            

            all_rhlv, pre_rhlv, mid_rhlv, post_rhlv, relative_height_label = calculate_rhlv(
                segmentation_fake, segmentation_label, center_z, length, vertebra,height_threshold
            )
            results.append({
                "Vertebra": vertebra,
                "Label": label,
                "Dataset": dataset_type,
                "All RHLV": all_rhlv,
                "Pre RHLV": pre_rhlv,
                "Mid RHLV": mid_rhlv,
                "Post RHLV": post_rhlv,
                "Relative Height Label": relative_height_label
            })

    # Create a DataFrame from results and save to Excel
    df = pd.DataFrame(results)
    df.to_excel(output_file, index=False)

def main():
    with open('vertebra_data_local.json', 'r') as file:
       json_data = json.load(file)
    
    label_folder = '/dssg/home/acct-milesun/zhangqi/Dataset/HealthiVert_straighten/label'
    output_folder = '/dssg/home/acct-milesun/zhangqi/Project/HealthiVert-GAN_eval/output'
    result_folder = '/dssg/home/acct-milesun/zhangqi/Project/HealthiVert-GAN_eval/evaluation/RHLV_quantification'
    for root, dirs, files in os.walk(output_folder):
        for dir in dirs:
            exp_folder = os.path.join(root,dir)
            fake_folder = os.path.join(exp_folder,'label_fake')
            result_file = os.path.join(result_folder,dir+'.xlsx')
            process_datasets_to_excel(json_data, label_folder, fake_folder, result_file, length_divisor=5, height_threshold=0.7)

if __name__ == "__main__":
    main()