Merge pull request #38 from PommeDroid/ignore_size

Fix check_shape() with resize options

config.py

@@ -44,6 +44,7 @@ class Config:
 
     # Image requirement
     desired_size = 512
+    desired_shape = 512, 512, 3
 
     # Argparser dict
     args = {}

main.py

@@ -76,6 +76,11 @@ def select_phases():
         phases = add_tail(phases, ImageToResizedCrop)
     elif conf.args['auto_rescale']:
         phases = add_tail(phases, ImageToRescale)
+    elif os.path.isfile(conf.args['input']):
+        if not conf.args['ignore_size']:
+            check_shape(conf.args['input'])
+        else:
+            conf.log.warn('Image Size Requirements Unchecked.')
     return phases
 
 

processing/__init__.py

@@ -127,11 +127,6 @@ class ImageTransform(Process):
         self.__starting_step = self._args['steps'][0] if self._args['steps'] else 0
         self.__ending_step = self._args['steps'][1] if self._args['steps'] else None
 
-        if not conf.args['ignore_size']:
-            check_shape(read_image(input_path))
-        else:
-            conf.log.warn('Image Size Requirements Unchecked.')
-
         conf.log.debug("All Phases : {}".format(self.__phases))
         conf.log.debug("To Be Executed Phases : {}".format(self.__phases[self.__starting_step:self.__ending_step]))
 

utils.py

@@ -6,6 +6,7 @@ from re import finditer
 
 import coloredlogs
 import cv2
+import imageio
 import numpy as np
 from PIL import Image
 from config import Config as conf
@@ -52,16 +53,21 @@ def write_image(image, path):
         sys.exit(1)
 
 
-def check_shape(image, shape=(512, 512, 3)):
+def check_shape(path, shape=conf.desired_shape):
     """
     Valid the shape of an image
     :param image: <RGB> Image to check
     :param shape: <(int,int,int)> Valid shape
     :return: None
     """
-    if image.shape != shape:
-        conf.log.error("Image is not 512 x 512, got shape: {}".format(image.shape))
-        conf.log.error("You should use one of the rescale options".format(image.shape))
+    if os.path.splitext(path)[1] != ".gif":
+        img_shape = read_image(path).shape
+    else:
+        img_shape = imageio.mimread(path)[0][:, :, :3].shape
+
+    if img_shape != shape:
+        conf.log.error("Image is not 512 x 512, got shape: {}".format(img_shape))
+        conf.log.error("You should use one of the rescale options".format(img_shape))
         sys.exit(1)