-
Notifications
You must be signed in to change notification settings - Fork 0
/
Attention.py
96 lines (75 loc) · 2.5 KB
/
Attention.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
import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.models import Sequential
# Habana specific libraries
from habana_frameworks.tensorflow import load_habana_module
load_habana_module()
import pathlib
dataset_url = '...tgz'
data_dir = tf.keras.utils.get_file('attention_data', origin=dataset_url, untar=True)
data_dir = pathlib.Path(data_dir)
image_count = len(list(data_dir.glob('*/*.png')))
# print(image_count)
# Create dataset
batch_size = 32
img_height = 180
img_width = 180
train_ds = tf.keras.preprocessing.image_dataset_from_directory(
data_dir,
validation_split=0.2,
subset="training",
seed=123,
image_size=(img_height, img_width),
batch_size=batch_size)
val_ds = tf.keras.preprocessing.image_dataset_from_directory(
data_dir,
validation_split=0.2,
subset="validation",
seed=123,
image_size=(img_height, img_width),
batch_size=batch_size)
class_names = train_ds.class_names
# print(class_names)
# Create model
normalization_layer = layers.experimental.preprocessing.Rescaling(1. / 255)
num_classes = 3
model = Sequential([
layers.experimental.preprocessing.Rescaling(1. / 255, input_shape=(img_height, img_width, 3)),
layers.Conv2D(16, 3, padding='same', activation='relu'),
layers.MaxPooling2D(),
layers.Conv2D(32, 3, padding='same', activation='relu'),
layers.MaxPooling2D(),
layers.Conv2D(64, 3, padding='same', activation='relu'),
layers.MaxPooling2D(),
layers.Flatten(),
layers.Dense(128, activation='relu'),
layers.Dense(num_classes)
])
model.compile(optimizer='adam',
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=['accuracy'])
model.summary()
epochs = 10
def train():
model.fit(
train_ds,
validation_data=val_ds,
epochs=epochs
)
model.save_weights('/media/classifier.h5')
train() # comment out if only running inference
def eval(image):
test_path = image
img = keras.preprocessing.image.load_img(
test_path, target_size=(img_height, img_width)
)
img_array = keras.preprocessing.image.img_to_array(img)
img_array = tf.expand_dims(img_array, 0)
model.load_weights('/Users/joshuabelofsky/Desktop/AWS/classifier.h5')
predictions = model.predict(img_array)
score = tf.nn.softmax(predictions[0])
print("Classification: " + (class_names[np.argmax(score)]))
attention = (class_names[np.argmax(score)])
return attention