Added check in RGB branch for Alpha-channel. Remove Alpha-Channel if …

…present
nekitmm · Mar 28, 2021 · b82d697 · b82d697
1 parent d9aa3ff
commit b82d697
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