def __getitem__(self, i):
g = imread(random.choice(self.d)) / 255
g = rgb2gray(g).astype(np.float32)
g = rand_crop(g, self.gs + self.ks - 1)
g = torch.from_numpy(g).view(1, *g.shape)
k = gen_kernel(self.ks).astype(np.float32)
k = torch.from_numpy(k)
s = random.uniform(1, 3)
s = torch.tensor((s, ), dtype=torch.float)
y = F.conv2d(g.view(1, *g.shape), k.view(1, 1, *k.shape)).squeeze(0)
assert y.shape[-1] == self.gs
# noise
if self.noise:
y += torch.randn_like(y) * s / 255
# show(y[0])
# edgetaping, todo convert to torch, and move to model stage
if self.edgetaper:
y = y.permute(1, 2, 0)
y = to_tensor(
edgetaper(pad_for_kernel(y.numpy(), k.numpy(), "edge"),
k.numpy())).astype(np.float32)
y = torch.from_numpy(y).squeeze(-1)
g = center_crop(g, self.gs)
# [1,284,284] [1,320,320] [37,37] [1]
return g, y, k, s
def forward(self, x):
blocks = []
for i, down in enumerate(self.down_path):
x = down(x)
if i != len(self.down_path)-1:
blocks.append(x)
if self.do_max_pool:
x = F.max_pool2d(x, 2)
else:
x = self.downsample_convs[i](x)
for i, up in enumerate(self.up_path):
x = up(x, blocks[-i-1])
seg_x = self.seg_conv(x)
if self.do_soft_max:
seg = self.soft_max(seg_x)
else:
seg = seg_x
if self.num_lands > 0:
if self.lands_block is not None:
x = self.lands_block(x)
x = torch.cat((x, util.center_crop(seg_x, x.shape)), dim=1)
heat_maps = self.lands_1x1(x)
return (seg, heat_maps)
else:
return seg
def test(m):
m.eval()
with torch.no_grad():
d = DataLoader(Sun(), 1)
losss = []
for i in tqdm(d):
g, y, k, s = [x.to(o.device) for x in i]
out = m([y, y, k, s])
out = crop(out, k)
out = center_crop(out, *g.shape[-2:])
loss = npsnr(out, g)
losss.append(-loss.detach().item())
log("psnr", losss[-1])
show(
torch.cat(
(center_crop(y, *g.shape[-2:])[0, 0], g[0, 0], out[0, 0]),
1))
log("psnr avg", sum(losss) / len(losss))
def train(m, p=None):
d = DataLoader(BSD3000(noise=False, edgetaper=False),
o.batch_size,
num_workers=o.num_workers)
optimizer = torch.optim.Adam(m.parameters(), lr=o.lr)
iter_num = len(d)
num = 0
losss = []
stage = 1 if not p else p.stage + 1
for epoch in range(o.epoch):
for i in tqdm(d):
g, y, k, s = [x.to(o.device) for x in i]
k = k.flip(1, 2)
x = torch.tensor(y, requires_grad=True)
if p:
with torch.no_grad():
x = p([x, y, k, s])
optimizer.zero_grad()
out = m([x, y, k, s])
log("out", out)
out = center_crop(out, *g.shape[-2:])
loss = npsnr(out, g)
loss.backward()
optimizer.step()
losss.append(loss.detach().item())
assert not isnan(losss[-1])
print("stage", stage, "epoch", epoch + 1)
log("loss", mean(losss[-5:]))
num += 1
# if num > (o.epoch * iter_num - 4):
if num % 20 == 0:
show(
torch.cat((center_crop(
y, *g.shape[-2:])[0, 0], g[0, 0], out[0, 0]), 1),
save=f"save/{stage:02}{epoch:02}.png",
)
plt.clf()
plt.plot(range(len(losss)), losss)
plt.xlabel("batch")
plt.ylabel("loss")
plt.title(f"{iter_num} iter x {o.epoch} epoch")
plt.savefig(f"save/{stage:02}loss.png")
def __getitem__(self, i):
# [img1xkernel1, img1xkernel2,...,img2xkernel1]
g = imread(self.d[i // 8]) / 255
g = rgb2gray(g).astype(np.float32)
k = self.k[i % 8]
g = torch.from_numpy(g)
s = torch.tensor((self.s, ), dtype=torch.float)
g = g.view(1, *g.shape)
# blur
y = F.conv2d(g.view(1, *g.shape), k.view(1, 1, *k.shape))
g = center_crop(g, 250)
y = center_crop(y, 250)[0]
# show(torch.cat((g.detach().cpu()[0], y.detach().cpu()[0,0]), 0))
# noise
# y += torch.randn_like(y) * s / 255
# edgetaping, todo convert to torch, and move to model stage
y = y.permute(1, 2, 0)
y = to_tensor(
edgetaper(pad_for_kernel(y.numpy(), k.numpy(), "edge"),
k.numpy())).astype(np.float32)
y = torch.from_numpy(y).squeeze(-1)
# [1,250,250] [1,267,267] [1,13,13] [1]
return g, y, k, s
def __init__(self, type="train"):
d = Path(f"data/BSR/BSDS500/data/images/{type}").glob("**/*")
self.d = [i for i in d if i.is_file()]
self.k = []
for i in range(1, 9):
k = imread(f"data/kernel/kernel{i}_groundtruth_kernel.png").astype(
np.float32) / 255
# crop to 13x13
k = np.clip(k, 0, 1)
k = center_crop(k, 13)
k /= np.sum(k)
self.k.append(torch.from_numpy(k))
# or pad to 27x27
# k = torch.from_numpy(k)
# k = k.clamp(0, 1)
# k = center_pad(k, 27)
# k /= k.sum()
# self.k.append(k)
self.s = 0.01 * 255
def img_process(self, imgs, frames_num):
images = np.zeros((frames_num, 224, 224, 3))
orig_imgs = np.zeros_like(images)
for i in range(frames_num):
next_image = imgs[i]
next_image = np.uint8(next_image)
scaled_img = cv2.resize(
next_image, (256, 256),
interpolation=cv2.INTER_LINEAR) # resize to 256x256
cropped_img = center_crop(scaled_img) # center crop 224x224
final_img = cv2.cvtColor(cropped_img, cv2.COLOR_BGR2RGB)
images[i] = final_img
orig_imgs[i] = cropped_img
torch_imgs = torch.from_numpy(images.transpose(3, 0, 1, 2))
torch_imgs = torch_imgs.float() / 255.0
mean_3d = [124 / 255, 117 / 255, 104 / 255]
std_3d = [0.229, 0.224, 0.225]
for t, m, s in zip(torch_imgs, mean_3d, std_3d):
t.sub_(m).div_(s)
return np.expand_dims(orig_imgs, 0), torch_imgs.unsqueeze(0)
t1_idx = t0_idx + 1
dither_t0 = load_dither(fnames[t0_idx])
rgb_t1 = load_rgb(fnames[t1_idx])
yield fnames[t1_idx], rgb_t1, dither_t0
def just_fname(full_name):
m = re.match(".*/(f_.*jpg)", full_name)
if not m:
raise Exception("unexpected filename [%s]" % full_name)
return m.group(1)
dataset = (tf.data.Dataset.from_generator(fname_rgb_t1_dither_t0,
output_types=(tf.string, tf.float32,
tf.float32))
.batch(opts.batch_size)
.prefetch(1))
for fnames, rgb_imgs_t1, dither_t0 in dataset:
fnames = fnames.numpy()
rgb_imgs_t1 = rgb_imgs_t1.numpy()
dither_t0 = dither_t0.numpy()
pred_dithers = generator(rgb_imgs_t1, dither_t0)
for dither, full_fname in zip(pred_dithers, fnames):
dest_fname = just_fname(full_fname.decode()).replace("jpg", "png")
dither_pil = u.dither_to_pil(dither)
dither_pil = u.center_crop(dither_pil, 1448, 1072)
dither_pil.save("%s/%s" % (opts.output_dir, dest_fname))
