-
Notifications
You must be signed in to change notification settings - Fork 1
/
prediction.py
108 lines (88 loc) · 3.67 KB
/
prediction.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
#!/usr/bin/python
# -*- coding: utf-8 -*-
from cv2 import imread
from keras.models import load_model
from PIL import Image
import numpy as np
import vlc
"""
Predict
This script loads a pre-trained CNN model and classifies whale blowholes based
on a single image
Isaac Vandor
"""
from argparse import ArgumentParser
def parse_command_line(raw_args=None):
"""
Parse command-line arguments.
Use the built-in Python argparse facility to parse the command-line
arguments as well as construct and provide meaningful help and feedback
should bad data be entered or help be requested. By doing this we'll
automatically get a traditional usage message as well as a help
response.
Args:
raw_args: The arguments as obtained from the command line.
Normally this is left at None unless one wants to
load in values for unit testing.
Returns:
The input filename.
"""
assert raw_args is None or isinstance(raw_args, list)
# Use the Python argparser to parse command line input and provide
# usage information and feedback should bad input be given or help
# requested.
default_filename = '../6AK2017_BlowholeCompilation_0929.mp4'
parser = ArgumentParser(
description="Grab still images from a video."
)
parser.add_argument(
'--input', '-i', default=default_filename,
nargs='?', const=default_filename, metavar="input_filename",
help='Optional filename for input video file. '
'Default: {0}'.format(default_filename)
)
args = parser.parse_args(raw_args)
return args.input
def load(trained_model):
""" Loads a pre-trained model. """
model = load_model(trained_model)
return model
def predict(trained_model, test_image):
""" Loads an image, resizes it to the size model was trained on,
corrects the color channels to be similar to the model's channels
and predicts the blowhole """
img = Image.open(input_filename)
img = img.resize((75,75), resample=0) # resize to 75x75 px
img = img.save('Data/OutputData/temp.jpg')
img = imread('Data/OutputData/temp.jpg')
img = img.astype(np.float32)/255.0 # convert to float32
#img = np.array(img).astype(np.float32)
# turn image into a 1-element batch :
#img = np.expand_dims(img, axis=0)
img = img[:,:,::-1] # convert from RGB to BGR
# prediction probability vector :
#result = model.predict(img)
result = trained_model.predict(np.expand_dims(img, axis=0))[0]
return result
def find_blowhole(list, dict):
""" Finds the biggest element in the list and looks for the corresponding
key in the dictionary
result: list whose biggest element we're trying to find
list: dictionary whose key corresponds to the largest element """
idx = list.argmax(axis=0) # find the index of the biggest argument
# most probable item :
#best_index = np.argmax(result, axis=1)[0]
# look for the key corresponding to the biggest argument
decoded = [key for key, value in dict.items() if value == idx]
return decoded[0]
#return best_index
if __name__ == "__main__":
input_filename = parse_command_line()
model = load(trained_model='models/model.h5')
result = predict(trained_model=model, test_image='Data/OutputData/scaledoutput.jpg')
whale_types = {"A": 0, "B":1, "C": 2, "D":3, "E": 4, "F": 5,
"G": 6, "H": 7, "I": 8, "J": 9, "K": 10, "L": 11,
"M": 12, "N": 13, "O": 14, "P": 15, "Q": 16, "S": 17,
"T": 18, "U": 19, "V": 20, "W": 21, "X": 22, "Y": 23}
alphabet = find_blowhole(list=result, dict=whale_types)
print("The blowhole is: ", alphabet)