Revise paste CNN Keras

#1. IMPORT FILE
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import Sequential
from tensorflow.keras.layers import Dense, Flatten, Dropout
from tensorflow.keras.preprocessing import image
import numpy as np

#2. IMPORT DATASET (MNIST)
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

#3. VARIABLE DECLARATION / PRE-PROCESSING
# Normalize & reshape dataset
x_train = x_train.reshape(-1, 28, 28, 1).astype("float32") / 255.0
x_test  = x_test.reshape(-1, 28, 28, 1).astype("float32") / 255.0

# Create TensorFlow datasets
train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(32)
val_ds   = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(32)

#4. MODEL CREATION (your ANN model approach)
model = Sequential([
    Flatten(input_shape=(28, 28, 1)),
    Dense(512, activation="relu"),
    Dropout(0.2),
    Dense(10, activation="softmax")  # output for 0–9 MNIST
])

# OPTIONAL
# model.summary()

#5. MODEL COMPILE AND TRAINING
model.compile(optimizer="adam",
              loss="sparse_categorical_crossentropy",
              metrics=["accuracy"])

model.fit(train_ds, epochs=10, validation_data=val_ds)

#6. IMAGE ARRAY
img = image.load_img("E:/Maroon/College/S9/DL/Fonts/test1.png",
                     target_size=(28, 28), color_mode="grayscale")
img_array = np.expand_dims(image.img_to_array(img) / 255.0, axis=0)

#7. MODEL PREDICTION
pred = model.predict(img_array)
pred_class = np.argmax(pred, axis=1)[0]

#8. VALUE GENERATION
class_labels = [str(i) for i in range(10)]  # labels 0-9

#9. PRINT
print("Predicted label:", class_labels[pred_class])