train_stgan.py

"""
train_stgan.py — Training script for the Spatial Transformer GAN.

Usage
-----
.. code-block:: bash

    python train_stgan.py --group 0 --model STGAN --warpN 1 --toIt 50000

    # Resume from checkpoint
    python train_stgan.py --group 0 --model STGAN --warpN 1 --fromIt 20000 --toIt 50000

    # Multi-stage warp (stage 2 initialised from stage 1 checkpoint)
    python train_stgan.py --group 0 --model STGAN --warpN 2 --toIt 50000

Training Loop
-------------
Each iteration:
1. Sample a random batch (no-glasses BG, glasses-face real, glasses FG).
2. Apply small random perturbations to BG and FG for augmentation.
3. Run ``opt.updateGP`` steps of the geometric predictor (adversarial loss
   + warp-update norm regularisation).
4. Run ``opt.updateD`` steps of the discriminator (WGAN loss + gradient
   penalty), using a replay buffer to stabilise training.

TensorBoard summaries are written to ``summary_{group}/{model}/`` every
20/200/500 iterations for losses, gradients, and composite image grids.
Model checkpoints are saved to ``models_{group}/`` every 2000 iterations.

Requirements
------------
TensorFlow 1.x (session-based API).  For TF 2.x, run with::

    import tensorflow.compat.v1 as tf
    tf.disable_v2_behavior()
"""

import sys
import time

import numpy as np

# Add src/ to path so `stgan` package is importable
sys.path.insert(0, "src")

from stgan import utils

utils.mkdir  # noqa: B018 — ensure utils is importable before TF setup


