import numpy as np
import tensorflow as tf
from tensorflow.keras.datasets import mnist
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.layers import GlobalAveragePooling2D, Dense
from tensorflow.keras.models import Model
from tensorflow.keras.utils import to_categorical
from tensorflow.keras.applications.mobilenet_v2 import preprocess_input

# 1. Načítanie MNIST dát
(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]

# 2. Úprava obrázkov – z 28x28 (1 kanál) na 224x224 (3 kanály)
x_train_rgb = np.repeat(x_train[..., np.newaxis], 3, -1)  # z 28x28 → 28x28x3
x_test_rgb = np.repeat(x_test[..., np.newaxis], 3, -1)

x_train_resized = tf.image.resize(x_train_rgb, [224, 224]).numpy()
x_test_resized = tf.image.resize(x_test_rgb, [224, 224]).numpy()

# 3. Normalizácia pre MobileNetV2
x_train_resized = preprocess_input(x_train_resized)
x_test_resized = preprocess_input(x_test_resized)

# 4. Konverzia výstupov na one-hot vektory
y_train_cat = to_categorical(y_train, 10)
y_test_cat = to_categorical(y_test, 10)

# 5. Vytvorenie modelu s MobileNetV2 ako základom
base_model = MobileNetV2(weights='imagenet', include_top=False, input_shape=(224, 224, 3))
base_model.trainable = False  # zmrazíme váhy predtrénovanej časti

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

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

# 6. Kompilácia modelu
model.compile(optimizer='adam',
              loss='categorical_crossentropy',
              metrics=['accuracy'])

# 7. Tréning modelu
print("Trénujem model...")
model.fit(x_train_resized, y_train_cat, epochs=3, batch_size=32)

# 8. Vyhodnotenie modelu
print("Vyhodnocujem model...")
loss, acc = model.evaluate(x_test_resized, y_test_cat, verbose=2)
print(f"Presnosť modelu: {acc:.4f}")


