convert_data.py 文件源码

def read_images_labels(path):
    # Read the images and labels from a folder containing both dicom files
    for subdir, dirs, files in os.walk(path):
        dcms = glob.glob(os.path.join(subdir, '*.dcm'))
        if len(dcms) == 1:
            structure = dicom.read_file(dcms[0])
            contours = read_structure(structure)
        elif len(dcms) > 1:
            slices = [dicom.read_file(dcm) for dcm in dcms]
            slices.sort(key = lambda x: float(x.ImagePositionPatient[2]))
            images = np.stack([s.pixel_array for s in slices], axis=0).astype(np.float32)
            images = images + slices[0].RescaleIntercept

    images = normalize(images)
    labels = get_labels(contours, images.shape, slices)
    inplane_scale = slices[0].PixelSpacing[0] / PIXEL_SPACING
    inplane_size = int(np.rint(inplane_scale * slices[0].Rows / 2) * 2)
    z_scale = slices[0].SliceThickness / SLICE_THICKNESS
    z_size = int(np.rint(z_scale * images.shape[0]))

    if inplane_size != INPLANE_SIZE or z_scale != 1:
        images = resize(images, (z_size, inplane_size, inplane_size), mode='constant')
        new_labels = np.zeros_like(images, dtype=np.float32)
        for z in range(N_CLASSES):
            roi = resize((labels == z + 1).astype(np.float32), (z_size, inplane_size, inplane_size), mode='constant')
            new_labels[roi >= 0.5] = z + 1
        labels = new_labels
        if inplane_size != INPLANE_SIZE:
            if inplane_size > INPLANE_SIZE:
                crop = int((inplane_size - INPLANE_SIZE) / 2)
                images = images[:, crop : crop + INPLANE_SIZE, crop : crop + INPLANE_SIZE]
                labels = labels[:, crop : crop + INPLANE_SIZE, crop : crop + INPLANE_SIZE]
            else:
                pad = int((INPLANE_SIZE - new_size) / 2)
                images = np.pad(images, ((0, 0), (pad, pad), (pad, pad)), 'constant')
                labels = np.pad(labels, ((0, 0), (pad, pad), (pad, pad)), 'constant')

    return images, labels