From 9dd3f68da87e8c89e6610c327b517779f5822e29 Mon Sep 17 00:00:00 2001 From: Jeremy Hummel Date: Wed, 14 Aug 2019 11:48:57 -0700 Subject: [PATCH] update for only no ct --- models/Model_SAE/Model.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/models/Model_SAE/Model.py b/models/Model_SAE/Model.py index f1165f1..90d0bfa 100644 --- a/models/Model_SAE/Model.py +++ b/models/Model_SAE/Model.py @@ -433,13 +433,19 @@ class SAEModel(ModelBase): self.AE_view = K.function([warped_src, warped_dst], [pred_src_src[-1], pred_dst_dst[-1], pred_src_dst[-1]]) - else: # Updated "pred_src_dst[-1]" to "pred_src_dst[-1][:4]", as it was breaking on images with 4th channel if self.options['learn_mask']: - self.AE_convert = K.function([warped_dst], [pred_src_dst[-1][:4], pred_dst_dstm[-1], pred_src_dstm[-1]]) + if self.options['apply_random_ct']: + params = [pred_src_dst[-1], pred_dst_dstm[-1], pred_src_dstm[-1]] + else: + params = [pred_src_dst[-1][:4], pred_dst_dstm[-1], pred_src_dstm[-1]] else: - self.AE_convert = K.function([warped_dst], [pred_src_dst[-1][:4]]) + if self.options['apply_random_ct']: + params = [pred_src_dst[-1]] + else: + params = [pred_src_dst[-1][:4]] + self.AE_convert = K.function([warped_dst], params) if self.is_training_mode: self.src_sample_losses = []