Merge pull request #20 from Epoxid2/master

Added check in RGB branch for Alpha-channel.
nekitmm · Mar 31, 2021 · adb22aa · adb22aa
2 parents d9aa3ff + b82d697
commit adb22aa
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/starnet_v1_TF2.py b/starnet_v1_TF2.py
@@ -323,6 +323,10 @@ def transform(self, in_name, out_name):
         if self.mode == 'RGB' and len(image.shape) == 2:
             raise ValueError('You loaded RGB model, but the image is Greyscale!')
 
+        if self.mode == 'RGB' and image.shape[2] == 4:
+            print("Input image has 4 channels. Removing Alpha-Channel")
+            image=image[:,:,[0,1,2]]
+
         offset = int((self.window_size - self.stride) / 2)
 
         h, w, _ = image.shape