def main(): logger = SummaryWriter('/home/xiu/databag/deblur/pretrain/finetune/') # ---------------------------------------- # load kernels # ---------------------------------------- device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') PSF_grid = np.load('./data/ZEMAX-AC254-075-A-new.npz')['PSF'] PSF_grid = PSF_grid.astype(np.float32) kw, kh = PSF_grid.shape[:2] for w_ in range(kw): for h_ in range(kh): PSF_grid[w_, h_] = PSF_grid[w_, h_] / np.sum(PSF_grid[w_, h_], axis=(0, 1)) # ---------------------------------------- # build # ---------------------------------------- model = net(n_iter=8, h_nc=64, in_nc=4, out_nc=3, nc=[64, 128, 256, 512], nb=3, act_mode="R", downsample_mode='strideconv', upsample_mode="convtranspose") model.load_state_dict( torch.load('/home/xiu/databag/deblur/usrnet_ours_epoch10.pth'), strict=True) model.train() for _, v in model.named_parameters(): v.requires_grad = True model = model.to(device) train_imgs = glob.glob('/home/xiu/databag/deblur/images/*/**.png', recursive=True) train_imgs.sort() n_stage = 5 n_batch = 3 n_epoch = 200 w_patch = 128 n_patch = 2 ab_buffer = np.ones( (n_batch, PSF_grid.shape[0] * PSF_grid.shape[1], 2 * n_stage + 1, 3), dtype=np.float32) * 0.1 ab_param = torch.tensor(ab_buffer, device=device, requires_grad=True) params = [] params += [{"params": [ab_param], "lr": 1e-4}] for key, value in model.named_parameters(): params += [{"params": [value], "lr": 1e-4}] optimizer = torch.optim.Adam(params, lr=1e-4) img_index = np.arange(len(train_imgs)) global_iter = 0 expand = PSF_grid.shape[2] // 2 for epoch in range(n_epoch): np.random.shuffle(img_index) for iteration in range(len(train_imgs) // n_batch): imgs = [] for ii in range(n_batch): imgs.append( cv2.imread(train_imgs[img_index[iteration * n_batch + ii]])) global_iter += 1 x = [] y = [] k = [] vis_L = [] vis_H = [] vis_E = [] for img in imgs: w, h = img.shape[:2] x_start = np.random.randint( 0, w - w_patch * n_patch - expand * 2 + 1) y_start = np.random.randint( 0, h - w_patch * n_patch - expand * 2 + 1) px = np.random.randint(0, PSF_grid.shape[0] - n_patch) py = np.random.randint(0, PSF_grid.shape[1] - n_patch) patch_H = img[x_start:x_start + w_patch * n_patch + expand * 2, y_start:y_start + w_patch * n_patch + expand * 2] patch_L = util_deblur.blockConv2d( patch_H, PSF_grid[px:px + n_patch, py:py + n_patch], expand) vis_L.append(patch_L) vis_H.append(patch_H[expand:-expand, expand:-expand]) patch_L_wrap = util_deblur.wrap_boundary_liu( patch_L, (w_patch * n_patch + expand * 2, w_patch * n_patch + expand * 2)) patch_L_wrap = np.hstack( (patch_L_wrap[:, -expand:, :], patch_L_wrap[:, :w_patch * n_patch + expand, :])) patch_L_wrap = np.vstack( (patch_L_wrap[-expand:, :, :], patch_L_wrap[:w_patch * n_patch + expand, :, :])) x_L = util.uint2single(patch_L_wrap) x_blocky = [] for h_ in range(n_patch): for w_ in range(n_patch): x_blocky.append(x_L[w_*w_patch:w_*w_patch+w_patch+expand*2,\ h_*w_patch:h_*w_patch+w_patch+expand*2:]) x_blocky = [util.single2tensor4(el) for el in x_blocky] x_blocky = torch.cat(x_blocky, dim=0) k_all = [] for w_ in range(n_patch): for h_ in range(n_patch): k_all.append( util.single2tensor4(PSF_grid[h_ + px, w_ + py])) k_all = torch.cat(k_all, dim=0) x_gt = util.uint2single(patch_H[expand:-expand, expand:-expand]) x_gt = util.single2tensor4(x_gt) y.append(x_blocky) x.append(x_gt) k.append(k_all) ab = F.softplus(ab_param) loss = 0 for i in range(n_batch): yy = y[i].to(device) kk = k[i].to(device) xx = x[i].to(device) xE = model.forward_patchdeconv(yy, kk, ab[i], [n_patch, n_patch], w_patch) loss += F.l1_loss(xE[-2], xx) vis_E.append(util.tensor2uint(xE[-2])) optimizer.zero_grad() loss.backward() optimizer.step() if global_iter % 10 == 0: print('iter {}: loss{}.'.format(global_iter, loss.item())) logger.add_scalar('finetune-loss', loss.item(), global_iter) for i in range(n_batch): show1 = np.hstack((vis_H[i], vis_L[i], vis_E[i])) logger.add_image('show-{}'.format(i), util.uint2tensor3(show1[:, :, ::-1])) logger.flush()
def main(): logger = SummaryWriter('/home/xiu/databag/deblur/pretrain/full/') # ---------------------------------------- # load kernels # ---------------------------------------- device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') all_PSFs = load_kernels('./data') # ---------------------------------------- # build # ---------------------------------------- model = net(n_iter=8, h_nc=64, in_nc=4, out_nc=3, nc=[64, 128, 256, 512], nb=3, act_mode="R", downsample_mode='strideconv', upsample_mode="convtranspose") model.train() for _, v in model.named_parameters(): v.requires_grad = True model = model.to(device) train_imgs = glob.glob('/home/xiu/databag/deblur/images/*/**.png', recursive=True) train_imgs.sort() n_stage = 5 n_batch = 3 n_epoch = 200 w_patch = 128 n_patch = 2 ab_buffer = np.ones((n_batch, n_patch * n_patch, 2 * n_stage + 1, 3), dtype=np.float32) * 0.1 ab_param = torch.tensor(ab_buffer, device=device, requires_grad=True) params = [] params += [{"params": [ab_param], "lr": 1e-4}] for key, value in model.named_parameters(): params += [{"params": [value], "lr": 1e-4}] optimizer = torch.optim.Adam(params, lr=1e-4) img_index = np.arange(len(train_imgs)) global_iter = 0 PSFs = [] for i in range(n_batch): if i % 2 == 0: PSFs.append(get_kernels(all_PSFs, n_patch)) else: PSFs.append(rand_kernels(n_patch)) expands = [] for i in range(n_batch): expands.append(PSFs[i].shape[2] // 2) for epoch in range(n_epoch): np.random.shuffle(img_index) for iteration in range(len(train_imgs) // n_batch): imgs = [] for ii in range(n_batch): imgs.append( cv2.imread(train_imgs[img_index[iteration * n_batch + ii]])) global_iter += 1 if global_iter % 100 == 0: PSFs = [] for i in range(n_batch): if i % 2 == 0: PSFs.append(get_kernels(all_PSFs, n_patch)) else: PSFs.append(rand_kernels(n_patch)) expands = [] for i in range(n_batch): expands.append(PSFs[i].shape[2] // 2) x = [] y = [] k = [] vis_L = [] vis_H = [] vis_E = [] for img, expand, PSF in zip(imgs, expands, PSFs): w, h = img.shape[:2] x_start = np.random.randint( 0, w - w_patch * n_patch - expand * 2 + 1) y_start = np.random.randint( 0, h - w_patch * n_patch - expand * 2 + 1) patch_H = img[x_start:x_start + w_patch * n_patch + expand * 2, y_start:y_start + w_patch * n_patch + expand * 2] patch_L = util_deblur.blockConv2d(patch_H, PSF, expand) vis_L.append(patch_L) vis_H.append(patch_H[expand:-expand, expand:-expand]) patch_L_wrap = util_deblur.wrap_boundary_liu( patch_L, (w_patch * n_patch + expand * 2, w_patch * n_patch + expand * 2)) patch_L_wrap = np.hstack( (patch_L_wrap[:, -expand:, :], patch_L_wrap[:, :w_patch * n_patch + expand, :])) patch_L_wrap = np.vstack( (patch_L_wrap[-expand:, :, :], patch_L_wrap[:w_patch * n_patch + expand, :, :])) x_L = util.uint2single(patch_L_wrap) x_blocky = [] for h_ in range(n_patch): for w_ in range(n_patch): x_blocky.append(x_L[w_*w_patch:w_*w_patch+w_patch+expand*2,\ h_*w_patch:h_*w_patch+w_patch+expand*2:]) x_blocky = [util.single2tensor4(el) for el in x_blocky] x_blocky = torch.cat(x_blocky, dim=0) k_all = [] for w_ in range(n_patch): for h_ in range(n_patch): k_all.append(util.single2tensor4(PSF[h_, w_])) k_all = torch.cat(k_all, dim=0) x_gt = util.uint2single(patch_H[expand:-expand, expand:-expand]) x_gt = util.single2tensor4(x_gt) y.append(x_blocky) x.append(x_gt) k.append(k_all) ab = F.softplus(ab_param) loss = 0 for i in range(n_batch): yy = y[i].to(device) kk = k[i].to(device) xx = x[i].to(device) xE = model.forward_patchdeconv(yy, kk, ab[i], [n_patch, n_patch], w_patch) loss += F.l1_loss(xE[-2], xx) vis_E.append(util.tensor2uint(xE[-2])) optimizer.zero_grad() loss.backward() optimizer.step() if global_iter % 10 == 0: print('iter {}: loss{}.'.format(global_iter, loss.item())) logger.add_scalar('train-loss', loss.item(), global_iter) for i in range(n_batch): show1 = np.hstack((vis_H[i], vis_L[i], vis_E[i])) logger.add_image('show-{}'.format(i), util.uint2tensor3(show1[:, :, ::-1])) logger.flush() ab_numpy = ab.detach().cpu().numpy()[:, :, 0, 0] ab_numpy = ab_numpy.flatten() torch.save(model.state_dict(), 'usrnet_ours_epoch{}.pth'.format(epoch)) np.savetxt('ab_ours.txt', ab_numpy)
def main(): # ---------------------------------------- # load kernels # ---------------------------------------- #PSF_grid = np.load('./data/Schuler_PSF01.npz')['PSF'] #PSF_grid = np.load('./data/Schuler_PSF_bench.npz')['PSF'] PSF_grid = np.load('./data/Heidel_PSF_plano.npz')['PSF'] #PSF_grid = np.load('./data/ZEMAX-LA1608.npz')['PSF'] #PSF_grid = np.load('./data/Schuler_PSF03.npz')['PSF'] #PSF_grid = np.load('./data/PSF.npz')['PSF'] PSF_grid = PSF_grid.astype(np.float32) gx, gy = PSF_grid.shape[:2] for xx in range(gx): for yy in range(gy): PSF_grid[xx, yy] = PSF_grid[xx, yy] / np.sum(PSF_grid[xx, yy], axis=(0, 1)) # ---------------------------------------- # load model # ---------------------------------------- device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = net(n_iter=8, h_nc=64, in_nc=4, out_nc=3, nc=[64, 128, 256, 512], nb=3, act_mode="R", downsample_mode='strideconv', upsample_mode="convtranspose") #model.train() #model.load_state_dict(torch.load('./data/usrgan.pth'), strict=True) #model.load_state_dict(torch.load('./data/usrnet.pth'), strict=True) model.load_state_dict( torch.load('/home/xiu/databag/deblur/usrnet_ours_epoch10.pth'), strict=True) for _, v in model.named_parameters(): v.requires_grad = True model = model.to(device) imgs = glob.glob('/home/xiu/databag/deblur/images/*/**.png', recursive=True) imgs.sort() patch_size = 3 * 128 num_patch = 3 expand = PSF_grid.shape[2] // 2 #positional alpha-beta parameters for HQS stage = 5 ab_buffer = np.ones((gx, gy, 2 * stage + 1, 3), dtype=np.float32) * 0.1 ab_buffer[:, :, 0, :] = 0.01 ab = torch.tensor(ab_buffer, device=device, requires_grad=True) params = [] params += [{"params": [ab], "lr": 1e-3}] for key, value in model.named_parameters(): params += [{"params": [value], "lr": 1e-4}] optimizer = torch.optim.Adam(params, lr=1e-4) running = True while running: #alpha.beta img_idx = np.random.randint(len(imgs)) img = imgs[img_idx] img_H = cv2.imread(img) w, h = img_H.shape[:2] px_start = np.random.randint(0, gx - num_patch + 1) py_start = np.random.randint(0, gy - num_patch + 1) x_start = np.random.randint(0, w - patch_size - expand * 2 + 1) y_start = np.random.randint(0, h - patch_size - expand * 2 + 1) PSF_patch = PSF_grid[px_start:px_start + num_patch, py_start:py_start + num_patch] patch_H = img_H[x_start:x_start + patch_size + expand * 2, y_start:y_start + patch_size + expand * 2] patch_L = util_deblur.blockConv2d(patch_H, PSF_patch, expand) block_size = patch_size // num_patch block_expand = max(patch_size // 16, expand) if block_expand > 0: patch_L_wrap = util_deblur.wrap_boundary_liu( patch_L, (patch_size + block_expand * 2, patch_size + block_expand * 2)) #centralize patch_L_wrap = np.hstack( (patch_L_wrap[:, -block_expand:, :], patch_L_wrap[:, :patch_size + block_expand, :])) patch_L_wrap = np.vstack( (patch_L_wrap[-block_expand:, :, :], patch_L_wrap[:patch_size + block_expand, :, :])) else: patch_L_wrap = patch_L if block_expand > 0: x = util.uint2single(patch_L_wrap) else: x = util.uint2single(patch_L) x_blocky = [] for h_ in range(num_patch): for w_ in range(num_patch): x_blocky.append(x[w_*block_size:w_*block_size+block_size+block_expand*2,\ h_*block_size:h_*block_size+block_size+block_expand*2:]) x_blocky = [util.single2tensor4(el) for el in x_blocky] x_blocky = torch.cat(x_blocky, dim=0) k_all = [] for w_ in range(num_patch): for h_ in range(num_patch): k_all.append(util.single2tensor4(PSF_patch[h_, w_])) k = torch.cat(k_all, dim=0) x_gt = util.uint2single(patch_H[expand:-expand, expand:-expand]) x_gt = util.single2tensor4(x_gt) [x_blocky, x_gt, k] = [el.to(device) for el in [x_blocky, x_gt, k]] #cd = F.softplus(ab[px_start:px_start+num_patch,py_start:py_start+num_patch].reshape(num_patch**2,2*stage+1,1,1)) #for n_iter in range(optim_iter): cd = F.softplus(ab[px_start:px_start + num_patch, py_start:py_start + num_patch]) cd = cd.view(num_patch**2, 2 * stage + 1, 3) x_E = model.forward_patchdeconv(x_blocky, k, cd, [num_patch, num_patch], patch_sz=patch_size // num_patch) loss = 0 #for xx in x_E[::2]: loss = F.l1_loss(x_E[-2], x_gt) optimizer.zero_grad() loss.backward() optimizer.step() print('loss {}'.format(loss.item())) x_E = x_E[:-1] patch_L = patch_L_wrap.astype(np.uint8) patch_E = util.tensor2uint(x_E[-1]) patch_E_all = [util.tensor2uint(pp) for pp in x_E] patch_E_z = np.hstack((patch_E_all[::2])) patch_E_x = np.hstack((patch_E_all[1::2])) patch_E_show = np.vstack((patch_E_z, patch_E_x)) if block_expand > 0: show = np.hstack((patch_H[expand:-expand, expand:-expand], patch_L[block_expand:-block_expand, block_expand:-block_expand], patch_E)) else: show = np.hstack((patch_H[expand:-expand, expand:-expand], patch_L, patch_E)) #get kernel #cv2.imshow('stage',patch_E_show) #rgb = np.hstack((patch_E[:,:,0],patch_E[:,:,1],patch_E[:,:,2])) cv2.imshow('HL', show) #cv2.imshow('RGB',rgb) key = cv2.waitKey(1) if key == ord('q'): running = False break ab_numpy = ab.detach().cpu().numpy().flatten() #.reshape(-1,2*stage+1) torch.save(model.state_dict(), 'usrnet_plano_ours.pth') np.savetxt('ab_plano_ours.txt', ab_numpy)
def main(): # ---------------------------------------- # load kernels # ---------------------------------------- #PSF_grid = np.load('./data/Schuler_PSF01.npz')['PSF'] PSF_grid = np.load('./data/Schuler_PSF_facade.npz')['PSF'] #PSF_grid = np.load('./data/Schuler_PSF03.npz')['PSF'] #PSF_grid = np.load('./data/PSF.npz')['PSF'] #print(PSF_grid.shape) PSF_grid = PSF_grid.astype(np.float32) gx,gy = PSF_grid.shape[:2] for xx in range(gx): for yy in range(gy): PSF_grid[xx,yy] = PSF_grid[xx,yy]/np.sum(PSF_grid[xx,yy],axis=(0,1)) # ---------------------------------------- # load model # ---------------------------------------- device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = net(n_iter=8, h_nc=64, in_nc=4, out_nc=3, nc=[64, 128, 256, 512], nb=2, act_mode="R", downsample_mode='strideconv', upsample_mode="convtranspose") model.load_state_dict(torch.load('usrnet_facade_gan.pth'), strict=True) model.eval() #model.train() for _, v in model.named_parameters(): v.requires_grad = False # v.requires_grad = False model = model.to(device) #imgs = glob.glob('/home/xiu/databag/deblur/synthetic/synthetic_dataset/ground_truth/*.png') #imgs = glob.glob('/home/xiu/databag/coco/val2014/*.jpg') #imgs.sort() #img_L = cv2.imread('/home/xiu/workspace/dwgan/MPI_data/facade/blurry.jpg') img_H = cv2.imread('./suo/5.jpg') img_H = cv2.resize(img_H,dsize=None,fx=0.5,fy=0.5) w_crop = img_H.shape[0] - 8*128 h_crop = img_H.shape[1] - 12*128 to_w = w_crop//2 to_h = h_crop//2 img_H = img_H[to_w:-(w_crop-to_w),to_h:-(h_crop-to_h),:] #img_H = img_H[:8*128,:12*128] img_L = util_deblur.blockConv2d(img_H,PSF_grid) img_E = np.zeros_like(img_L) img_E_deconv = [] img_E_denoise = [] for i in range(5): img_E_deconv.append(np.zeros_like(img_L)) img_E_denoise.append(np.zeros_like(img_L)) weight_E = np.zeros_like(img_L) patch_size = 2*128 num_patch = 2 p_size = patch_size//num_patch expand = PSF_grid.shape[2]//2 #positional alpha-beta parameters for HQS #ab_numpy = np.ones((num_patch*num_patch,17,1,1),dtype=np.float32)*0.1 #ab_numpy[:,0,:,:] = 0.01 ab_numpy = np.loadtxt('ab_facade_gan.txt').astype(np.float32) ab_numpy = ab_numpy[...,None,None] ab = torch.tensor(ab_numpy,device=device,requires_grad=False) #save img_L #while running: for px_start in range(0,8-2+1,2): for py_start in range(0,12-2+1,2): #px_start = np.random.randint(0,PSF_grid.shape[0]+1-num_patch) #py_start = np.random.randint(0,PSF_grid.shape[1]+1-num_patch) PSF_patch = PSF_grid[px_start:px_start+num_patch,py_start:py_start+num_patch] # x = util.uint2single(patch_L) block_size = patch_size//num_patch patch_L = img_L[px_start*p_size:(px_start+num_patch)*p_size,py_start*p_size:py_start*p_size+num_patch*p_size,:] #patch_H = img_H[px_start*p_size:(px_start+num_patch)*p_size,py_start*p_size:py_start*p_size+num_patch*p_size,:] #block_expand = expand*2 block_expand = expand #block_expand = 1 if block_expand > 0: patch_L_wrap = util_deblur.wrap_boundary_liu(patch_L,(patch_size+block_expand*2,patch_size+block_expand*2)) #centralize patch_L_wrap = np.hstack((patch_L_wrap[:,-block_expand:,:],patch_L_wrap[:,:patch_size+block_expand,:])) patch_L_wrap = np.vstack((patch_L_wrap[-block_expand:,:,:],patch_L_wrap[:patch_size+block_expand,:,:])) else: patch_L_wrap = patch_L if block_expand>0: x = util.uint2single(patch_L_wrap) else: x = util.uint2single(patch_L) x_blocky = [] for h_ in range(num_patch): for w_ in range(num_patch): x_blocky.append(x[w_*block_size:w_*block_size+block_size+block_expand*2,\ h_*block_size:h_*block_size+block_size+block_expand*2:]) x_blocky = [util.single2tensor4(el) for el in x_blocky] x_blocky = torch.cat(x_blocky,dim=0) # x = util.single2tensor4(x) # x_blocky = torch.cat(torch.chunk(x,num_patch,dim=2),dim=0) # x_blocky = torch.cat(torch.chunk(x_blocky,num_patch,dim=3),dim=0) k_all = [] for w_ in range(num_patch): for h_ in range(num_patch): k_all.append(util.single2tensor4(PSF_patch[h_,w_])) k = torch.cat(k_all,dim=0) [x_blocky,k] = [el.to(device) for el in [x_blocky,k]] cd = F.softplus(ab) beta = cd[:,9:,0,0].detach().cpu().numpy() x_E = model.forward_patchdeconv(x_blocky,k,cd,[num_patch,num_patch],patch_sz=patch_size//num_patch) x_E = x_E[:-1] patch_L = patch_L_wrap.astype(np.uint8) patch_E = util.tensor2uint(x_E[-1]) patch_E_all = [util.tensor2uint(pp) for pp in x_E] #patch_E_z = np.hstack((patch_E_all[::2])) #patch_E_x = np.hstack((patch_E_all[1::2])) #patch_E_show = np.vstack((patch_E_z,patch_E_x)) #if block_expand>0: # show = np.hstack((patch_L[block_expand:-block_expand,block_expand:-block_expand],patch_E)) #else: # show = np.hstack((patch_L,patch_E)) #get kernel for i in range(5): img_E_deconv[i][px_start*p_size:(px_start+num_patch)*p_size,py_start*p_size:py_start*p_size+num_patch*p_size,:] += patch_E_all[2*i] img_E_denoise[i][px_start*p_size:(px_start+num_patch)*p_size,py_start*p_size:py_start*p_size+num_patch*p_size,:] += patch_E_all[2*i+1] weight_E[px_start*p_size:(px_start+num_patch)*p_size,py_start*p_size:py_start*p_size+num_patch*p_size,:] += 1.0 #cv2.imshow('stage',patch_E_show) #cv2.imshow('HL',show) #cv2.imshow('RGB',rgb) #key = cv2.waitKey(-1) #if key==ord('n'): # break img_E = img_E/weight_E img_E_deconv = [pp/weight_E for pp in img_E_deconv] img_E_denoise = [pp/weight_E for pp in img_E_denoise] # img_L = img_L.astype(np.uint8) # img_E = img_E.astype(np.uint8) # img_E_deconv = img_E_deconv.astype(np.uint8) # img_E_denoise = img_E_denoise.astype(np.uint8) # cv2.imshow('imE',img_E) # cv2.imshow('imE_deconv',img_E_deconv) # cv2.imshow('imE_denoise',img_E_denoise) # cv2.imshow('imL',img_L) for i in range(5): zk = img_E_deconv[i] xk = img_E_denoise[i] zk = zk.astype(np.uint8) xk = xk.astype(np.uint8) cv2.imwrite('/home/xiu/workspace/dwgan/suo/z5-{}.png'.format(i),zk) cv2.imwrite('/home/xiu/workspace/dwgan/suo/x5-{}.png'.format(i),xk) print(i)