How to Restore Original Mask from Overlayed Image Using CNN?

[StormLord07]

I have a 120x120 mask generated from a seed, which is then tiled and overlayed onto a 1920x1080 screenshot. My goal is to use a CNN to reverse this operation and restore the original mask from the overlayed image. I've tried using a simple CNN architecture but it's not working as expected.

# My CNN architecture
model = models.Sequential()
model.add(layers.Input(shape=(None, None,
model.add(layers.Conv2D(32, (3, 3), activation='relu'))
model.add(layers.MaxPooling2D((2, 2)))
model.add(layers.Conv2D(64, (3, 3), activation='relu'))
model.add(layers.MaxPooling2D((2, 2)))
model.add(layers.Conv2D(128, (3, 3), activation='relu'))
model.add(layers.GlobalAveragePooling2D())
model.add(layers.Dense(512, activation='relu'))
model.add(layers.Dense(120*120*3, activation='linear'))
model.add(layers.Reshape((120, 120, 3)))

model.compile(optimizer='adam', loss='mse')

I have also prepared the dataset, but the CNN isn't producing results close to the original mask. It's not a case of overfitting because I've tried reducing epochs and altering the dataset. After Alterting there was no even remotely close result.

Questions:

What could be wrong with my current approach?

Is the U-Net architecture more suited for this type of problem?

I'm training my model on an NVIDIA RTX 4090. Considering the computational power I have, how can I appropriately increase the number of layers or the sizes of the convolutional layers in my existing architecture to potentially improve performance?

Any better approach to the problem?

Any literature or resources to better understand the problem?

P.S. Idea was taken from this article , second neural network (first D decoder)