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()