Update dataset_tool.py

drboog · Mar 27, 2022 · acf5010 · acf5010
1 parent 5b48186
commit acf5010
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/dataset_tool.py b/dataset_tool.py
@@ -18,7 +18,7 @@
 import torch.nn.functional as F
 import torchvision.transforms as T
 import torch
-
+import cv2
 
 def custom_reshape(img, mode='bicubic', ratio=0.99):   # more to be implemented here
     full_size = img.shape[-2]
@@ -90,6 +90,10 @@ def iterate_images():
             arch_fname = arch_fname.replace('\\', '/')
             try:
                 img = np.array(PIL.Image.open(fname))
+
+                if img.shape[2] == 4:
+                    img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
+
                 try:
                     with open(fname[:-4] + '.txt', 'r') as file:
                         txt = file.read().split('\n')