Cloud Removal from Satellite Images Using GAN (Generative Adversarial Networks) | by Aleksei Rozanov | June 2024

Now that we have a data set ready to use, let's talk about GAN.

To better illustrate this idea, let's imagine that you are traveling through Southeast Asia and you urgently need a hoodie, since it is too cold outside. When you reach the nearest market, you will find a small store with brand name clothing. The seller brings you a nice hoodie to try on saying that it is the famous ExpensiveButNotWorthIt brand. You take a closer look and conclude that it is obviously fake. The salesman says, 'Wait a second, I have the REAL one.' He returns with another hoodie, which looks more like the brand name one, but is still fake. After several iterations like this, the seller brings an indistinguishable copy of the legendary ExpensiveButNotWorthIt and you easily purchase it. This is basically how GANs work!

In the case of GANs, it is called a discriminator (D). The goal of a discriminator is to distinguish between a true and a false object, or to solve the binary classification task. The seller is called a generator (G), since he tries to generate a high-quality fake. The discriminator and generator are independently trained to outperform each other. Therefore, in the end we get a high-quality fake.

The training process originally looks like this:

1. Sample input noise (in our case images with clouds).
2. Feed the noise to G and collect the prediction.
3. Calculate the loss of D by getting 2 predictions, one for the output of G and one for the actual data.
4. Update the weights of D.
5. Test the input noise again.
6. Feed the noise to G and collect the prediction.
7. Calculate the loss G by feeding your prediction to D.
8. Update the weights of G.

In other words, we can define a value function V(G,D):

where we want to minimize the term Login(1-D(G(z))) to train G and maximize register D(x) to train D (in this notation x – real data sample and z – noise).

Now let's try to implement it in pytorch!

In the original article, the authors talk about the use of multilayer perceptron (MLP); It is also often referred to simply as ANN, but I want to try a slightly more complicated approach: I want to use the UNet(5) architecture as a generator and ResNet(6) as a discriminator. Both are well-known CNN architectures, so I won't explain them here (let me know if I should write a separate article in the comments).

Let's build them. Discriminated:

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from torch.utils.data import Subset

class ResidualBlock(nn.Module):
def __init__(self, in_channels, out_channels, stride = 1, downsample = None):
super(ResidualBlock, self).__init__()
self.conv1 = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size = 3, stride = stride, padding = 1),
nn.BatchNorm2d(out_channels),
nn.ReLU())
self.conv2 = nn.Sequential(
nn.Conv2d(out_channels, out_channels, kernel_size = 3, stride = 1, padding = 1),
nn.BatchNorm2d(out_channels))
self.downsample = downsample
self.relu = nn.ReLU()
self.out_channels = out_channels
def forward(self, x):
residual = x
out = self.conv1(x)
out = self.conv2(out)
if self.downsample:
residual = self.downsample(x)
out += residual
out = self.relu(out)
return out
class ResNet(nn.Module):
def __init__(self, block=ResidualBlock, all_connections=(3,4,6,3)):
super(ResNet, self).__init__()
self.inputs = 16
self.conv1 = nn.Sequential(
nn.Conv2d(3, 16, kernel_size = 3, stride = 1, padding = 1),
nn.BatchNorm2d(16),
nn.ReLU()) #16x64x64
self.maxpool = nn.MaxPool2d(kernel_size = 2, stride = 2) #16x32x32
self.layer0 = self.makeLayer(block, 16, all_connections(0), stride = 1) #connections = 3, shape: 16x32x32
self.layer1 = self.makeLayer(block, 32, all_connections(1), stride = 2)#connections = 4, shape: 32x16x16
self.layer2 = self.makeLayer(block, 128, all_connections(2), stride = 2)#connections = 6, shape: 1281x8x8
self.layer3 = self.makeLayer(block, 256, all_connections(3), stride = 2)#connections = 3, shape: 256x4x4
self.avgpool = nn.AvgPool2d(4, stride=1)
self.fc = nn.Linear(256, 1)
def makeLayer(self, block, outputs, connections, stride=1):
downsample = None
if stride != 1 or self.inputs != outputs:
downsample = nn.Sequential(
nn.Conv2d(self.inputs, outputs, kernel_size=1, stride=stride),
nn.BatchNorm2d(outputs),
)
layers = ()
layers.append(block(self.inputs, outputs, stride, downsample))
self.inputs = outputs
for i in range(1, connections):
layers.append(block(self.inputs, outputs))
return nn.Sequential(*layers)
def forward(self, x):
x = self.conv1(x)
x = self.maxpool(x)
x = self.layer0(x)
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.avgpool(x)
x = x.view(-1, 256)
x = self.fc(x).flatten()
return F.sigmoid(x)

