fix: Left some images out of training

Author: augustoicaro

train.py

@@ -25,12 +25,10 @@
 from keras.layers.advanced_activations import LeakyReLU
 from keras.models import Model
 from math import sqrt, ceil
-from tqdm import tqdm_notebook as tqdm
 import cv2
 from sklearn.utils import shuffle
 from tqdm import tqdm
 import tifffile as tif
-from model import msrf
 from model import *
 from tensorflow.keras.callbacks import *
 import skimage.io
@@ -90,7 +88,7 @@
 
 def train(epochs, batch_size,output_dir, model_save_dir):
 
-    batch_count = int(len(train_x) / batch_size)
+    batch_count = ceil(len(train_x) / batch_size)
     max_val_dice= -1
     G = msrf()
     G.summary()
@@ -99,7 +97,7 @@ def train(epochs, batch_size,output_dir, model_save_dir):
     for e in range(1, epochs+1):
         print ('-'*15, 'Epoch %d' % e, '-'*15,batch_size)
         #sp startpoint
-        for sp in range(0,batch_count,1):
+        for sp in tqdm(range(0,batch_count,1), desc="Batches"):
             if (sp+1)*batch_size>len(train_x):
                 batch_end = len(train_x)
             else: