Update train.py

tohinz · Mar 1, 2021 · 5e95373 · 5e95373
1 parent a14b006
commit 5e95373
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/train.py b/train.py
@@ -30,7 +30,7 @@
     pass
 
 # get resolution of training images
-_training_imgs = glob.glob(os.path.join(opt.dataroot, "*.jpg")) + glob.glob(os.path.join(opt.dataroot, "*.png"))
+_training_imgs = glob.glob(os.path.join(opt.dataroot, "*.jpg")) + glob.glob(os.path.join(opt.dataroot, "*.png")) + glob.glob(os.path.join(opt.dataroot, "*.jpeg"))
 _img = img.imread(_training_imgs[0])
 opt.image_size_x = _img.shape[0]
 opt.image_size_y = _img.shape[1]