Using DCGAN to Generate Synthetic Marine Plastic Images

The paper starts by highlighting the growing problem of marine plastic pollution and the difficulty of obtaining large, annotated image datasets for detection. Existing datasets such as JAMSTEC‑JEDI are limited, prompting the use of generative adversarial networks (GANs) to synthesize additional data.

It introduces the original GAN concept by Ian Goodfellow (2014) and explains that a GAN consists of a generator that creates fake data and a discriminator that tries to distinguish fake from real data. The DCGAN variant replaces fully‑connected layers with convolutional layers in both components, enabling unsupervised representation learning for images.

The implementation uses PyTorch and the DeepTrash dataset (also referenced as DeepPlastic). After importing required libraries, the script sets reproducible random seeds and defines hyper‑parameters such as batch size (128), image size (64), number of channels (3), latent vector size (100), and learning rate (0.0002).

from __future__ import print_function
import argparse, os, random, torch, torch.nn as nn
import torch.backends.cudnn as cudnn, torch.optim as optim
import torchvision.datasets as dset, torchvision.transforms as transforms
import torchvision.utils as vutils, numpy as np, matplotlib.pyplot as plt
manualSeed = 999
random.seed(manualSeed)
torch.manual_seed(manualSeed)

The generator and discriminator architectures are defined with transposed convolutions and batch‑normalization for the generator, and standard convolutions with leaky ReLU and sigmoid output for the discriminator.

# Generator
class Generator(nn.Module):
    def __init__(self, ngpu):
        super(Generator, self).__init__()
        self.ngpu = ngpu
        self.main = nn.Sequential(
            nn.ConvTranspose2d(nz, ngf*8, 4, 1, 0, bias=False),
            nn.BatchNorm2d(ngf*8), nn.ReLU(True),
            nn.ConvTranspose2d(ngf*8, ngf*4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf*4), nn.ReLU(True),
            nn.ConvTranspose2d(ngf*4, ngf*2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf*2), nn.ReLU(True),
            nn.ConvTranspose2d(ngf*2, nc, 4, 2, 1, bias=False),
            nn.Tanh()
        )
    def forward(self, input):
        return self.main(input)

# Discriminator
class Discriminator(nn.Module):
    def __init__(self, ngpu):
        super(Discriminator, self).__init__()
        self.ngpu = ngpu
        self.main = nn.Sequential(
            nn.Conv2d(nc, ndf, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(ndf, ndf*2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf*2), nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(ndf*2, ndf*4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf*4), nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(ndf*4, 1, 4, 1, 0, bias=False),
            nn.Sigmoid()
        )
    def forward(self, input):
        return self.main(input)

A training function is provided that performs one epoch of alternating discriminator and generator updates, uses label smoothing (real = 0.9, fake = 0.1), logs losses with Wandb, and periodically saves generated image grids.

def train(args, gen, disc, device, dataloader, optimizerG, optimizerD, criterion, epoch, iters):
    gen.train(); disc.train()
    real_label = 0.9; fake_label = 0.1
    for i, data in enumerate(dataloader, 0):
        # Update Discriminator with real data
        disc.zero_grad()
        real_cpu = data[0].to(device)
        b_size = real_cpu.size(0)
        label = torch.full((b_size,), real_label, device=device)
        output = disc(real_cpu).view(-1)
        errD_real = criterion(output, label)
        errD_real.backward()
        D_x = output.mean().item()
        # Update Discriminator with fake data
        noise = torch.randn(b_size, nz, 1, 1, device=device)
        fake = gen(noise)
        label.fill_(fake_label)
        output = disc(fake.detach()).view(-1)
        errD_fake = criterion(output, label)
        errD_fake.backward()
        D_G_z1 = output.mean().item()
        errD = errD_real + errD_fake
        optimizerD.step()
        # Update Generator
        gen.zero_grad()
        label.fill_(real_label)
        output = disc(fake).view(-1)
        errG = criterion(output, label)
        errG.backward()
        D_G_z2 = output.mean().item()
        optimizerG.step()
        if i % 50 == 0:
            print('[%d/%d][%d/%d]\tLoss_D: %.4f\tLoss_G: %.4f\tD(x): %.4f\tD(G(z)): %.4f / %.4f' %
                  (epoch, args.epochs, i, len(dataloader), errD.item(), errG.item(), D_x, D_G_z1, D_G_z2))
            wandb.log({"Gen Loss": errG.item(), "Disc Loss": errD.item()})
        if (iters % 500 == 0) or (epoch == args.epochs-1 and i == len(dataloader)-1):
            with torch.no_grad():
                fake = gen(fixed_noise).detach().cpu()
            img_list.append(wandb.Image(vutils.make_grid(fake, padding=2, normalize=True)))
            wandb.log({"Generated Images": img_list})
        iters += 1

The script initializes Wandb for experiment tracking, loads the dataset (the example uses CIFAR‑10 for illustration), creates the generator and discriminator, applies a custom weight initialization, and runs the training loop for the specified number of epochs. After training, the model checkpoint is saved.

Results are visualized by plotting generator and discriminator losses over iterations and by displaying side‑by‑side real and generated images. The loss curves show the typical adversarial dynamics, while the final generated samples resemble marine plastic debris, demonstrating that the DCGAN can augment existing marine‑plastic datasets.

In summary, the article presents a complete PyTorch pipeline—from data preparation and model definition to training, monitoring, and result visualization—for generating synthetic marine‑plastic images with a DCGAN, offering a practical way to expand limited datasets.