def test(test_data_loader, netG):
psnrs = []
ssims= []
for iteration, batch in enumerate(test_data_loader, 1):
netG.eval()
steps = iteration
data_clean, data_mask, data_shadow = \
Variable(batch[0]), \
Variable(batch[1]), \
Variable(batch[2], requires_grad=False)
if opt.cuda:
data_clean = data_clean.cuda()
data_mask = data_mask.cuda()
data_shadow = data_shadow.cuda()
else:
data_clean = data_clean.cuda()
data_mask = data_mask.cuda()
data_shadow = data_shadow.cuda()
data_predict = netG(data_clean, data_mask)
with torch.no_grad():
psnr = (10 * torch.log10(1.0 / loss_mse(rgb2yuv(data_predict)[:,0].unsqueeze(1), rgb2yuv(data_shadow)[:,0].unsqueeze(1)).detach())).mean()
psnrs.append(psnr)
mse = nn.MSELoss()(rgb2yuv(data_predict)[:,0,:,:], rgb2yuv(data_shadow)[:,0,:,:])
ssim = 1 - instance_ssim(data_predict, data_shadow).mean().item()
ssims.append(ssim)
print("PSNR {} SSIM {}".format(sum(psnrs)/len(psnrs), sum(ssims)/len(ssims)))
def process(model, img):
if img is not None:
# print(img)
# print(img.shape)
img = misc.imresize(img[:, :, :], (66, 200, 3))
# print(img.shape)
img = utils.rgb2yuv(img)
# print(img.shape)
img = np.array([img])
# print(img.shape)
# print('\n\n')
# steering_angle = model.predict(img[None, :, :, :])[0][0]
steering_angle = float(model.predict(img, batch_size=1))
print(steering_angle)
pub_steering(steering_angle)
def telemetry(sid, data):
if data:
# The current steering angle of the car
steering_angle = data["steering_angle"]
# The current throttle of the car
throttle = data["throttle"]
# The current speed of the car
speed = data["speed"]
# The current image from the center camera of the car
imgString = data["image"]
image = Image.open(BytesIO(base64.b64decode(imgString)))
image_array = np.asarray(image)
## modified lines ##
image_array = utils.crop_img(image_array)
image_array = utils.blur_img(image_array)
image_array = utils.resize_img(image_array)
image_array = utils.rgb2yuv(image_array)
# This model currently assumes that the features of the model are just the images. Feel free to change this.
steering_angle = float(
model.predict(image_array[None, :, :, :], batch_size=1))
# Speed control using PI Controller
throttle = controller.update(float(speed))
print(steering_angle, throttle)
send_control(steering_angle, throttle)
# save frame
if args.image_folder != '':
timestamp = datetime.utcnow().strftime('%Y_%m_%d_%H_%M_%S_%f')[:-3]
image_filename = os.path.join(args.image_folder, timestamp)
image.save('{}.jpg'.format(image_filename))
else:
# NOTE: DON'T EDIT THIS.
sio.emit('manual', data={}, skip_sid=True)
if __name__ == "__main__":
filenames = glob.glob("demo/*")
with open("model/vgg16-20160129.tfmodel", mode='rb') as f:
fileContent = f.read()
graph_def = tf.GraphDef()
graph_def.ParseFromString(fileContent)
batch_size = 4
num_epochs = 1e+9
colorimage = input_pipeline(filenames, batch_size, num_epochs=num_epochs)
grayscale = tf.image.rgb_to_grayscale(colorimage)
grayscale_rgb = tf.image.grayscale_to_rgb(grayscale)
grayscale_yuv = rgb2yuv(grayscale_rgb)
grayscale = tf.concat([grayscale, grayscale, grayscale],3)
tf.import_graph_def(graph_def, input_map={"images": grayscale})
graph = tf.get_default_graph()
# Saver.
with tf.Session() as sess:
saver = tf.train.import_meta_graph('checkpoints/color_net_model200500.ckpt.meta')
saver.restore(sess, "checkpoints/color_net_model200500.ckpt")
pred_rgb_ = tf.placeholder(dtype=tf.float32)
pred_rgb_ = sess.run(["pred_rgb:0"], feed_dict={"phase_train:0": False, "uv:0": 3})
'wc2': tf.Variable(tf.truncated_normal([3, 3, 256, 128], stddev=0.01)),
# 3x3 conv, 256 inputs, 64 outputs
'wc3': tf.Variable(tf.truncated_normal([3, 3, 128, 64], stddev=0.01)),
# 3x3 conv, 128 inputs, 3 outputs
'wc4': tf.Variable(tf.truncated_normal([3, 3, 64, 3], stddev=0.01)),
# 3x3 conv, 6 inputs, 3 outputs
'wc5': tf.Variable(tf.truncated_normal([3, 3, 3, 3], stddev=0.01)),
# 3x3 conv, 3 inputs, 2 outputs
'wc6': tf.Variable(tf.truncated_normal([3, 3, 3, 2], stddev=0.01)),
}
# pre process input image
reader = tf.WholeFileReader()
colorimage = input_pipeline(filenames, batch_size, num_epochs=num_epochs)
colorimage_yuv = rgb2yuv(colorimage)
grayscale = tf.image.rgb_to_grayscale(colorimage)
grayscale_rgb = tf.image.grayscale_to_rgb(grayscale)
grayscale_yuv = rgb2yuv(grayscale_rgb)
grayscale = tf.concat(3, [grayscale, grayscale, grayscale])
# build tensorflow graph
tf.import_graph_def(graph_def, input_map={"images": grayscale})
graph = tf.get_default_graph()
# get weights from vgg pre trained model
with tf.variable_scope('vgg'):
conv1_2 = graph.get_tensor_by_name("import/conv1_2/Relu:0")
conv2_2 = graph.get_tensor_by_name("import/conv2_2/Relu:0")
def train_color_net(graph, phase_train, uv, grayscale):
pred_rgb = tf.placeholder(tf.float32, name="pred_rgb")
pred = color_net(graph, phase_train, grayscale)
pred_yuv = tf.concat([tf.split(grayscale_yuv, 3, 3)[0], pred],3)
pred_rgb = yuv2rgb(pred_yuv)
colorimage_yuv = rgb2yuv(colorimage)
loss = tf.square(tf.subtract(pred, tf.concat([tf.split(colorimage_yuv, 3, 3)[1], tf.split(colorimage_yuv, 3, 3)[2]],3)))
if uv == 1:
loss = tf.split(loss, 2,3)[0]
elif uv == 2:
loss = tf.split(loss, 2, 3)[1]
else:
loss = (tf.split(loss,2, 3)[0] + tf.split(loss,2, 3)[1]) / 2
global_step = tf.Variable(0, name='global_step', trainable=False)
if phase_train is not None:
optimizer = tf.train.GradientDescentOptimizer(0.0001)
opt = optimizer.minimize(loss, global_step=global_step, gate_gradients=optimizer.GATE_NONE)
# Saver.
saver = tf.train.Saver()
sess = tf.Session()
# Initialize the variables.
sess.run(tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()))
# Start input enqueue threads.
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess=sess, coord=coord)
try:
while not coord.should_stop():
# Run training steps
training_opt = sess.run(opt, feed_dict={phase_train: True, uv: 1})
training_opt = sess.run(opt, feed_dict={phase_train: True, uv: 2})
step = sess.run(global_step)
if step % 1 == 0:
pred_, pred_rgb_, colorimage_, grayscale_rgb_, cost = sess.run(
[pred, pred_rgb, colorimage, grayscale_rgb, loss], feed_dict={phase_train: False, uv: 3})
print({"step": step, "cost": np.mean(cost)})
if step % 10 == 0:
summary_image = concat_images(grayscale_rgb_[0], pred_rgb_[0])
summary_image = concat_images(summary_image, colorimage_[0])
plt.imsave("summary/" + str(step) + "_0.jpg", summary_image)
if step % 100000 == 500:
if not os.path.exists(checkpoints_dir):
os.makedirs(checkpoints_dir)
save_path = saver.save(sess, checkpoints_dir + "/color_net_model"+str(step)+".ckpt")
print("Model saved in file: %s" % save_path)
except tf.errors.OutOfRangeError:
print('Done training -- epoch limit reached')
finally:
# When done, ask the threads to stop.
coord.request_stop()
# Wait for threads to finish.
coord.join(threads)
sess.close()
def show_imgs(images):
for index, img in enumerate(images):
show_img(img)
# Preprocess for model - START:
image = utils.load_image('data', img_path)
image_flip, steering_angle = utils.random_flip(image, steering_angle)
image_translate, steering_angle = utils.random_translate(
image_flip, steering_angle, range_x, range_y)
# image_shadow = utils.random_shadow(image_translate)
# image_brightness = utils.random_brightness(image_shadow)
# image_crop = utils.crop(image_brightness)
image_crop = utils.crop(image_translate)
image_resize = utils.resize(image_crop)
image_rgb2yuv = utils.rgb2yuv(image_resize)
# Preprocess for model - END
images = [
image,
image_flip,
image_translate,
# image_shadow,
# image_brightness,
image_crop,
image_resize,
image_rgb2yuv,
]
show_imgs(images)
def main():
global opt, name, logger, netG, netD, vgg, curriculum_ssim, loss_mse, rgb2yuv, instance_ssim, loss_bce
opt = parser.parse_args()
name = "ShadowRemoval"
print(opt)
# Tag_ResidualBlocks_BatchSize
cuda = opt.cuda
if cuda and not torch.cuda.is_available():
raise Exception("No GPU found, please run without --cuda")
seed = 1334
torch.manual_seed(seed)
if 'WORLD_SIZE' in os.environ:
opt.distributed = int(os.environ['WORLD_SIZE']) > 1
if cuda and not torch.cuda.is_available():
raise Exception("No GPU found, please run without --cuda")
if opt.parallel:
opt.gpu = opt.local_rank
torch.cuda.set_device(opt.gpu)
torch.distributed.init_process_group(backend='nccl',
init_method='env://')
opt.world_size = torch.distributed.get_world_size()
if cuda:
torch.cuda.manual_seed(seed)
cudnn.benchmark = True
print("==========> Loading datasets")
test_dataset = DatasetFromFolder(opt.test,
transform=Compose([ToTensor()]),
training=False,
experiments="ShadowRemoval")
data_loader = DataLoader(dataset=test_dataset,
num_workers=4,
batch_size=opt.batchSize,
pin_memory=True,
shuffle=False)
print("==========> Building model")
netG = ShadowRemoval(channels=64)
print("=========> Building criterion")
loss_smooth_l1 = nn.SmoothL1Loss()
loss_l1 = nn.L1Loss()
loss_mse = torch.nn.MSELoss()
loss_bce = torch.nn.BCELoss()
loss_perceptual = perceptual()
instance_ssim = SSIM(reduction='mean', window_size=7)
rgb2yuv = rgb2yuv()
curriculum_ssim_mask = CLBase(lossfunc=nn.BCELoss(reduce=False))
curriculum_ssim_clean = CLBase()
# optionally copy weights from a checkpoint
if opt.pretrained and opt.continue_training:
if os.path.isfile(opt.pretrained):
print("=> loading model '{}'".format(opt.pretrained))
weights = torch.load(opt.pretrained)
netG.load_state_dict(weights['state_dict'])
else:
print("=> no model found at '{}'".format(opt.pretrained))
print("==========> Setting Optimizer")
optimizerG = optim.Adam(filter(lambda p: p.requires_grad,
netG.parameters()),
lr=opt.lr_g,
betas=(0.9, 0.999))
print("==========> Setting GPU")
if cuda:
netG = netG.cuda()
instance_ssim = instance_ssim.cuda()
loss_smooth_l1 = loss_smooth_l1.cuda()
loss_mse = loss_mse.cuda()
loss_l1 = loss_l1.cuda()
loss_bce = loss_bce.cuda()
curriculum_ssim_mask = curriculum_ssim_mask.cuda()
curriculum_ssim_clean = curriculum_ssim_clean.cuda()
loss_perceptual = loss_perceptual.cuda()
rgb2yuv = rgb2yuv.cuda()
if opt.acceleration:
print("FP 16 Trianing")
amp.register_float_function(torch, 'sigmoid')
netG, optimizerG = amp.initialize(netG,
optimizerG,
opt_level=opt.opt_level)
else:
netG = netG.cpu()
instance_ssim = instance_ssim.cpu()
loss_smooth_l1 = loss_smooth_l1.cpu()
loss_mse = loss_mse.cpu()
loss_l1 = loss_l1.cpu()
loss_bce = loss_bce.cpu()
curriculum_ssim = curriculum_ssim.cpu()
loss_perceptual = loss_perceptual.cpu()
rgb2yuv = rgb2yuv.cpu()
test(data_loader, netG)
def test(test_data_loader, netG):
psnrs_clean = []
psnrs_mask = []
ssims_clean = []
ssims_mask = []
bces_mask = []
if not os.path.exists("datasets/removals/"):
os.mkdir("datasets/removals/")
if not os.path.exists("datasets/removals/clean/"):
os.mkdir("datasets/removals/clean/")
if not os.path.exists("datasets/removals/mask/"):
os.mkdir("datasets/removals/mask")
if not os.path.exists("datasets/removals/shadow/"):
os.mkdir("datasets/removals/shadow/")
for iteration, batch in enumerate(test_data_loader, 1):
netG.eval()
steps = iteration
data_shadow, data_mask, data_clean = \
Variable(batch[0]), \
Variable(batch[1]), \
Variable(batch[2], requires_grad=False)
if opt.cuda:
data_clean = data_clean.cuda()
data_mask = data_mask.cuda()
data_shadow = data_shadow.cuda()
else:
data_clean = data_clean.cpu()
data_mask = data_mask.cpu()
data_shadow = data_shadow.cpu()
with torch.no_grad():
predict_clean, predict_mask = netG(data_shadow)
psnr_mask = (10 * torch.log10(
1.0 / loss_mse(predict_mask, data_mask.detach()))).mean()
psnr_clean = (10 * torch.log10(1.0 / loss_mse(
rgb2yuv(predict_clean)[:, 0].unsqueeze(1),
rgb2yuv(data_clean)[:, 0].unsqueeze(1)).detach())).mean()
psnrs_mask.append(psnr_mask)
psnrs_clean.append(psnr_clean)
ssim_mask = 1 - instance_ssim(predict_mask, data_mask).mean()
ssim_clean = 1 - instance_ssim(
rgb2yuv(predict_clean)[:, 0].unsqueeze(1),
rgb2yuv(data_clean)[:, 0].unsqueeze(1)).mean()
ssims_mask.append(ssim_mask)
ssims_clean.append(ssim_clean)
bce_mask = loss_bce(predict_mask.clamp(0, 1), data_mask)
bces_mask.append(bce_mask)
data_clean = Image.fromarray(
np.uint8(
torch.cat((data_clean, data_mask.expand_as(data_clean),
predict_clean),
dim=3).cpu().data[0].permute(1, 2, 0).mul(255)))
data_mask = Image.fromarray(
np.uint8(
torch.cat((data_mask.expand_as(data_shadow),
predict_mask.expand_as(data_shadow)),
dim=3).cpu().data[0].permute(1, 2, 0).mul(255)))
data_shadow = Image.fromarray(
np.uint8(data_shadow.cpu().data[0].permute(1, 2, 0).mul(255)))
data_shadow.save("datasets/removals/shadow/{}_{}_{}.jpg".format(
iteration, psnr_clean, ssim_clean))
data_clean.save("datasets/removals/clean/{}_{}_{}.jpg".format(
iteration, psnr_clean, ssim_clean))
data_mask.save("datasets/removals/mask/{}_{}_{}.jpg".format(
iteration, psnr_clean, ssim_clean))
print("processing {}th".format(iteration))
print("psnr_mask:{} psnr_clean:{} ssim_mask:{} ssim_clean:{} bce_mask:{}".
format(
sum(psnrs_mask) / len(psnrs_mask),
sum(psnrs_clean) / len(psnrs_clean),
sum(ssims_mask) / len(ssims_mask),
sum(ssims_clean) / len(ssims_clean),
sum(bces_mask) / len(bces_mask)))
def test(test_data_loader, netG):
psnrs_clean = []
psnrs_mask = []
ssims_clean = []
ssims_mask = []
bces_mask = []
for iteration, batch in enumerate(test_data_loader, 1):
netG.eval()
steps = iteration
data_shadow, data_mask, data_clean = \
Variable(batch[0]), \
Variable(batch[1]), \
Variable(batch[2], requires_grad=False)
if opt.cuda:
data_clean = data_clean.cuda()
data_mask = data_mask.cuda()
data_shadow = data_shadow.cuda()
else:
data_clean = data_clean.cpu()
data_mask = data_mask.cpu()
data_shadow = data_shadow.cpu()
predict_clean, predict_mask = netG(data_shadow)
with torch.no_grad():
psnr_mask = (10 * torch.log10(
1.0 / loss_mse(predict_mask, data_mask.detach()))).mean()
psnr_clean = (10 * torch.log10(1.0 / loss_mse(
rgb2yuv(predict_clean)[:, 0].unsqueeze(1),
rgb2yuv(data_clean)[:, 0].unsqueeze(1)).detach())).mean()
psnrs_mask.append(psnr_mask)
psnrs_clean.append(psnr_clean)
ssim_mask = 1 - instance_ssim(predict_mask, data_mask).mean()
ssim_clean = 1 - instance_ssim(
rgb2yuv(predict_clean)[:, 0].unsqueeze(1),
rgb2yuv(data_clean)[:, 0].unsqueeze(1)).mean()
ssims_mask.append(ssim_mask)
ssims_clean.append(ssim_clean)
bce_mask = loss_bce(predict_mask.clamp(0, 1), data_mask)
bces_mask.append(bce_mask)
print("psnr_mask:{} psnr_clean:{} ssim_mask:{} ssim_clean:{} bce_mask:{}".
format(
sum(psnrs_mask) / len(psnrs_mask),
sum(psnrs_clean) / len(psnrs_clean),
sum(ssims_mask) / len(ssims_mask),
sum(ssims_clean) / len(ssims_clean),
sum(bces_mask) / len(bces_mask)))
def train(train_data_loader, netG, netD, optimizerG, optimizerD, epoch,
logger):
print("epoch =", epoch, "lr =", optimizerG.param_groups[0]["lr"])
Wasserstein_D = torch.zeros(1)
penalty_lambda = 100
psnrs_clean = []
psnrs_mask = []
ssims_clean = []
ssims_mask = []
bces_mask = []
atomAge = 1.0 / len(train_data_loader)
loss_function_mask = loss_function(smoothl1=False,
l1=False,
mse=False,
instance_ssim=False,
bce=True,
perceptual_loss=False,
yuv=False)
loss_function_clean = loss_function(smoothl1=False,
l1=True,
mse=False,
instance_ssim=False,
bce=False,
perceptual_loss=True,
yuv=True)
for iteration, batch in enumerate(train_data_loader, 1):
netG.train()
netD.train()
optimizerD.zero_grad()
steps = len(train_data_loader) * (epoch - 1) + iteration
data_shadow, data_mask, data_clean = \
Variable(batch[0]), \
Variable(batch[1]), \
Variable(batch[2], requires_grad=False)
if opt.cuda:
data_clean = data_clean.cuda()
data_mask = data_mask.cuda()
data_shadow = data_shadow.cuda()
else:
data_clean = data_clean.cuda()
data_mask = data_mask.cuda()
data_shadow = data_shadow.cuda()
########################
# (1) Update D network Every Two Iteration#
########################
# train with fake
predict_clean, predict_mask = netG(data_shadow)
# curriculum learning
c_predict_clean, c_data_clean = predict_clean, data_clean
c_predict_mask, c_data_mask = predict_mask, data_mask
if iteration % 2 == 0:
for p in netD.parameters(): # reset requires_grad to True
p.requires_grad = True # they are set to False again in netG update.
# train with real
D_real = netD(
torch.cat((data_shadow, c_data_clean), dim=1).detach())
D_fake = netD(
torch.cat((data_shadow, c_predict_clean), dim=1).detach())
# train with gradient penalty and curriculum regularization
gradient_penalty = calc_gradient_penalty(
netD, data_shadow.data, data_clean.data, data_mask.data,
c_predict_clean.data.detach(), c_predict_mask.data.detach(),
penalty_lambda)
D_loss = 0.5 * (
(D_real - 1).mean() + D_fake.mean()) + gradient_penalty
Wasserstein_D = (D_fake.mean() - D_real.mean())
netD.zero_grad()
if opt.acceleration:
with amp.scale_loss(D_loss, optimizerD) as D_loss_scaled:
D_loss_scaled.clamp(1e-8, 1e8)
D_loss_scaled.backward(retain_graph=True)
torch.nn.utils.clip_grad_norm_(amp.master_params(optimizerD),
5)
else:
D_loss.backward(retain_graph=True)
optimizerD.step()
########################
# (2) update G network #
########################
optimizerG.zero_grad()
for p in netD.parameters(): # reset requires_grad to False
p.requires_grad = False # they are set to True again in netD update
netG.zero_grad()
G_loss = (netD(torch.cat(
(data_shadow, c_predict_clean), dim=1)) - 1).mean()
loss_clean = loss_function_clean(reconstructed=c_predict_clean,
target=c_data_clean,
original_r=predict_clean,
original_t=data_clean).mean()
loss_mask = loss_function_mask(reconstructed=c_predict_mask.clamp(
0, 1),
target=c_data_mask.clamp(0, 1),
original_r=predict_mask.clamp(0, 1),
original_t=data_mask.clamp(0,
1)).mean()
with torch.no_grad():
psnr_mask = 10 * torch.log10(
1.0 / loss_mse(predict_mask, data_mask.detach())).mean()
psnr_clean = (10 * torch.log10(1.0 / loss_mse(
rgb2yuv(predict_clean)[:, 0].unsqueeze(1),
rgb2yuv(data_clean)[:, 0].unsqueeze(1)).detach())).mean()
psnrs_mask.append(psnr_mask.mean())
psnrs_clean.append(psnr_clean.mean())
ssim_mask = 1 - instance_ssim(predict_mask, data_mask).mean()
ssim_clean = 1 - instance_ssim(rgb2yuv(predict_clean),
rgb2yuv(data_clean)).mean()
ssims_mask.append(ssim_mask)
ssims_clean.append(ssim_clean)
bce_mask = loss_bce(predict_mask.clamp(0, 1), data_mask).mean()
bces_mask.append(bce_mask)
loss = 100 * (0.2 * loss_clean) + loss_mask + G_loss
if opt.acceleration:
with amp.scale_loss(loss, optimizerG) as loss_scaled:
loss_scaled.clamp(1e-8, 1e8)
loss_scaled.backward()
torch.nn.utils.clip_grad_norm_(amp.master_params(optimizerG), 5)
else:
loss.backward()
optimizerG.step()
if (iteration - 1) % 10 == 0:
print(
"===> Epoch[{}]({}/{}): Loss_clean:{:.6f}: Loss_mask:{:.6f} : Psnr_clean:{:.2f} : Psnr_mask:{:.2f} : SSIM_clean:{:2f} : SSIM_mask:{:2f} : BCE_mask:{:2f} : Wasserstein:{} : Lr:{:6f}"
.format(epoch, iteration, len(train_data_loader),
loss_clean.mean(), loss_mask.mean(), psnr_clean.mean(),
psnr_mask.mean(), ssim_clean.mean(), ssim_mask.mean(),
bce_mask.mean(), Wasserstein_D.mean(),
optimizerG.param_groups[0]['lr']))
logger.add_scalar('loss', loss.item(), steps)
logger.add_scalar('psnr', psnr_clean.item(), steps)
logger.add_scalar('ssim', ssim_clean.item(), steps)
logger.add_scalar('bce', bce_mask.item(), steps)
show = []
if (iteration - 1) % 10 == 0:
for idx, tensor in enumerate(
zip(
data_clean.data.cpu(),
data_mask.data.cpu().expand_as(data_clean),
data_shadow.data.cpu(),
predict_clean.data.cpu().clamp(0, 1),
predict_mask.data.cpu().expand_as(predict_clean).clamp(
0, 1))):
if idx > 1:
break
show.extend(
[tensor[0], tensor[1], tensor[2], tensor[3], tensor[4]])
show = torch.stack(show, 0)
show = make_grid(show, nrow=5, padding=5)
logger.add_image('Comparison_nEpochs:{}'.format(epoch), show)
def main():
global opt, name, logger, netG, netD, vgg, curriculum_ssim_mask, curriculum_ssim_clean, loss_mse, rgb2yuv, instance_ssim, loss_bce
opt = parser.parse_args()
name = "ShadowRemoval"
print(opt)
# Tag_ResidualBlocks_BatchSize
logger = SummaryWriter("./runs_sr/" +
time.strftime("/%Y-%m-%d-%H/", time.localtime()))
cuda = opt.cuda
if 'WORLD_SIZE' in os.environ:
opt.distributed = int(os.environ['WORLD_SIZE']) > 1
if cuda and not torch.cuda.is_available():
raise Exception("No GPU found, please run without --cuda")
if opt.distributed:
opt.gpu = opt.local_rank
torch.cuda.set_device(opt.gpu)
torch.distributed.init_process_group(backend='nccl',
init_method='env://')
opt.world_size = torch.distributed.get_world_size()
seed = 1334
torch.manual_seed(seed)
if cuda:
torch.cuda.manual_seed(seed)
cudnn.benchmark = True
print("==========> Loading datasets")
train_dataset = DatasetFromFolder(
opt.train,
transform=Compose([ToTensor()]),
training=True,
experiments="ShadowRemoval",
)
test_dataset = DatasetFromFolder(opt.test,
transform=Compose([ToTensor()]),
training=False,
experiments="ShadowRemoval")
train_data_loader = DataLoader(dataset=train_dataset,
num_workers=4,
batch_size=opt.batchSize,
pin_memory=True,
shuffle=True)
test_data_loader = DataLoader(dataset=test_dataset,
num_workers=4,
batch_size=opt.batchSize,
pin_memory=True,
shuffle=False)
print("==========> Building model")
netG = ShadowRemoval(channels=64)
netD = Discrimator(in_channels=6, channels=64, depth=3)
print("=========> Building criterion")
loss_smooth_l1 = nn.SmoothL1Loss()
loss_l1 = nn.L1Loss()
loss_mse = torch.nn.MSELoss()
loss_bce = torch.nn.BCELoss()
loss_perceptual = perceptual()
instance_ssim = SSIM(reduction='mean', window_size=7)
rgb2yuv = rgb2yuv()
curriculum_ssim_mask = CLBase(lossfunc=nn.BCELoss(reduce=False))
curriculum_ssim_clean = CLBase()
# optionally copy weights from a checkpoint
if opt.pretrained and opt.continue_training:
if os.path.isfile(opt.pretrained):
print("=> loading model '{}'".format(opt.pretrained))
weights = torch.load(opt.pretrained)
netG.load_state_dict(weights['state_dict'])
else:
print("=> no model found at '{}'".format(opt.pretrained))
print("==========> Setting Optimizer")
optimizerG = optim.Adam(filter(lambda p: p.requires_grad,
netG.parameters()),
lr=opt.lr_g,
betas=(0.9, 0.999))
#optimizerD = optim.Adam(filter(lambda p: p.requires_grad, netD.module.parameters() if opt.parallel else netD.parameters()), lr=opt.lr_d, betas = (0.5, 0.999))
optimizerD = optim.SGD(filter(lambda p: p.requires_grad,
netD.parameters()),
lr=opt.lr_d)
print("==========> Setting GPU")
if cuda:
netG = netG.cuda()
netD = netD.cuda()
instance_ssim = instance_ssim.cuda()
loss_smooth_l1 = loss_smooth_l1.cuda()
loss_mse = loss_mse.cuda()
loss_l1 = loss_l1.cuda()
loss_bce = loss_bce.cuda()
curriculum_ssim_mask = curriculum_ssim_mask.cuda()
curriculum_ssim_clean = curriculum_ssim_clean.cuda()
loss_perceptual = loss_perceptual.cuda()
rgb2yuv = rgb2yuv.cuda()
if opt.acceleration:
print("FP 16 Trianing")
amp.register_float_function(torch, 'sigmoid')
[netD,
netG], [optimizerD,
optimizerG] = amp.initialize([netD, netG],
[optimizerD, optimizerG],
opt_level=opt.opt_level)
if opt.parallel:
print("Parallel Training")
netG = nn.DataParallel(netG)
netD = nn.DataParallel(netD)
elif opt.distributed:
netG = DDP(netG, delay_allreduce=True)
netD = DDP(netD, delay_allreduce=True)
else:
netG = netG.cpu()
netD = netD.cpu()
instance_ssim = instance_ssim.cpu()
loss_smooth_l1 = loss_smooth_l1.cpu()
loss_mse = loss_mse.cpu()
loss_l1 = loss_l1.cpu()
loss_bce = loss_bce.cpu()
curriculum_ssim = curriculum_ssim.cpu()
loss_perceptual = loss_perceptual.cpu()
rgb2yuv = rgb2yuv.cpu()
lr_schedulerG = optim.lr_scheduler.CosineAnnealingLR(optimizerG,
opt.epoch,
eta_min=1e-7)
lr_schedulerD = optim.lr_scheduler.CosineAnnealingLR(optimizerD,
opt.epoch,
eta_min=1e-7)
print("==========> Training")
for epoch in range(opt.epoch + 1):
train(train_data_loader,
netG,
netD,
optimizerG,
optimizerD,
epoch,
logger=logger)
#test(test_data_loader, netG)
if epoch % opt.save_model_freq == 0:
save_checkpoint(netG, epoch, name, opt)
lr_schedulerG.step()
lr_schedulerD.step()
logger.close()
def main():
global opt, name, logger, netG, netD, vgg, curriculum_ssim, loss_mse, rgb2yuv, instance_ssim
opt = parser.parse_args()
name = "ShadowSyns"
print(opt)
# Tag_ResidualBlocks_BatchSize
cuda = opt.cuda
if cuda and not torch.cuda.is_available():
raise Exception("No GPU found, please run without --cuda")
seed = 1334
torch.manual_seed(seed)
if cuda:
torch.cuda.manual_seed(seed)
cudnn.benchmark = True
print("==========> Loading datasets")
train_dataset = DatasetFromFolder(
opt.train,
transform=Compose([ToTensor()]),
training = True,
)
test_dataset = DatasetFromFolder(
opt.test,
transform=Compose([ToTensor()]),
training = False,
)
train_data_loader = DataLoader(dataset=train_dataset, num_workers=4, batch_size=opt.batchSize,
pin_memory=True, shuffle=True)
test_data_loader = DataLoader(dataset=test_dataset, num_workers=4, batch_size=opt.batchSize,
pin_memory=True, shuffle=True)
print("==========> Building model")
netG = ShadowMattingNet(channels =64, depth = 9)
netD = Discrimator(in_channels = 7, channels = 64, depth = 5)
print("=========> Building criterion")
loss_smooth_l1 = nn.SmoothL1Loss()
loss_l1 = nn.L1Loss()
loss_mse = torch.nn.MSELoss()
instance_ssim = SSIM(reduction = 'mean', window_size = 7)
curriculum_ssim = CLBase()
loss_perceptual = perceptual()
rgb2yuv = rgb2yuv()
# optionally copy weights from a checkpoint
if opt.pretrained:
if os.path.isfile(opt.pretrained):
print("=> loading model '{}'".format(opt.pretrained))
weights = torch.load(opt.pretrained)
netG.load_state_dict(weights['state_dict_g'])
else:
print("=> no model found at '{}'".format(opt.pretrained))
print("==========> Setting GPU")
if cuda:
if opt.parallel:
netG = nn.DataParallel(netG, [0, 1, 2, 3]).cuda()
netD = nn.DataParallel(netD, [0, 1, 2, 3]).cuda()
instance_ssim = nn.DataParallel(instance_ssim, [0, 1, 2, 3]).cuda()
loss_smooth_l1 = nn.DataParallel(loss_smooth_l1, [0, 1, 2, 3]).cuda()
loss_mse = nn.DataParallel(loss_mse, [0, 1, 2, 3]).cuda()
loss_l1 = nn.DataParallel(loss_l1, [0, 1, 2, 3]).cuda()
curriculum_ssim = nn.DataParallel(curriculum_ssim, [0, 1, 2, 3]).cuda()
rgb2yuv = nn.DataParallel(rgb2yuv, [0, 1, 2, 3]).cuda()
else:
netG = netG.cuda()
netD = netD.cuda()
instance_ssim = instance_ssim.cuda()
loss_smooth_l1 = loss_smooth_l1.cuda()
loss_mse = loss_mse.cuda()
loss_l1 = loss_l1.cuda()
curriculum_ssim = curriculum_ssim.cuda()
loss_perceptual = loss_perceptual.cuda()
rgb2yuv = rgb2yuv.cuda()
else:
netG = netG.cpu()
netD = netD.cpu()
instance_ssim = instance_ssim.cpu()
loss_smooth_l1 = loss_smooth_l1.cpu()
loss_mse = loss_mse.cpu()
loss_l1 = loss_l1.cpu()
curriculum_ssim = curriculum_ssim.cpu()
loss_perceptual = loss_perceptual.cpu()
rgb2yuv = rgb2yuv.cpu()
print("==========> Setting Optimizer")
optimizerG = optim.Adam(filter(lambda p: p.requires_grad, netG.module.parameters() if opt.parallel else netG.parameters()), lr=opt.lr_g, betas = (0.5, 0.99))
#optimizerD = optim.Adam(filter(lambda p: p.requires_grad, netD.module.parameters() if opt.parallel else netD.parameters()), lr = opt.lr_d, betas = (0.5, 0.999))
optimizerD = optim.SGD(filter(lambda p: p.requires_grad, netD.module.parameters() if opt.parallel else netD.parameters()), lr = opt.lr_d)
lr_schedulerG = optim.lr_scheduler.CosineAnnealingLR(optimizerG, opt.epoch, eta_min = 1e-7)
lr_schedulerD = optim.lr_scheduler.CosineAnnealingLR(optimizerD, opt.epoch, eta_min = 1e-7)
print("==========> Training")
for epoch in range(opt.epoch + 1):
train(train_data_loader, netG, netD, optimizerG, optimizerD, epoch)
test(test_data_loader, netG)
if epoch % opt.save_model_freq == 0:
save_checkpoint(netG, epoch, name)
lr_schedulerG.step()
lr_schedulerD.step()
def train(training_data_loader, netG, netD, optimizerG, optimizerD, epoch):
print("epoch =", epoch, "lr =", optimizerG.param_groups[0]["lr"])
Wasserstein_D = torch.zeros(1)
penalty_lambda = 50
psnrs = []
ssims = []
atomAge = 1.0 / len(training_data_loader)
logger = SummaryWriter("./runs_ss/" + time.strftime("/%Y-%m-%d-%H/", time.localtime()))
loss_function_main = loss_function(smoothl1 = False, l1 = True, mse = False, instance_ssim = True, perceptual_loss = True)
for iteration, batch in enumerate(training_data_loader, 1):
netG.train()
netD.train()
optimizerD.zero_grad()
steps = len(training_data_loader) * (epoch-1) + iteration
data_clean, data_mask, data_shadow = \
Variable(batch[0], requires_grad=True), \
Variable(batch[1], requires_grad=True), \
Variable(batch[2], requires_grad=False)
if opt.cuda:
data_clean = data_clean.cuda()
data_mask = data_mask.cuda()
data_shadow = data_shadow.cuda()
else:
data_clean = data_clean.cpu()
data_mask = data_mask.cpu()
data_shadow = data_shadow.cpu()
########################
# (1) Update D network Every Two Iteration#
########################
# train with fake
data_predict = netG(data_clean, data_mask)
# curriculum learning
#c_data_predict, c_data_shadow, weight = curriculum_ssim(data_predict, data_shadow.detach(), epoch + atomAge * iteration)
c_data_predict, c_data_shadow = data_predict, data_shadow
if iteration % 2 == 0:
for p in netD.parameters(): # reset requires_grad to True
p.requires_grad = True # they are set to False again in netG update.
# train with real
D_real = netD(torch.cat((data_mask, data_clean, c_data_shadow), dim=1))
D_fake = netD(torch.cat((data_mask, data_clean, c_data_predict.detach()), dim=1))
#D_real = netD(((c_data_shadow)))
#D_fake = netD(((c_data_predict.detach())))
# train with gradient penalty and curriculum regularization
gradient_penalty = calc_gradient_penalty(netD, data_clean.data, data_mask.data, c_data_predict.data, c_data_shadow.data, penalty_lambda)
D_loss = (0.5*((D_real-1).mean() + D_fake.mean()) + gradient_penalty)
Wasserstein_D = (D_fake.mean() - D_real.mean())
netD.zero_grad()
D_loss.backward(retain_graph = True)
optimizerD.step()
########################
# (2) update G network #
########################
optimizerG.zero_grad()
for p in netD.parameters(): # reset requires_grad to False
p.requires_grad = False # they are set to True again in netD update
netG.zero_grad()
G_loss = (netD(torch.cat((data_mask, data_clean, c_data_predict), dim=1)) - 1).mean()
#G_loss = (netD(c_data_predict) - 1).mean().pow(2)
loss_shadow = loss_function_main(reconstructed = c_data_predict, target = c_data_shadow, original_r = data_predict, original_t = data_shadow).mean()
with torch.no_grad():
psnr = (10 * torch.log10(1.0 / loss_mse(rgb2yuv(data_predict)[:,0].unsqueeze(1), rgb2yuv(data_shadow)[:,0].unsqueeze(1)).detach())).mean()
psnrs.append(psnr)
ssim = 1 - instance_ssim(data_predict, data_shadow).mean()
ssims.append(ssim)
loss = 100*loss_shadow + (G_loss)
loss.backward()
optimizerG.step()
if (iteration-1) % 10 == 0:
print("===> Epoch[{}]({}/{}): Loss_shadow:{:.6f} : Psnr:{:.2f} : SSIM:{:2f} : Wasserstein:{} : Lr:{:6f}".format(epoch, iteration, len(training_data_loader), loss_shadow, psnr, ssim, Wasserstein_D.mean(), optimizerG.param_groups[0]['lr']))
logger.add_scalar('loss', loss.item(), steps)
logger.add_scalar('psnr', psnr.item(), steps)
logger.add_scalar('ssim', ssim.item(), steps)
show = []
if (iteration-1) % 10 == 0:
for idx, tensor in enumerate(zip(data_clean.data.cpu(), data_mask.data.cpu().expand_as(data_clean), data_shadow.data.cpu(), data_predict.data.cpu().clamp(0,1))):
if idx >1:
break
show.extend([tensor[0], tensor[1], tensor[2], tensor[3]])
show = torch.stack(show,0)
show = make_grid(show, nrow = 4, padding = 5)
logger.add_image('Comparison_nEpochs:{}'.format(epoch), show)
logger.close()
