#!/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)])