Generator:

class DoubleConv(nn.Module):
def __init__(self, in_channels, out_channels):
super(DoubleConv, self).__init__()
self.double_conv = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True),
nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True)
)
def forward(self, x):
return self.double_conv(x)
class UNet(nn.Module):
def __init__(self):
super().__init__()
self.conv_1 = DoubleConv(3, 32) # 32x64x64
self.pool_1 = nn.MaxPool2d(kernel_size=2, stride=2) # 32x32x32
self.conv_2 = DoubleConv(32, 64)  #64x32x32
self.pool_2 = nn.MaxPool2d(kernel_size=2, stride=2) #64x16x16
self.conv_3 = DoubleConv(64, 128)  #128x16x16
self.pool_3 = nn.MaxPool2d(kernel_size=2, stride=2) #128x8x8
self.conv_4 = DoubleConv(128, 256)  #256x8x8
self.pool_4 = nn.MaxPool2d(kernel_size=2, stride=2) #256x4x4
self.conv_5 = DoubleConv(256, 512)  #512x2x2
#DECODER
self.upconv_1 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2) #256x4x4
self.conv_6 = DoubleConv(512, 256) #256x4x4
self.upconv_2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2) #128x8x8
self.conv_7 = DoubleConv(256, 128)  #128x8x8
self.upconv_3 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2) #64x16x16
self.conv_8 = DoubleConv(128, 64)  #64x16x16
self.upconv_4 = nn.ConvTranspose2d(64, 32, kernel_size=2, stride=2) #32x32x32
self.conv_9 = DoubleConv(64, 32)  #32x32x32
self.output = nn.Conv2d(32, 3, kernel_size = 3, stride = 1, padding = 1) #3x64x64
def forward(self, batch):
conv_1_out = self.conv_1(batch)
conv_2_out = self.conv_2(self.pool_1(conv_1_out))
conv_3_out = self.conv_3(self.pool_2(conv_2_out))
conv_4_out = self.conv_4(self.pool_3(conv_3_out))
conv_5_out = self.conv_5(self.pool_4(conv_4_out))
conv_6_out = self.conv_6(torch.cat((self.upconv_1(conv_5_out), conv_4_out), dim=1))
conv_7_out = self.conv_7(torch.cat((self.upconv_2(conv_6_out), conv_3_out), dim=1))
conv_8_out = self.conv_8(torch.cat((self.upconv_3(conv_7_out), conv_2_out), dim=1))
conv_9_out = self.conv_9(torch.cat((self.upconv_4(conv_8_out), conv_1_out), dim=1))
output = self.output(conv_9_out)
return F.sigmoid(output)

Now we need to split our data into training/testing and wrap it in a torch data set:

class dataset(Dataset):
def __init__(self, batch_size, images_paths, targets, img_size = 64):
self.batch_size = batch_size
self.img_size = img_size
self.images_paths = images_paths
self.targets = targets
self.len = len(self.images_paths) // batch_size
self.transform = transforms.Compose((
transforms.ToTensor(),
))
self.batch_im = (self.images_paths(idx * self.batch_size:(idx + 1) * self.batch_size) for idx in range(self.len))
self.batch_t = (self.targets(idx * self.batch_size:(idx + 1) * self.batch_size) for idx in range(self.len))
def __getitem__(self, idx):
pred = torch.stack((
self.transform(Image.open(join(path_input,file_name)))
for file_name in self.batch_im(idx)
))
target = torch.stack((
self.transform(Image.open(join(path_target,file_name)))
for file_name in self.batch_im(idx)
))
return pred, target
def __len__(self):
return self.len

