Spaces:
Sleeping
Sleeping
refactored
Browse files- lungtumormask/mask.py +2 -3
lungtumormask/mask.py
CHANGED
@@ -11,7 +11,6 @@ def load_model():
|
|
11 |
gpu_device = T.device('cpu')
|
12 |
model = UNet_double(3, 1, 1, tuple([64, 128, 256, 512, 1024]), tuple([2 for i in range(4)]), num_res_units = 0)
|
13 |
state_dict = T.hub.load_state_dict_from_url("https://github.com/VemundFredriksen/LungTumorMask/releases/download/0.0/dc_student.pth", progress=True, map_location=gpu_device)
|
14 |
-
#model.load_state_dict(T.load("D:\\OneDrive\\Skole\\Universitet\\10. Semester\\Masteroppgave\\bruk_for_full_model.pth", map_location="cuda:0"))
|
15 |
model.load_state_dict(state_dict)
|
16 |
model.eval()
|
17 |
return model
|
@@ -28,8 +27,8 @@ def mask(image_path, save_path):
|
|
28 |
right = model(preprocess_dump['right_lung']).squeeze(0).squeeze(0).detach().numpy()
|
29 |
|
30 |
print("Post-processing image...")
|
31 |
-
|
32 |
|
33 |
print(f"Storing segmentation at {save_path}")
|
34 |
-
nimage = nibabel.Nifti1Image(
|
35 |
nibabel.save(nimage, save_path)
|
|
|
11 |
gpu_device = T.device('cpu')
|
12 |
model = UNet_double(3, 1, 1, tuple([64, 128, 256, 512, 1024]), tuple([2 for i in range(4)]), num_res_units = 0)
|
13 |
state_dict = T.hub.load_state_dict_from_url("https://github.com/VemundFredriksen/LungTumorMask/releases/download/0.0/dc_student.pth", progress=True, map_location=gpu_device)
|
|
|
14 |
model.load_state_dict(state_dict)
|
15 |
model.eval()
|
16 |
return model
|
|
|
27 |
right = model(preprocess_dump['right_lung']).squeeze(0).squeeze(0).detach().numpy()
|
28 |
|
29 |
print("Post-processing image...")
|
30 |
+
inferred = post_process(left, right, preprocess_dump).astype("uint8")
|
31 |
|
32 |
print(f"Storing segmentation at {save_path}")
|
33 |
+
nimage = nibabel.Nifti1Image(inferred, preprocess_dump['org_affine'])
|
34 |
nibabel.save(nimage, save_path)
|