prepare_for_submission.py 文件源码

def batch_works(k):
    if k == n_processes - 1:
        paths = all_paths[k * int(len(all_paths) / n_processes) : ]
    else:
        paths = all_paths[k * int(len(all_paths) / n_processes) : (k + 1) * int(len(all_paths) / n_processes)]

    for path in paths:
        probs = np.load(os.path.join(input_path, path))
        pred = np.argmax(probs, axis=3)
        fg_prob = 1 - probs[..., 0]
        pred = clean_contour(fg_prob, pred)
        seg = np.zeros(pred.shape, dtype=np.int16)
        seg[pred == 1] = 1
        seg[pred == 2] = 2
        seg[pred == 3] = 4
        img = nib.Nifti1Image(seg, np.eye(4))
        nib.save(img, os.path.join(output_path, path.replace('_probs.npy', '.nii.gz')))