test.py

from api import remove_watermark

remove_watermark(
    image_path = I2MAGE_NAME,
    mask_path = MASK_NAME,
    max_dim = MAX_DIM,
    show_step = SHOW_STEPS,
    reg_noise = REG_NOISE,
    input_depth = INPUT_DEPTH,
    lr = LR,
    training_steps = TRAINING_STEPS,
    tqdm_length = 900
)