55 lines
1.6 KiB
Python
55 lines
1.6 KiB
Python
|
#!/usr/bin/python3
|
||
|
|
|
||
|
|
img_size = (150, 150)
|
||
|
|
|
||
|
|
from tensorflow.keras import models as m
|
||
|
|
from tensorflow.keras import layers as l
|
||
|
|
from tensorflow.keras import optimizers as o
|
||
|
|
from PIL import Image
|
||
|
|
from sys import argv
|
||
|
|
from os import listdir as ls
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
model = m.Sequential([
|
||
|
|
l.Input(shape = (*img_size, 3)),
|
||
|
|
|
||
|
|
l.Conv2D(96, (11, 11), strides = 4, activation = "relu"),
|
||
|
|
l.BatchNormalization(),
|
||
|
|
l.MaxPooling2D((3, 3), strides = 2),
|
||
|
|
|
||
|
|
l.Conv2D(192, (5, 5), activation = "relu", padding = "same"),
|
||
|
|
l.BatchNormalization(),
|
||
|
|
l.MaxPooling2D((3, 3), strides = 2),
|
||
|
|
|
||
|
|
l.Conv2D(256, (3, 3), activation = "relu", padding = "same"),
|
||
|
|
l.Conv2D(256, (3, 3), activation = "relu", padding = "same"),
|
||
|
|
l.Conv2D(160, (3, 3), activation = "relu", padding = "same"),
|
||
|
|
l.BatchNormalization(),
|
||
|
|
l.MaxPooling2D((3, 3), strides = 2),
|
||
|
|
|
||
|
|
l.Flatten(),
|
||
|
|
l.Dense(1024, activation = "relu"),
|
||
|
|
l.Dropout(0.5),
|
||
|
|
l.Dense(1024, activation = "relu"),
|
||
|
|
l.Dropout(0.5),
|
||
|
|
l.Dense(10, activation = "softmax"),
|
||
|
|
])
|
||
|
|
|
||
|
|
model.compile(optimizer = o.Adam(learning_rate = 0.0001),
|
||
|
|
loss = "categorical_crossentropy",
|
||
|
|
metrics = ["accuracy"])
|
||
|
|
|
||
|
|
model.load_weights("w2.weights.h5")
|
||
|
|
|
||
|
|
if len(argv) >= 2:
|
||
|
|
for i in argv[1:]:
|
||
|
|
with Image.open(i) as im:
|
||
|
|
im = im.resize((150, 150), Image.Resampling.LANCZOS)
|
||
|
|
im = np.divide(np.array(im),
|
||
|
|
np.array(255.))
|
||
|
|
|
||
|
|
res = model.predict(np.array([im]))
|
||
|
|
print(res)
|
||
|
|
print(np.argmax(res))
|
||
|
|
print(sorted(ls("../ds/raw-img"))[np.argmax(res)])
|