Browse Source

Support input resizing in CLI

master
Ringo 3 years ago committed by PommeDroid
parent
commit
d808406ec2
  1. 26
      main.py
  2. 39
      src/cli/utils.py

26
main.py

@ -8,6 +8,7 @@ import os @@ -8,6 +8,7 @@ import os
import imageio
import sentry_sdk
import rook
import utils
from run import process, process_gif
from multiprocessing import freeze_support
@ -47,6 +48,24 @@ parser.add_argument( @@ -47,6 +48,24 @@ parser.add_argument(
parser.add_argument(
"--gif", action="store_true", default=False, help="Run the processing on a gif"
)
parser.add_argument(
"--auto-resize",
action="store_true",
default=False,
help="Scale and pad image to 512x512. Maintains aspect ratio. Doesn't support gifs for now.",
)
parser.add_argument(
"--auto-resize-crop",
action="store_true",
default=False,
help="Scale and crop image to 512x512. Maintains aspect ratio. Doesn't support gifs for now.",
)
parser.add_argument(
"--auto-rescale",
action="store_true",
default=False,
help="Scale image to 512x512. Doesn't support gifs for now.",
)
args = parser.parse_args()
"""
@ -72,6 +91,13 @@ def main(): @@ -72,6 +91,13 @@ def main():
# Read input image
image = cv2.imread(args.input)
if args.auto_resize:
image = utils.resize_input(image)
elif args.auto_resize_crop:
image = utils.resize_crop_input(image)
elif args.auto_rescale:
image = utils.rescale_input(image)
# Process
result = process(image, gpu_ids, args.enablepubes)

39
src/cli/utils.py

@ -0,0 +1,39 @@ @@ -0,0 +1,39 @@
import cv2
import math
desired_size = 512
def resize_input(img):
old_size = img.shape[:2]
ratio = float(desired_size)/max(old_size)
new_size = tuple([int(x*ratio) for x in old_size])
img = cv2.resize(img, (new_size[1], new_size[0]))
delta_w = desired_size - new_size[1]
delta_h = desired_size - new_size[0]
top, bottom = delta_h//2, delta_h-(delta_h//2)
left, right = delta_w//2, delta_w-(delta_w//2)
return cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=[255, 255, 255])
def resize_crop_input(img):
old_size = img.shape[:2]
ratio = float(desired_size)/min(old_size)
new_size = tuple([int(x*ratio) for x in old_size])
img = cv2.resize(img, (new_size[1], new_size[0]))
delta_w = new_size[1] - desired_size
delta_h = new_size[0] - desired_size
top = delta_h//2
left = delta_w//2
return img[top:desired_size+top, left:desired_size+left]
def rescale_input(img):
return cv2.resize(img, (desired_size, desired_size))
Loading…
Cancel
Save