def main() -> None:
    """Entry point: parse arguments, build graph, and run the training loop."""
    print(utils.toYellow("======================================================="))
    print(utils.toYellow("train_stgan.py — ST-GAN with homography"))
    print(utils.toYellow("======================================================="))

    import tensorflow as tf
    from stgan import data, graph, options, warp

    opt = options.set(training=True)

    utils.mkdir(f"models_{opt.group}")

    print(utils.toMagenta("building graph..."))
    tf.reset_default_graph()

    with tf.device(opt.GPUdevice):
        # Input placeholders
        imageRealData    = tf.placeholder(tf.float32, shape=[opt.batchSize, opt.dataH, opt.dataW, 3])
        imageBGfakeData  = tf.placeholder(tf.float32, shape=[opt.batchSize, opt.dataH, opt.dataW, 3])
        imageFGfake      = tf.placeholder(tf.float32, shape=[opt.batchSize, opt.H, opt.W, 4])
        PH = [imageBGfakeData, imageRealData, imageFGfake]

        # Data augmentation: random perturbation of BG and FG
        imageReal   = data.perturbBG(opt, imageRealData)
        imageBGfake = data.perturbBG(opt, imageBGfakeData)
        pPertFG = opt.pertFG * tf.random_normal([opt.batchSize, opt.warpDim])

        # Geometric Predictor + Discriminator
        geometric     = graph.geometric_multires
        discriminator = graph.discriminator

        # Multi-stage spatial transformer prediction
        imageFGwarpAll, pAll, dp = geometric(opt, imageBGfake, imageFGfake, pPertFG)
        pWarp = pAll[-1]
        dp_sqnorm = tf.reduce_sum(dp ** 2 + 1e-8, reduction_indices=[1])

        # Build composite at each warp stage for TensorBoard
        summaryImageTrain, summaryImageTest = [], []
        summaryImageTrain.append(utils.imageSummary(opt, imageReal, "TRAIN_real", opt.H, opt.W))
        summaryImageTest.append(utils.imageSummary(opt, imageReal, "TEST_real", opt.H, opt.W))

        for stage in range(opt.warpN + 1):
            imageFGwarp = imageFGwarpAll[stage]
            imageComp   = graph.composite(opt, imageBGfake, imageFGwarp)
            summaryImageTrain.append(utils.imageSummary(opt, imageComp, f"TRAIN_compST{stage}", opt.H, opt.W))
            summaryImageTest.append(utils.imageSummary(opt, imageComp, f"TEST_compST{stage}", opt.H, opt.W))

        # Interpolated image for gradient penalty
        alpha      = tf.random_uniform(shape=[opt.batchSize, 1, 1, 1])
        imageIntp  = alpha * imageReal + (1 - alpha) * imageComp

        # Discriminator outputs
        outReal  = discriminator(opt, imageReal)
        outComp  = discriminator(opt, imageComp, reuse=True)
        outIntp  = discriminator(opt, imageIntp, reuse=True)

        # Gradient penalty (WGAN-GP)
        grad_D_fake      = tf.gradients(outIntp, imageIntp)[0]
        grad_D_norm      = tf.sqrt(tf.reduce_sum(grad_D_fake ** 2 + 1e-8, reduction_indices=[1, 2, 3]))
        grad_D_norm_mean = tf.reduce_mean(grad_D_norm)

        # Losses
        loss_D        = tf.reduce_mean(outComp) - tf.reduce_mean(outReal)
        loss_GP       = -tf.reduce_mean(outComp) + opt.dplambda * tf.reduce_mean(dp_sqnorm)
        loss_D_grad   = tf.reduce_mean((grad_D_norm - 1) ** 2)
        loss_D       += opt.gradlambda * loss_D_grad

        # Optimisers
        varsGP    = [v for v in tf.global_variables() if "geometric" in v.name]
        varsGPcur = [v for v in varsGP if f"geometric/warp{opt.warpN - 1}" in v.name]
        varsGPprev= [v for v in varsGP if v not in varsGPcur]
        varsD     = [v for v in tf.global_variables() if "discrim" in v.name]

        lrGP_PH = tf.placeholder(tf.float32, shape=[])
        lrD_PH  = tf.placeholder(tf.float32, shape=[])
        with tf.name_scope("adam"):
            optimGP = tf.train.AdamOptimizer(learning_rate=lrGP_PH).minimize(loss_GP, var_list=varsGPcur)
            optimD  = tf.train.AdamOptimizer(learning_rate=lrD_PH).minimize(loss_D,  var_list=varsD)

        # TensorBoard summaries
        summaryLossTrain = tf.summary.merge([
            tf.summary.scalar("TRAIN_loss_D",  loss_D),
            tf.summary.scalar("TRAIN_loss_GP", loss_GP),
        ])
        summaryGradTrain  = tf.summary.scalar("TRAIN_grad_D", grad_D_norm_mean)
        summaryImageTrain = tf.summary.merge(summaryImageTrain)
        summaryImageTest  = tf.summary.merge(summaryImageTest)

    # Load dataset
    print(utils.toMagenta("loading training data..."))
    trainData = data.load(opt)
    print(utils.toMagenta("loading test data..."))
    testData = data.load(opt, test=True)

    saver_GP = tf.train.Saver(var_list=varsGP, max_to_keep=20)
    saver_D  = tf.train.Saver(var_list=varsD,  max_to_keep=20)
    if opt.warpN > 1:
        saver_GPprev = tf.train.Saver(var_list=varsGPprev)
        varsGPdict = {}
        for v in varsGPcur:
            scopes = v.op.name.split("/")
            scopes[1] = f"warp{opt.warpN - 2}" if opt.loadGP == "prev" else "warp0"
            varsGPdict["/".join(scopes)] = v
        saver_GPcur = tf.train.Saver(varsGPdict)
    else:
        saver_GPcur = tf.train.Saver(var_list=varsGPcur)

    summaryWriter = tf.summary.FileWriter(f"summary_{opt.group}/{opt.model}")

    print(utils.toYellow("======= TRAINING START ======="))
    t_start = time.time()

    tfConfig = tf.ConfigProto(allow_soft_placement=True)
    tfConfig.gpu_options.allow_growth = True

    with tf.Session(config=tfConfig) as sess:
        sess.run(tf.global_variables_initializer())
        summaryWriter.add_graph(sess.graph)

        if opt.fromIt != 0:
            utils.restoreModelFromIt(opt, sess, saver_GP, "GP", opt.fromIt)
            utils.restoreModelFromIt(opt, sess, saver_D,  "D",  opt.fromIt)
            print(utils.toMagenta(f"resuming from iteration {opt.fromIt}..."))
        elif opt.warpN > 1:
            utils.restoreModelPrevStage(opt, sess, saver_GPprev, "GP")
            print(utils.toMagenta("loading GP from previous warp stage..."))
            if opt.loadGP == "prev":
                utils.restoreModelPrevStage(opt, sess, saver_GPcur, "GP")
            elif opt.loadGP is not None:
                utils.restoreModel(opt, sess, saver_GPcur, opt.loadGP, "GP")
            utils.restoreModelPrevStage(opt, sess, saver_D, "D")
        else:
            if opt.loadGP:
                utils.restoreModel(opt, sess, saver_GPcur, opt.loadGP, "GP")
            if opt.loadD:
                utils.restoreModel(opt, sess, saver_D, opt.loadD, "D")

        print(utils.toMagenta("start training..."))

        for it in range(opt.fromIt, opt.toIt):
            lr_GP = opt.lrGP * opt.lrGPdecay ** (it // opt.lrGPstep)
            lr_D  = opt.lrD  * opt.lrDdecay  ** (it // opt.lrDstep)

            batch = data.makeBatch(opt, trainData, PH)
            batch[lrGP_PH] = lr_GP
            batch[lrD_PH]  = lr_D

            # Update Geometric Predictor
            for _ in range(opt.updateGP):
                _, lg, ic = sess.run([optimGP, loss_GP, imageComp], feed_dict=batch)

            # Update Discriminator (with replay buffer)
            batch[imageComp] = data.updateHistory(opt, ic)
            for _ in range(opt.updateD):
                _, ld, gdn = sess.run([optimD, loss_D, grad_D_norm_mean], feed_dict=batch)

            global_it = (opt.warpN - 1) * opt.toIt + it + 1

            if (it + 1) % 10 == 0:
                print(
                    f"it.{utils.toCyan(str(global_it))}/{opt.warpN * opt.toIt}  "
                    f"lr={utils.toYellow(f'{lr_GP:.0e}')}(GP),"
                    f"{utils.toYellow(f'{lr_D:.0e}')}(D)  "
                    f"loss={utils.toRed(f'{lg:.4f}')}(GP),"
                    f"{utils.toRed(f'{ld:.4f}')}(D)  "
                    f"norm={utils.toBlue(f'{gdn:.4f}')}  "
                    f"time={utils.toGreen(f'{time.time() - t_start:.2f}')}"
                )
            if (it + 1) % 20 == 0:
                sl, sg = sess.run([summaryLossTrain, summaryGradTrain], feed_dict=batch)
                summaryWriter.add_summary(sl, global_it)
                summaryWriter.add_summary(sg, global_it)
            if (it + 1) % 200 == 0:
                si = sess.run(summaryImageTrain, feed_dict=batch)
                summaryWriter.add_summary(si, global_it)
            if (it + 1) % 500 == 0:
                batch_test = data.makeBatch(opt, testData, PH)
                si = sess.run(summaryImageTest, feed_dict=batch_test)
                summaryWriter.add_summary(si, global_it)
            if (it + 1) % 2000 == 0:
                utils.saveModel(opt, sess, saver_GP, "GP", it + 1)
                utils.saveModel(opt, sess, saver_D,  "D",  it + 1)
                print(utils.toGreen(f"model saved: {opt.group}/{opt.model}, it.{it + 1}"))

    print(utils.toYellow("======= TRAINING DONE ======="))


if __name__ == "__main__":
    main()