d

RobinNet/nets/Unet.py

@@ -115,7 +115,7 @@ def Unet(x, weights, filter_multiplier, training=True):
 def train_unet(model, inputs, gt, weights, optimizer, filter_multiplier):
     with tf.GradientTape() as tape:
         pred = model(inputs, weights, filter_multiplier, training=True)
-        fig = plt.figure()
+        '''fig = plt.figure()
         fig.add_subplot(2, 3, 1)
         plt.imshow(inputs[1, :, :, 0], cmap=plt.cm.bone)
         plt.title('Input')
@@ -125,7 +125,7 @@ def train_unet(model, inputs, gt, weights, optimizer, filter_multiplier):
         fig.add_subplot(2, 3, 3)
         plt.imshow(gt[1,:,:,0], cmap=plt.cm.bone)
         plt.title('True')
-        plt.show()
+        plt.show()'''
 
         current_loss = dice_loss(pred, gt, axis=(1, 2, 3))
         grads = tape.gradient(current_loss, weights)