Browse Source

input not 512x512 error message #26

master
PommeDroid 3 years ago
parent
commit
f29e4289b7
  1. 5
      main.py
  2. 3
      transform/gan/__init__.py
  3. 3
      utils.py

5
main.py

@ -6,7 +6,7 @@ from multiprocessing import freeze_support @@ -6,7 +6,7 @@ from multiprocessing import freeze_support
import argv
from config import Config as conf
from utils import setup_log
from utils import setup_log, read_image, check_shape
from processing.gif import SimpleGIFTransform
from processing.image import SimpleImageTransform, MultipleImageTransform
@ -77,7 +77,8 @@ def select_phases(): @@ -77,7 +77,8 @@ def select_phases():
phases = add_tail(phases, ImageToResizedCrop())
elif conf.args['auto_rescale']:
phases = add_tail(phases, ImageToRescale())
else:
check_shape(read_image(conf.args['input']))
return phases

3
transform/gan/__init__.py

@ -10,6 +10,7 @@ from torchvision import transforms as transforms @@ -10,6 +10,7 @@ from torchvision import transforms as transforms
from config import Config as conf
from transform import ImageTransform
from utils import check_shape
class ImageTransformGAN(ImageTransform):
@ -34,6 +35,8 @@ class ImageTransformGAN(ImageTransform): @@ -34,6 +35,8 @@ class ImageTransformGAN(ImageTransform):
:param image: <RGB> image to be transform
:return: None
"""
check_shape(image)
if self.__gpu_ids:
conf.log.debug("GAN Processing Using GPU IDs: {}".format(self.__gpu_ids))
else:

3
utils.py

@ -52,11 +52,12 @@ def check_shape(image, shape=(512, 512, 3)): @@ -52,11 +52,12 @@ def check_shape(image, shape=(512, 512, 3)):
"""
Valid the shape of an image
:param image: <RGB> Image to check
:param shape: <(int,int,int)> Valid hape
: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))
sys.exit(1)

Loading…
Cancel
Save