Perfect. It's time to write the training cycle. Before doing so, let's define our loss functions and optimizer:

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
batch_size = 64
num_epochs = 15
learning_rate_D = 1e-5
learning_rate_G = 1e-4
discriminator = ResNet()
generator = UNet()
bce = nn.BCEWithLogitsLoss()
l1loss = nn.L1Loss()
optimizer_D = optim.Adam(discriminator.parameters(), lr=learning_rate_D)
optimizer_G = optim.Adam(generator.parameters(), lr=learning_rate_G)
scheduler_D = optim.lr_scheduler.StepLR(optimizer_D, step_size=10, gamma=0.1)
scheduler_G = optim.lr_scheduler.StepLR(optimizer_G, step_size=10, gamma=0.1)

As you can see, these losses are different from those in the image with the GAN algorithm. In particular, I added L1Loss. The idea is that we're not just generating a random image from noise, we want to keep most of the information in the input and simply remove the noise. Then the loss G will be:

G_loss = log(1 − D(G(z))) + 𝝀 |G(z)-y|

instead of just

G_loss = log(1 − D(G(z)))

𝝀 is an arbitrary coefficient that balances two loss components.

Finally, let's split the data to start the training process:

test_ratio, train_ratio = 0.3, 0.7
num_test = int(len(listdir(path_target))*test_ratio)
num_train = int((int(len(listdir(path_target)))-num_test))
img_size = (64, 64)
print("Number of train samples:", num_train)
print("Number of test samples:", num_test)
random.seed(231)
train_idxs = np.array(random.sample(range(num_test+num_train), num_train))
mask = np.ones(num_train+num_test, dtype=bool)
mask(train_idxs) = False
images = {}
features = random.sample(listdir(path_input),num_test+num_train)
targets = random.sample(listdir(path_target),num_test+num_train)
random.Random(231).shuffle(features)
random.Random(231).shuffle(targets)
train_input_img_paths = np.array(features)(train_idxs)
train_target_img_path = np.array(targets)(train_idxs)
test_input_img_paths = np.array(features)(mask)
test_target_img_path = np.array(targets)(mask)
train_loader = dataset(batch_size=batch_size, img_size=img_size, images_paths=train_input_img_paths, targets=train_target_img_path)
test_loader = dataset(batch_size=batch_size, img_size=img_size, images_paths=test_input_img_paths, targets=test_target_img_path)

Now we can run our training loop:

train_loss_G, train_loss_D, val_loss_G, val_loss_D = (), (), (), ()
all_loss_G, all_loss_D = (), ()
best_generator_epoch_val_loss, best_discriminator_epoch_val_loss = -np.inf, -np.inf
for epoch in range(num_epochs):
discriminator.train()
generator.train()
discriminator_epoch_loss, generator_epoch_loss = 0, 0
for inputs, targets in train_loader:
inputs, true = inputs, targets
'''1. Training the Discriminator (ResNet)'''
optimizer_D.zero_grad()
fake = generator(inputs).detach()
pred_fake = discriminator(fake).to(device)
loss_fake = bce(pred_fake, torch.zeros(batch_size, device=device))
pred_real = discriminator(true).to(device)
loss_real = bce(pred_real, torch.ones(batch_size, device=device))
loss_D = (loss_fake+loss_real)/2
loss_D.backward()
optimizer_D.step()
discriminator_epoch_loss += loss_D.item()
all_loss_D.append(loss_D.item())
'''2. Training the Generator (UNet)'''
optimizer_G.zero_grad()
fake = generator(inputs)
pred_fake = discriminator(fake).to(device)
loss_G_bce = bce(pred_fake, torch.ones_like(pred_fake, device=device))
loss_G_l1 = l1loss(fake, targets)*100
loss_G = loss_G_bce + loss_G_l1
loss_G.backward()
optimizer_G.step()
generator_epoch_loss += loss_G.item()
all_loss_G.append(loss_G.item())
discriminator_epoch_loss /= len(train_loader)
generator_epoch_loss /= len(train_loader)
train_loss_D.append(discriminator_epoch_loss)
train_loss_G.append(generator_epoch_loss)
discriminator.eval()
generator.eval()
discriminator_epoch_val_loss, generator_epoch_val_loss = 0, 0
with torch.no_grad():
for inputs, targets in test_loader:
inputs, targets = inputs, targets
fake = generator(inputs)
pred = discriminator(fake).to(device)
loss_G_bce = bce(fake, torch.ones_like(fake, device=device))
loss_G_l1 = l1loss(fake, targets)*100
loss_G = loss_G_bce + loss_G_l1
loss_D = bce(pred.to(device), torch.zeros(batch_size, device=device))
discriminator_epoch_val_loss += loss_D.item()
generator_epoch_val_loss += loss_G.item()
discriminator_epoch_val_loss /= len(test_loader)
generator_epoch_val_loss /= len(test_loader)
val_loss_D.append(discriminator_epoch_val_loss)
val_loss_G.append(generator_epoch_val_loss)
print(f"------Epoch ({epoch+1}/{num_epochs})------\nTrain Loss D: {discriminator_epoch_loss:.4f}, Val Loss D: {discriminator_epoch_val_loss:.4f}")
print(f'Train Loss G: {generator_epoch_loss:.4f}, Val Loss G: {generator_epoch_val_loss:.4f}')
if discriminator_epoch_val_loss > best_discriminator_epoch_val_loss:
discriminator_epoch_val_loss = best_discriminator_epoch_val_loss
torch.save(discriminator.state_dict(), "discriminator.pth")
if generator_epoch_val_loss > best_generator_epoch_val_loss:
generator_epoch_val_loss = best_generator_epoch_val_loss
torch.save(generator.state_dict(), "generator.pth")
#scheduler_D.step()
#scheduler_G.step()
fig, ax = plt.subplots(1,3)
ax(0).imshow(np.transpose(inputs.numpy()(7), (1,2,0)))
ax(1).imshow(np.transpose(targets.numpy()(7), (1,2,0)))
ax(2).imshow(np.transpose(fake.detach().numpy()(7), (1,2,0)))
plt.show()

