This repository has been archived by the owner on Dec 4, 2019. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 61
/
helpers.py
128 lines (107 loc) · 3.66 KB
/
helpers.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
# Copyright (c) Facebook, Inc. and its affiliates.
#
# 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.
## @package helpers
# Module caffe2.python.tutorials.helpers
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals
import numpy as np
import skimage.io
import skimage.transform
def crop_center(img, cropx, cropy):
y, x, c = img.shape
startx = x // 2 - (cropx // 2)
starty = y // 2 - (cropy // 2)
return img[starty:starty + cropy, startx:startx + cropx]
def rescale(img, input_height, input_width):
# print("Original image shape:" + str(img.shape) + " --> it should be in H, W, C!")
# print("Model's input shape is %dx%d") % (input_height, input_width)
aspect = img.shape[1] / float(img.shape[0])
# print("Orginal aspect ratio: " + str(aspect))
if(aspect > 1):
# landscape orientation - wide image
res = int(aspect * input_height)
imgScaled = skimage.transform.resize(
img,
(input_width, res),
preserve_range=False)
if(aspect < 1):
# portrait orientation - tall image
res = int(input_width / aspect)
imgScaled = skimage.transform.resize(
img,
(res, input_height),
preserve_range=False)
if(aspect == 1):
imgScaled = skimage.transform.resize(
img,
(input_width, input_height),
preserve_range=False)
return imgScaled
def load(img):
# load and transform image
img = skimage.img_as_float(skimage.io.imread(img)).astype(np.float32)
return img
def chw(img):
# switch to CHW
img = img.swapaxes(1, 2).swapaxes(0, 1)
return img
def bgr(img):
# switch to BGR
img = img[(2, 1, 0), :, :]
return img
def removeMean(img, mean):
# remove mean for better results
img = img * 255 - mean
return img
def batch(img):
# add batch size
img = img[np.newaxis, :, :, :].astype(np.float32)
return img
def parseResults(results):
results = np.asarray(results)
results = np.delete(results, 1)
index = 0
highest = 0
arr = np.empty((0, 2), dtype=object)
arr[:, 0] = int(10)
arr[:, 1:] = float(10)
for i, r in enumerate(results):
# imagenet index begins with 1!
i = i + 1
arr = np.append(arr, np.array([[i, r]]), axis=0)
if (r > highest):
highest = r
index = i
# top 3 results
print("Raw top 3 results:", sorted(arr, key=lambda x: x[1], reverse=True)[:3])
# now we can grab the code list
with open('inference_codes.txt', 'r') as f:
for line in f:
code, result = line.partition(":")[::2]
if (code.strip() == str(index)):
answer = "The image contains a %s with a %s percent probability." \
% (result.strip()[1:-2], highest * 100)
f.closed
return answer
def loadToNCHW(img, mean, input_size):
img = load(img)
img = rescale(img, input_size, input_size)
img = crop_center(img, input_size, input_size)
img = chw(img)
img = bgr(img)
img = removeMean(img, mean)
img = batch(img)
return img