def test_model_3d_att(): # verifying if 3d attention model can be created a = [[[[[0 for i in range(48)] for j in range(48)] for k in range(16)]]] inp = torch.tensor(a).float() model = imed_model.Modified3DUNet(in_channel=1, out_channel=1, attention=True) inf = model(inp) assert(type(inf) == torch.Tensor)
def test_create_model_unet3d(): model = imed_models.Modified3DUNet(in_channel=1, out_channel=1, n_filters=8, attention=True) torch.save(model, "model_unet_3d_test.pt") os.makedirs("3d_test") command = "cp model_unet_3d_test.pt 3d_test/best_model.pt" subprocess.check_output(command, shell=True)
def test_onnx(download_data_testing_test_files): model = imed_models.Modified3DUNet(1, 1) if not PATH_MODEL.exists(): PATH_MODEL.mkdir() torch.save(model, PATH_MODEL_PT) img = nib.load(IMAGE_PATH).get_fdata().astype('float32')[:16, :64, :32] # Add batch and channel dimensions img_tensor = torch.tensor(img).unsqueeze(0).unsqueeze(0) dummy_input = torch.randn(1, 1, 32, 32, 32) imed_utils.save_onnx_model(model, dummy_input, str(PATH_MODEL_ONNX)) model = torch.load(PATH_MODEL_PT) model.eval() out_pt = model(img_tensor).detach().numpy() out_onnx = imed_inference.onnx_inference(str(PATH_MODEL_ONNX), img_tensor).numpy() shutil.rmtree(PATH_MODEL) assert np.allclose(out_pt, out_onnx, rtol=1e-3)
def test_segment_volume_3d(download_functional_test_files, center_crop): model = imed_models.Modified3DUNet(in_channel=1, out_channel=1, base_n_filter=1) if not os.path.exists(PATH_MODEL): os.mkdir(PATH_MODEL) torch.save(model, os.path.join(PATH_MODEL, "model_test.pt")) config = { "Modified3DUNet": { "applied": True, "length_3D": LENGTH_3D, "stride_3D": LENGTH_3D, "attention": False }, "loader_parameters": { "slice_filter_params": { "filter_empty_mask": False, "filter_empty_input": False }, "roi_params": { "suffix": None, "slice_filter_roi": None }, "slice_axis": "sagittal" }, "transformation": { "Resample": { "wspace": 1, "hspace": 1, "dspace": 2 }, "CenterCrop": { "size": center_crop }, "RandomTranslation": { "translate": [0.03, 0.03], "applied_to": ["im", "gt"], "dataset_type": ["training"] }, "NumpyToTensor": {}, "NormalizeInstance": { "applied_to": ["im"] } }, "postprocessing": {}, "training_parameters": { "batch_size": BATCH_SIZE } } PATH_CONFIG = os.path.join(PATH_MODEL, 'model_test.json') with open(PATH_CONFIG, 'w') as fp: json.dump(config, fp) nib_lst, _ = imed_inference.segment_volume(PATH_MODEL, [IMAGE_PATH]) nib_img = nib_lst[0] assert np.squeeze(nib_img.get_fdata()).shape == nib.load(IMAGE_PATH).shape assert (nib_img.dataobj.max() <= 1.0) and (nib_img.dataobj.min() >= 0.0) assert nib_img.dataobj.dtype == 'float32' shutil.rmtree(PATH_MODEL)