LR = 0.001
param_noise = False
elif NET_TYPE == 'ResNet':
net = ResNet(input_depth, img_np.shape[0], 8, 32, need_sigmoid=True, act_fun='LeakyReLU')
LR = 0.001
param_noise = False
else:
assert False
else:
assert False
net = net.type(dtype)
net_input = get_noise(input_depth, INPUT, img_np.shape[1:]).type(dtype)
# In[ ]:
# Compute number of parameters
s = sum(np.prod(list(p.size())) for p in net.parameters())
print ('Number of params: %d' % s)
# Loss
mse = torch.nn.MSELoss().type(dtype)
img_var = np_to_torch(img_np).type(dtype)
mask_var = np_to_torch(img_mask_np).type(dtype)
def perform_inpainting(self):
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype = torch.float32 # dtype = torch.cuda.FloatTensor
PLOT = self.args.plot
imsize = -1
dim_div_by = 64
# Path to input image, mask and output
img_path = (str(Path(__file__).resolve().parents[1]) +
"/data/Input data/" + self.image)
mask_path = (str(Path(__file__).resolve().parents[1]) +
"/data/Input data/" + self.mask)
if self.args.tuning == "basic":
if self.vase_or_kate_or_library == "library":
folder = (
str(Path(__file__).resolve().parents[1]) +
"/data/Output data/Hyperparameter optimization/Basic/" +
self.vase_or_kate_or_library + "/" + self.NET_TYPE + "/")
Path(folder).mkdir(parents=True, exist_ok=True)
else:
folder = (
str(Path(__file__).resolve().parents[1]) +
"/data/Output data/Hyperparameter optimization/Basic/" +
self.vase_or_kate_or_library + "/")
Path(folder).mkdir(parents=True, exist_ok=True)
outp_path = folder + "/plotout"
elif self.args.tuning == "advanced":
print(
"lr =",
self.lr,
"param_noise =",
self.param_noise,
"reg_noise_std =",
self.reg_noise_std,
)
folder = (
str(Path(__file__).resolve().parents[1]) +
"/data/Output data/Hyperparameter optimization/Advanced/" +
self.image + "/lr=" + str(self.lr) + ", param_noise=" +
str(self.param_noise) + ", reg_noise_std=" +
str(self.reg_noise_std) + "/")
Path(folder).mkdir(parents=True, exist_ok=True)
outp_path = folder + "/plotout"
else:
folder = (str(Path(__file__).resolve().parents[1]) +
"/data/Output data/" + self.image.split(".")[0] + "/")
Path(folder).mkdir(parents=True, exist_ok=True)
outp_path = folder + "/plotout"
# Load mask
img_pil, img_np = get_image(img_path, imsize)
img_mask_pil, img_mask_np = get_image(mask_path, imsize)
# Center crop
img_mask_pil = crop_image(img_mask_pil, dim_div_by)
img_pil = crop_image(img_pil, dim_div_by)
img_np = pil_to_np(img_pil)
img_mask_np = pil_to_np(img_mask_pil)
# Visualize
if PLOT:
plot_image_grid([img_np, img_mask_np, img_mask_np * img_np], 3, 11)
# Setup
pad = "reflection" # 'zero'
OPT_OVER = "net"
OPTIMIZER = "adam"
num_iter = self.args.num_iter
if self.args.tuning == "advanced":
save_every = int(num_iter / 4)
else:
save_every = self.args.save_every
if self.vase_or_kate_or_library == "vase":
INPUT = "meshgrid"
input_depth = 2
LR = self.lr if self.lr else 0.01
# num_iter = 5001
param_noise = self.param_noise if self.param_noise else False
# save_every = 50
figsize = 32 # changed from 5
reg_noise_std = self.reg_noise_std if self.reg_noise_std else 0.03
net = (skip(
input_depth,
img_np.shape[0],
num_channels_down=[128] * 5,
num_channels_up=[128] * 5,
num_channels_skip=[0] * 5,
upsample_mode="nearest",
filter_skip_size=1,
filter_size_up=3,
filter_size_down=3,
need_sigmoid=True,
need_bias=True,
pad=pad,
act_fun="LeakyReLU",
).type(dtype).to(device))
elif self.vase_or_kate_or_library == "kate":
# Same params and net as in super-resolution and denoising
INPUT = "noise"
input_depth = 32
# num_iter = 6001
LR = 0.01
param_noise = False
# save_every = 50
figsize = 5
reg_noise_std = 0.03
net = (skip(
input_depth,
img_np.shape[0],
num_channels_down=[128] * 5,
num_channels_up=[128] * 5,
num_channels_skip=[128] * 5,
filter_size_up=3,
filter_size_down=3,
upsample_mode="nearest",
filter_skip_size=1,
need_sigmoid=True,
need_bias=True,
pad=pad,
act_fun="LeakyReLU",
).type(dtype).to(device))
elif self.vase_or_kate_or_library == "library":
INPUT = "noise"
input_depth = 1
# num_iter = 3001
# save_every = 50
figsize = 8
reg_noise_std = 0.00
param_noise = True
if "skip" in self.NET_TYPE:
depth = int(self.NET_TYPE[-1])
net = skip(
input_depth,
img_np.shape[0],
num_channels_down=[16, 32, 64, 128, 128, 128][:depth],
num_channels_up=[16, 32, 64, 128, 128, 128][:depth],
num_channels_skip=[0, 0, 0, 0, 0, 0][:depth],
filter_size_up=3,
filter_size_down=5,
filter_skip_size=1,
upsample_mode="nearest", # downsample_mode='avg',
need1x1_up=False,
need_sigmoid=True,
need_bias=True,
pad=pad,
act_fun="LeakyReLU",
)
LR = 0.01
elif self.NET_TYPE == "UNET":
net = UNet(
num_input_channels=input_depth,
num_output_channels=3,
feature_scale=8,
more_layers=1,
concat_x=False,
upsample_mode="deconv",
pad="zero",
norm_layer=torch.nn.InstanceNorm2d,
need_sigmoid=True,
need_bias=True,
)
LR = 0.001
param_noise = False
elif self.NET_TYPE == "ResNet":
net = ResNet(
input_depth,
img_np.shape[0],
8,
32,
need_sigmoid=True,
act_fun="LeakyReLU",
)
LR = 0.001
param_noise = False
else:
assert False
else:
assert False
net = net.type(dtype).to(device)
net_input = (get_noise(input_depth, INPUT,
img_np.shape[1:]).type(dtype).to(device))
# Compute number of parameters
s = sum(np.prod(list(p.size())) for p in net.parameters())
# print("Number of params: %d" % s)
# Loss
mse = torch.nn.MSELoss().type(dtype).to(device)
img_var = np_to_torch(img_np).type(dtype).to(device)
mask_var = np_to_torch(img_mask_np).type(dtype).to(device)
# Main loop
def closure():
# global i
if param_noise:
for n in [x for x in net.parameters() if len(x.size()) == 4]:
n = n + n.detach().clone().normal_() * n.std() / 50
net_input = net_input_saved
if reg_noise_std > 0:
net_input = net_input_saved + (noise.normal_() * reg_noise_std)
out = net(net_input)
total_loss = mse(out * mask_var, img_var * mask_var)
total_loss.backward()
# print('Iteration %05d Loss %f' % (i, total_loss.item()), '\r', end='')
if self.i % save_every == 0 or self.i == num_iter - 1:
out_np = torch_to_np(out)
out_np = 255 * np.moveaxis(out_np, 0, 2)
out_np = out_np.astype(np.uint8)
filep = outp_path + str(self.i) + ".png"
image = Image.fromarray(out_np)
image.save(filep)
self.i += 1
return total_loss
net_input_saved = net_input.detach().clone()
noise = net_input.detach().clone()
p = get_params(OPT_OVER, net, net_input)
optimize(OPTIMIZER, p, closure, LR, num_iter)
