Switch to GradICON

uncbiag · Dec 8, 2023 · 3877afe · 3877afe
1 parent 13d09bf
commit 3877afe
Showing 1 changed file with 1 addition and 4 deletions.
diff --git a/oai_analysis/registration.py b/oai_analysis/registration.py
@@ -17,10 +17,7 @@ def register(self, fixed_image, moving_image):
 
 class ICON_Registration:
     def __init__(self):
-        self.register_module = icon_registration.pretrained_models.OAI_knees_registration_model(pretrained=True)
-        icon_registration.network_wrappers.adjust_batch_size(self.register_module, 1)
-        self.register_module.cpu()
-        #self.register_module.cuda()
+        self.register_module = icon_registration.pretrained_models.OAI_knees_gradICON_model(pretrained=True)
 
     def register(self, fixed_image, moving_image):
         print("fixed range", np.min(fixed_image), np.max(fixed_image))