59 lines
2 KiB
Python
59 lines
2 KiB
Python
# Copyright (c) 2018-2019, NVIDIA CORPORATION. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import numpy as np
|
|
import skimage
|
|
|
|
def load_image(image_path):
|
|
"""Code from Loading_Pretrained_Models.ipynb - a Caffe2 tutorial"""
|
|
mean, std = 128, 128
|
|
img = skimage.img_as_float(skimage.io.imread(image_path))
|
|
if len(img.shape) == 2:
|
|
img = np.array([img, img, img]).swapaxes(0,2)
|
|
return img
|
|
|
|
def rescale(img, input_height, input_width):
|
|
"""Code from Loading_Pretrained_Models.ipynb - a Caffe2 tutorial"""
|
|
aspect = img.shape[1]/float(img.shape[0])
|
|
if(aspect>1):
|
|
# landscape orientation - wide image
|
|
res = int(aspect * input_height)
|
|
imgScaled = skimage.transform.resize(img, (input_width, res))
|
|
if(aspect<1):
|
|
# portrait orientation - tall image
|
|
res = int(input_width/aspect)
|
|
imgScaled = skimage.transform.resize(img, (res, input_height))
|
|
if(aspect == 1):
|
|
imgScaled = skimage.transform.resize(img, (input_width, input_height))
|
|
return imgScaled
|
|
|
|
def crop_center(img,cropx,cropy):
|
|
"""Code from Loading_Pretrained_Models.ipynb - a Caffe2 tutorial"""
|
|
y,x,c = img.shape
|
|
startx = x//2-(cropx//2)
|
|
starty = y//2-(cropy//2)
|
|
return img[starty:starty+cropy,startx:startx+cropx]
|
|
|
|
def normalize(img, mean=128, std=128):
|
|
img = (img * 256 - mean) / std
|
|
return img
|
|
|
|
def prepare_input(img_uri):
|
|
img = load_image(img_uri)
|
|
img = rescale(img, 300, 300)
|
|
img = crop_center(img, 300, 300)
|
|
img = normalize(img)
|
|
|
|
return img
|