def demosaick_load_model(net_path=None, noiselevel=0.0, xtrans=False):
'''
computes the relative paths to the pretrained models (Caffe)
if net_path is specified then the provided path is used
if xtrans is set then the XtransNetwork model is used,
otherwise BayerNetwork or BayerNetworkNoise are used depending
on whether noiselevel>0
'''
here = os.path.dirname(os.path.abspath(__file__))
pretrained_base = here+'/pretrained_models/'
print("Loading Caffe weights")
if xtrans:
model_name = "XtransNetwork"
model_path = pretrained_base+'/xtrans/'
else:
if noiselevel==0.0:
model_name = "BayerNetwork"
model_path = pretrained_base+'/bayer/'
else:
model_name = "BayerNetworkNoise"
model_path = pretrained_base+'/bayer_noise/'
# if a network path is provided use it instead
if net_path: model_path = net_path
model_ref = modules.get({"model": model_name})
cvt = converter.Converter(model_path, model_name)
cvt.convert(model_ref)
for p in model_ref.parameters():
p.requires_grad = False
return model_ref
def load_net(net_path=None):
here = os.path.dirname(os.path.abspath(__file__))
pretrained_base = here + '/pretrained_models/'
print("Loading Caffe weights")
model_name = "BayerNetworkNoise"
model_path = pretrained_base + '/bayer_noise/'
# if a network path is provided use it instead
if net_path: model_path = net_path
model_ref = modules.get({"model": model_name})
cvt = converter.Converter(model_path, model_name)
cvt.convert(model_ref)
return model_ref
def main(args):
if args.fix_seed:
np.random.seed(0)
th.manual_seed(0)
meta_params = utils.Checkpointer.get_meta(args.model)
name = os.path.basename(args.model)
model = modules.get(meta_params["model"])
log.info("Model configuration: {}".format(meta_params["model"]))
checkpointer = utils.Checkpointer(
args.model, model, None,
meta_params=meta_params)
f, _ = checkpointer.load_latest()
log.info("Loading checkpoint {}".format(f))
for p in model.parameters():
p.requires_grad = False
model.cuda()
for path in os.listdir(args.data_dir):
if not os.path.splitext(path)[-1] == ".tif":
continue
im = cv2.imread(os.path.join(args.data_dir, path), -1).astype(np.float32)
im = im[args.offset_y:, args.offset_x:]
sz = 512
im = im[args.shift_x:args.shift_x+sz, args.shift_y:args.shift_y+sz]
im /= 2**16
im = np.stack([im]*3, 0)
im = dset.bayer_mosaic(im)
im = np.expand_dims(im, 0)
im = th.from_numpy(im).cuda()
h, w = im.shape[-2:]
tot_time = 0
sample = {"mosaic": im}
th.cuda.synchronize()
start = time.time()
kernels = []
out = model(sample, kernel_list=kernels)
r_kernels = []
for k in kernels:
bs, ksz, _, kh, kw = k.shape
k = k.permute([0, 3, 1, 4, 2]).contiguous().view(bs, 1, kh*ksz, kw*ksz)
r_kernels.append(k)
r_kernels = th.cat(r_kernels, 1).squeeze().permute([1, 2, 0])
th.cuda.synchronize()
tot_time += time.time()-start
oh, ow = out.shape[2:]
ch = (h-oh) // 2
cw = (w-ow) // 2
im = im[:, ch:-ch, cw:-cw]
out = out.squeeze(0)
tot_time *= 1000
log.info("Time {:.0f} ms".format(tot_time))
out = out.permute(1, 2, 0).cpu().numpy()
out = np.clip(out, 0, 1)
os.makedirs(args.output, exist_ok=True)
dst_path = os.path.join(args.output, path.replace(".tif", "_"+name+"_.png"))
skio.imsave(dst_path, out)
# r_kernels = th.clip(r_kernels, 0, 1)
r_kernels = r_kernels.cpu().numpy()
for c in range(3):
k_copy = np.copy(r_kernels[:, :, c])
k_copy /= np.abs(k_copy).max()
k_copy *= 0.5
k_copy += 0.5
k_copy = np.dstack([k_copy]*3)
dst_path = os.path.join(args.output, path.replace(".tif", "_kernels_{}.png".format(c)))
skio.imsave(dst_path, k_copy)
def main(args, model_params):
if args.fix_seed:
np.random.seed(0)
th.manual_seed(0)
# ------------ Set up datasets ----------------------------------------------
xforms = [dset.ToTensor()]
if args.green_only:
xforms.append(dset.GreenOnly())
xforms = transforms.Compose(xforms)
if args.xtrans:
data = dset.XtransDataset(args.data_dir,
transform=xforms,
augment=True,
linearize=args.linear)
else:
data = dset.BayerDataset(args.data_dir,
transform=xforms,
augment=True,
linearize=args.linear)
data[0]
if args.val_data is not None:
if args.xtrans:
val_data = dset.XtransDataset(args.val_data,
transform=xforms,
augment=False)
else:
val_data = dset.BayerDataset(args.val_data,
transform=xforms,
augment=False)
else:
val_data = None
# ---------------------------------------------------------------------------
model = modules.get(model_params)
log.info("Model configuration: {}".format(model_params))
if args.pretrained:
log.info("Loading Caffe weights")
if args.xtrans:
model_ref = modules.get({"model": "XtransNetwork"})
cvt = converter.Converter(args.pretrained, "XtransNetwork")
else:
model_ref = modules.get({"model": "BayerNetwork"})
cvt = converter.Converter(args.pretrained, "BayerNetwork")
cvt.convert(model_ref)
model_ref.cuda()
else:
model_ref = None
if args.green_only:
model = modules.GreenOnly(model)
model_ref = modules.GreenOnly(model_ref)
if args.subsample:
dx = 1
dy = 0
if args.xtrans:
period = 6
else:
period = 2
model = modules.Subsample(model, period, dx=dx, dy=dy)
model_ref = modules.Subsample(model_ref, period, dx=dx, dy=dy)
if args.linear:
model = modules.DeLinearize(model)
model_ref = modules.DeLinearize(model_ref)
name = os.path.basename(args.output)
cbacks = [
default_callbacks.LossCallback(env=name),
callbacks.DemosaicVizCallback(val_data,
model,
model_ref,
cuda=True,
shuffle=False,
env=name),
callbacks.PSNRCallback(env=name),
]
metrics = {"psnr": losses.PSNR(crop=4)}
log.info("Using {} loss".format(args.loss))
if args.loss == "l2":
criteria = {
"l2": losses.L2Loss(),
}
elif args.loss == "l1":
criteria = {
"l1": losses.L1Loss(),
}
elif args.loss == "gradient":
criteria = {
"gradient": losses.GradientLoss(),
}
elif args.loss == "laplacian":
criteria = {
"laplacian": losses.LaplacianLoss(),
}
elif args.loss == "vgg":
criteria = {
"vgg": losses.VGGLoss(),
}
else:
raise ValueError("not implemented")
optimizer = optim.Adam
optimizer_params = {}
if args.optimizer == "sgd":
optimizer = optim.SGD
optimizer_params = {"momentum": 0.9}
train_params = Trainer.Parameters(viz_step=100,
lr=args.lr,
batch_size=args.batch_size,
optimizer=optimizer,
optimizer_params=optimizer_params)
trainer = Trainer(data,
model,
criteria,
output=args.output,
params=train_params,
model_params=model_params,
verbose=args.debug,
callbacks=cbacks,
metrics=metrics,
valset=val_data,
cuda=True)
trainer.train()
def main(args):
if args.fix_seed:
np.random.seed(0)
th.manual_seed(0)
meta_params = utils.Checkpointer.get_meta(args.model)
name = os.path.basename(args.model)
model = modules.get(meta_params["model"])
log.info("Model configuration: {}".format(meta_params["model"]))
if args.pretrained is not None:
if args.xtrans:
model_ref = modules.get({"model": "XtransNetwork"})
cvt = converter.Converter(args.pretrained, "XtransNetwork")
else:
model_ref = modules.get({"model": "BayerNetwork"})
log.info("Loading Caffe weights")
cvt = converter.Converter(args.pretrained, "BayerNetwork")
cvt.convert(model_ref)
for p in model_ref.parameters():
p.requires_grad = False
model_ref.cuda()
checkpointer = utils.Checkpointer(args.model,
model,
None,
meta_params=meta_params)
f, _ = checkpointer.load_latest()
log.info("Loading checkpoint {}".format(f))
for p in model.parameters():
p.requires_grad = False
model.cuda()
for path in os.listdir(args.data_dir):
if not os.path.splitext(path)[-1] == ".tif":
continue
im = cv2.imread(os.path.join(args.data_dir, path),
-1).astype(np.float32)
im = im[args.offset_y:, args.offset_x:]
# im = im[:3000, :3000]
im /= 2**16
im = np.stack([im] * 3, 0)
if args.xtrans:
im = dset.bayer_mosaic(im)
else:
im = dset.xtrans_mosaic(im)
im = np.expand_dims(im, 0)
im = th.from_numpy(im).cuda()
import ipdb
ipdb.set_trace()
_, _, h, w = im.shape
out_im = th.zeros(3, h, w).cuda()
out_ref = th.zeros(3, h, w).cuda()
tile_size = min(min(args.tile_size, h), w)
tot_time = 0
tot_time_ref = 0
if args.xtrans:
mod = 6
else:
mod = 2
tile_step = args.tile_step - (args.tile_step % mod)
tile_step = args.tile_step - (args.tile_step % mod)
for start_x in range(0, w, tile_step):
end_x = start_x + tile_size
if end_x > w:
# keep mosaic period
end_x = w
start_x = end_x - tile_size
start_x = start_x - (start_x % mod)
end_x = start_x + tile_size
for start_y in range(0, h, tile_step):
end_y = start_y + tile_size
if end_y > h:
end_y = h
start_y = end_y - tile_size
start_y = start_y - (start_y % mod)
end_y = start_y + tile_size
print(start_x, start_y)
sample = {"mosaic": im[:, :, start_y:end_y, start_x:end_x]}
th.cuda.synchronize()
start = time.time()
out = model(sample)
th.cuda.synchronize()
tot_time += time.time() - start
th.cuda.synchronize()
start = time.time()
outr = model_ref(sample)
th.cuda.synchronize()
tot_time_ref += time.time() - start
oh, ow = out.shape[2:]
ch = (tile_size - oh) // 2
cw = (tile_size - ow) // 2
out_im[:, start_y + ch:start_y + ch + oh,
start_x + cw:start_x + cw + ow] = out[0]
oh, ow = outr.shape[2:]
ch = (tile_size - oh) // 2
cw = (tile_size - ow) // 2
out_ref[:, start_y + ch:start_y + ch + oh,
start_x + cw:start_x + cw + ow] = outr[0]
tot_time *= 1000
tot_time_ref *= 1000
log.info("Time {:.0f} ms (ref {:.0f} ms)".format(
tot_time, tot_time_ref))
out_im = out_im.permute(1, 2, 0).cpu().numpy()
out_im = np.clip(out_im, 0, 1)
os.makedirs(args.output, exist_ok=True)
dst_path = os.path.join(args.output,
path.replace(".tif", "_" + name + "_.png"))
skio.imsave(dst_path, out_im)
out_ref = out_ref.permute(1, 2, 0).cpu().numpy()
out_ref = np.clip(out_ref, 0, 1)
dst_path = os.path.join(args.output, path.replace(".tif", "_ref.png"))
skio.imsave(dst_path, out_ref)
def main(args, model_params):
if args.fix_seed:
np.random.seed(0)
th.manual_seed(0)
# ------------ Set up datasets ----------------------------------------------
xforms = dset.ToTensor()
data = dset.BayerDataset(args.data_dir, transform=xforms, augment=True)
data[0]
if args.val_data is not None:
val_data = dset.BayerDataset(args.val_data,
transform=xforms,
augment=False)
else:
val_data = None
# ---------------------------------------------------------------------------
model = modules.get(model_params)
log.info("Model configuration: {}".format(model_params))
if args.pretrained:
log.info("Loading Caffe weights")
cvt = converter.Converter(args.pretrained, model_params["model"])
cvt.convert(model)
name = os.path.basename(args.output)
cbacks = [
default_callbacks.LossCallback(env=name),
callbacks.DemosaicVizCallback(val_data,
model,
cuda=True,
shuffle=True,
env=name),
]
metrics = {"psnr": losses.PSNR()}
log.info("Using {} loss".format(args.loss))
if args.loss == "l2":
criteria = {
"l2": losses.L2Loss(),
}
elif args.loss == "vgg":
criteria = {
"vgg": losses.VGGLoss(),
}
else:
raise ValueError("not implemented")
train_params = Trainer.Parameters(viz_step=100,
lr=args.lr,
batch_size=args.batch_size)
# optimizer=toptim.SVAG)
trainer = Trainer(data,
model,
criteria,
output=args.output,
params=train_params,
model_params=model_params,
verbose=args.debug,
callbacks=cbacks,
metrics=metrics,
valset=val_data,
cuda=True)
trainer.train()