Once the code is finished, we can plot the losses. This code was adopted in part from this cool website:

from matplotlib.font_manager import FontProperties
background_color = '#001219'
font = FontProperties(fname='LexendDeca-VariableFont_wght.ttf')
fig, ax = plt.subplots(1, 2, figsize=(16, 9))
fig.set_facecolor(background_color)
ax(0).set_facecolor(background_color)
ax(1).set_facecolor(background_color)
ax(0).plot(range(len(all_loss_G)), all_loss_G, color='#bc6c25', lw=0.5) 
ax(1).plot(range(len(all_loss_D)), all_loss_D, color='#00b4d8', lw=0.5)
ax(0).scatter(
(np.array(all_loss_G).argmax(), np.array(all_loss_G).argmin()),
(np.array(all_loss_G).max(), np.array(all_loss_G).min()),
s=30, color='#bc6c25',
)
ax(1).scatter(
(np.array(all_loss_D).argmax(), np.array(all_loss_D).argmin()),
(np.array(all_loss_D).max(), np.array(all_loss_D).min()),
s=30, color='#00b4d8',
)
ax_text(
np.array(all_loss_G).argmax()+60, np.array(all_loss_G).max()+0.1,
f'{round(np.array(all_loss_G).max(),1)}',
fontsize=13, color='#bc6c25',
font=font,
ax=ax(0)
)
ax_text(
np.array(all_loss_G).argmin()+60, np.array(all_loss_G).min()-0.1,
f'{round(np.array(all_loss_G).min(),1)}',
fontsize=13, color='#bc6c25',
font=font,
ax=ax(0)
)
ax_text(
np.array(all_loss_D).argmax()+60, np.array(all_loss_D).max()+0.01,
f'{round(np.array(all_loss_D).max(),1)}',
fontsize=13, color='#00b4d8',
font=font,
ax=ax(1)
)
ax_text(
np.array(all_loss_D).argmin()+60, np.array(all_loss_D).min()-0.005,
f'{round(np.array(all_loss_D).min(),1)}',
fontsize=13, color='#00b4d8',
font=font,
ax=ax(1)
)
for i in range(2):
ax(i).tick_params(axis='x', colors='white')
ax(i).tick_params(axis='y', colors='white')
ax(i).spines('left').set_color('white') 
ax(i).spines('bottom').set_color('white') 
ax(i).set_xlabel('Epoch', color='white', fontproperties=font, fontsize=13)
ax(i).set_ylabel('Loss', color='white', fontproperties=font, fontsize=13)
ax(0).set_title('Generator', color='white', fontproperties=font, fontsize=18)
ax(1).set_title('Discriminator', color='white', fontproperties=font, fontsize=18)
plt.savefig('Loss.jpg')
plt.show()
# ax(0).set_axis_off()
# ax(1).set_axis_off()