bug fix

predict.py

@@ -58,6 +58,6 @@ def predict(img_name, outdir, model, scale, mask_threshold):
                            device = device)
         #out_fn = out_files[i]
         result = mask_to_image(mask)
-        result.save(out_files)
+        result.save(out_files[i])
 
         logging.info("Mask saved to {}".format(out_files[i]))
\ No newline at end of file

train.py

@@ -16,7 +16,7 @@ from torch.utils.tensorboard import SummaryWriter
 from utils.dataset import BasicDataset
 from torch.utils.data import DataLoader, random_split
 
-dir_img = 'data/imgs/'
+dir_img = 'data/train_imgs/'
 dir_mask = 'data/masks/'
 dir_checkpoint = 'checkpoints/'