import tensorflow as tf
from tensorflow.keras.applications import InceptionV3
from tensorflow.keras import layers, models
from tensorflow.keras.utils import to_categorical
from tensorflow.keras.datasets import mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data()

x_train= x_train[1000:2000]
y_train = y_train[1000:2000]

x_test= x_test[200:400]
y_test = y_test[200:400]

x_train = tf.image.resize(x_train[..., tf.newaxis], (299, 299))
x_test = tf.image.resize(x_test[..., tf.newaxis], (299, 299))
x_train = tf.image.grayscale_to_rgb(x_train) / 255.0
x_test = tf.image.grayscale_to_rgb(x_test) / 255.0

y_train_cat = to_categorical(y_train, 10)
y_test_cat = to_categorical(y_test, 10)

base_model = InceptionV3(weights='imagenet', include_top=False, input_shape=(299, 299, 3))
base_model.trainable = False

x = layers.GlobalAveragePooling2D()(base_model.output)
x = layers.Dense(128, activation='relu')(x)
output = layers.Dense(10, activation='softmax')(x)

model = models.Model(inputs=base_model.input, outputs=output)

model.compile(optimizer='adam', 
              loss='categorical_crossentropy', 
              metrics=['accuracy'])

print("Trénujem model...")
model.fit(x_train, y_train_cat, epochs=5, batch_size=32, verbose=2)

print("Vyhodnocujem model...")
loss, acc = model.evaluate(x_test, y_test_cat, verbose=2)
print(f"InceptionV3 - Presnosť: {acc:.4f}")
