generate_label_mask.py

import json
import numpy as np
from affine import Affine
import rasterio.features
import openslide 
import glob
import os
import cv2
from tqdm import tqdm
import h5py

# feat_path = '/ailab/group/pjlab-medai/zhouxiao/pathology/features/keep/wsi_features/tcga/BRCA/h5_files/'
feat_path = 'test/feat/'
# wsi_path = '/ailab/group/pjlab-medai/zhouxiao/pathology/data/tcga_WSI_data/BRCA/BRCA_test/'
wsi_path = 'test/wsi/'
# anno_path = '/ailab/group/pjlab-medai/zhouxiao/pathology/annotations/tcga/'
anno_path = 'test/anno/'
# save_path = '/ailab/group/pjlab-medai/zhouxiao/pathology/features/keep/wsi_features/tcga/BRCA/h5_with_label/coarse/'
save_path = 'test/save/'
dataset_name = 'BRCA'
# ext_name = '.jpg'

def get_magnification(wsi):
    """
    获取WSI的放大倍数，支持不同厂商的扫描仪
    """
    # 可能包含放大倍数信息的键名列表
    mag_keys = [
        'aperio.AppMag',                      # Aperio
        'openslide.objective-power',          # OpenSlide通用标准键
        'hamamatsu.SourceLens',               # Hamamatsu
        'hamamatsu.NominalMagnification',     # Hamamatsu
        'philips.DICOM_ACQUISITION.OPTICAL_PATH.OBJECTIVE_POWER',  # Philips
        'leica.objective',                    # Leica
        'mirax.GENERAL.OBJECTIVE_MAGNIFICATION',  # Mirax/3DHISTECH
        'ventana.Magnification',              # Ventana/Roche
        'tiff.ImageDescription',              # 可能包含放大倍数信息的TIFF标签
    ]
    
    # 尝试从不同键名中获取放大倍数
    for key in mag_keys:
        if key in wsi.properties:
            try:
                mag_value = wsi.properties[key]
                # 尝试将值转换为数字
                return int(float(mag_value))
            except (ValueError, TypeError):
                # 如果无法转换，继续尝试下一个键
                continue
    
    # 如果找不到放大倍数信息，尝试备用方法：
    # 1. 检查所有属性中是否包含"magnification"或"mag"
    for key, value in wsi.properties.items():
        if 'magnification' in key.lower() or 'mag' in key.lower() or 'objective' in key.lower():
            try:
                return int(float(value))
            except (ValueError, TypeError):
                continue
                
    # 如果仍然找不到，返回None或默认值
    print("警告：无法确定放大倍数，请手动检查WSI属性，按默认40x 512*512")
    return 40

def save_hdf5(output_path, asset_dict, attr_dict= None, mode='a'):
    file = h5py.File(output_path, mode)
    for key, val in asset_dict.items():
        data_shape = val.shape
        if key not in file:
            data_type = val.dtype
            chunk_shape = (1, ) + data_shape[1:]
            maxshape = (None, ) + data_shape[1:]
            dset = file.create_dataset(key, shape=data_shape, maxshape=maxshape, chunks=chunk_shape, dtype=data_type)
            dset[:] = val
            if attr_dict is not None:
                if key in attr_dict.keys():
                    for attr_key, attr_val in attr_dict[key].items():
                        dset.attrs[attr_key] = attr_val
        else:
            dset = file[key]
            dset.resize(len(dset) + data_shape[0], axis=0)
            dset[-data_shape[0]:] = val
    file.close()
    return output_path

