import tensorflow as tf # v2.5.0
from tensorflow import keras

tf.random.set_seed(1337)

def load_image(f):
  return tf.io.decode_image(tf.io.read_file(f), channels=1, dtype=tf.float32)

def save_image(f, d):
  tf.io.write_file(f, tf.io.encode_png(tf.cast(d*255, tf.uint8)))

(x, y), _test = keras.datasets.cifar10.load_data()
# color (0-255) -> grayscale (0.0-1.0)
x = tf.cast(x, dtype=tf.float32)/255.
x = tf.math.reduce_mean(x, axis=3, keepdims=True)

# Save sample images
for i in range(4):
  save_image(f"no_flag_{i}.png", x[i])
x = x[4:]
y = y[4:]

# Add flag images to the dataset
flag = load_image("flag.png") # flag.png is QR code
x = tf.concat([x, [flag]*5000], axis=0)
y = tf.concat([y, [tf.constant([10])]*5000], axis=0) # label 10 is "flag"

# Deep Learning!!
model = keras.Sequential([
  keras.Input(shape=(32, 32, 1)),
  keras.layers.Flatten(),
  keras.layers.Dense(units=512),
  keras.layers.BatchNormalization(),
  keras.layers.Dropout(rate=0.5),
  keras.layers.Dense(units=256),
  keras.layers.BatchNormalization(),
  keras.layers.Dropout(rate=0.5),
  keras.layers.Dense(units=128),
  keras.layers.BatchNormalization(),
  keras.layers.Dropout(rate=0.5),
  keras.layers.Dense(units=64),
  keras.layers.BatchNormalization(),
  keras.layers.Dropout(rate=0.5),
  keras.layers.Dense(units=32),
  keras.layers.BatchNormalization(),
  keras.layers.Dropout(rate=0.5),
  keras.layers.Dense(units=16),
  keras.layers.BatchNormalization(),
  keras.layers.Dropout(rate=0.5),
  keras.layers.Dense(units=11),
  keras.layers.BatchNormalization(),
  keras.layers.Softmax(),
])
model.summary()
model.compile(
  optimizer = "adam",
  loss = "sparse_categorical_crossentropy",
  metrics = ["accuracy"],
)
model.fit(
  x = x,
  y = y,
  batch_size = 1024,
  epochs = 16,
  shuffle = True,
)
model.save("model")