def generate_label_mask():
    wsi_list = glob.glob(wsi_path + '/*.svs')
    # wsi_list = ['E:/subtyping/BRCA/IDC/TCGA-A2-A0ER-01Z-00-DX1.18123C1B-A0CC-4957-9373-5FFAB985C2E8.svs']
    for wsi_dir in wsi_list:
        wsi_name = os.path.basename(wsi_dir).split('.svs')[0]
        print(f'Processing {wsi_name}')
        if os.path.exists(save_path + wsi_name + '.h5'):
            print(f'{wsi_name} has been processed')
            continue
        for i, name in enumerate(['IDC', 'ILC']):
            wsi_label = i + 1
            anno_dir = anno_path + dataset_name + '/' + name + '_Anno_coarse/' + wsi_name + '.geojson'
            if  not os.path.exists(anno_dir):
                print(f'{wsi_name} not belongs to {name}')
                continue
            else:
                print(f'{wsi_name} belongs to {name}')
                wsi = openslide.OpenSlide(wsi_dir)
                level0_mag = get_magnification(wsi)
                orig_width, orig_height = wsi.dimensions[0], wsi.dimensions[1]
                if level0_mag == 40:
                    target_width, target_height = wsi.level_dimensions[1][0], wsi.level_dimensions[1][1]
                    scale = round(wsi.level_downsamples[1])
                    patch_size = int(512/scale)
                elif level0_mag == 20:
                    target_width, target_height = wsi.level_dimensions[2][0]*2, wsi.level_dimensions[2][1]*2
                    scale = round(wsi.level_downsamples[0]*2)
                    patch_size = int(256/scale)
                print(f'level0_mag{level0_mag}, orig{(orig_width, orig_height)}, target{(target_width, target_height)}, scale{scale}, patch{patch_size}')
                wsi.close()

                # 计算缩放因子（每个输出像素对应原始多少像素）
                # scale = orig_width / target_width   # = 100

                # 定义仿射变换，将目标像素坐标映射到原始坐标
                # 注意，这里假设原点为 (0,0)
                transform = Affine(scale, 0, 0, 0, scale, 0)

                # 从文件加载 GeoJSON 标注（假设存储于 annotation.json 中）
                with open(anno_dir, 'r') as f:
                    geojson_data = json.load(f)

                # 构造 shapes 列表：
                # 每个元素为 (geometry, 1)，表示该几何区域的像素值为 1（肿瘤区域）
                shapes = [(feature['geometry'], 1) for feature in geojson_data['features']]

                # 利用 rasterio.features.rasterize 直接生成目标大小的 mask
                mask_low = rasterio.features.rasterize(
                    shapes,
                    out_shape=(target_height, target_width),  # 行数（高度），列数（宽度）
                    transform=transform,  # 根据 transform，将输入的原始坐标转换到低分辨率像素空间
                    fill=0,
                    dtype=np.uint8
                )

                # import matplotlib.pyplot as plt
                # plt.imshow(mask_low, cmap='gray')
                # plt.title("Low Resolution Mask")
                # plt.show()
                mask_low = np.array(mask_low, dtype=np.uint8)
                
                mask_recon = np.zeros((target_height, target_width), dtype=np.uint8)
                # print(mask_recon.shape)
                # mask_255 = np.where(mask_low == 1, 255, 0)
                # max_pixel = mask_low.max()
                
                # cv2.imshow('test',mask_low)
                # cv2.waitKey()
                # cv2.imwrite(root_path + dataset_name + '/' + wsi_name + ext_name, mask_255)
                # cv2.imwrite('mask_low.png', mask_255)
                with h5py.File(feat_path + wsi_name + '.h5', 'r') as f:
                    coords = f['coords'][:]
                    features = f['features'][:]
                asset_dict = {'features': features, 'coords': coords}
                patch_labels = []
                for coord in coords:
                    # 获取以coord为左上角的512×512区域
                    x, y = coord / scale
                    x = round(x)
                    y = round(y)
                    # 检查区域是否超出mask_low的边界
                    w = min(patch_size, mask_low.shape[1] - x)
                    h = min(patch_size, mask_low.shape[0] - y)
                    patch = mask_low[y:y+h, x:x+w]
                    # 计算比例
                    ratio = np.sum(patch) / (w*h)
                    # 如果比例大于0.5，标签为1，否则为0
                    if ratio > 0.5:
                        label = wsi_label
                        # mask_recon[y:y+h, x:x+w] = label
                    else:
                        label = 0
                patch_labels.append(label)
                asset_dict['labels'] = np.array(patch_labels)
                attr_dict = {'labels':{'level0_mag':level0_mag, 'patch_size':patch_size, 'scale':scale, 'origin_size':(orig_width, orig_height), 'target_size':(target_width, target_height), 'diagnosis':name}}
                # mask_255 = np.where(mask_recon == 1, 255, 0)
                # cv2.imwrite('mask_recon.png', mask_255)
                save_hdf5(save_path + wsi_name + '.h5', asset_dict, attr_dict)

def check_label_mask(h5_path='TCGA-D8-A1XB-01Z-00-DX2.B262C269-F22B-4306-8195-079C0874EF8E.h5'):
    with h5py.File(h5_path, 'r') as f:
        coords = f['coords'][:]
        labels = f['labels'][:]
        target_size = f['labels'].attrs['target_size']
        scale = f['labels'].attrs['scale']
    mask = np.zeros((target_size[1], target_size[0]), dtype=np.uint8)
    for coord, label in zip(coords, labels):
        x, y = coord / scale
        x = round(x)
        y = round(y)
        patch_size = int(512/scale)
        mask[y:y+patch_size, x:x+patch_size] = int(label)
    
    mask_255 = np.where(mask >= 1, 255, 0)
    print(sum(labels == 1))

    cv2.imwrite('mask.png', mask_255)

if __name__ == '__main__':
    generate_label_mask()