def __init__(self,
training_image_path,
paper_affine_generator=False,
transform=None,
cache_images=False,
use_cuda=False):
'''
:param training_image_path:
:param output_size:
:param transform:
:param cache_images: 如果数据量不是特别大可以缓存到内存里面加快读取速度
:param use_cuda:
'''
self.use_cuda = use_cuda
self.cache_images = cache_images
self.paper_affine_generator = paper_affine_generator
# read image file
self.training_image_path = training_image_path
self.train_data = sorted(os.listdir(self.training_image_path))
self.image_count = len(self.train_data)
# bi = np.floor(np.arange(n) / batch_size).astype(np.int) # batch index
# nb = bi[-1] + 1 # number of batches
self.imgs = [None] * self.image_count
self.transform = transform
self.control_size = 1000000000
# cache images into memory for faster training(~5GB)
if self.cache_images:
for i in tqdm(range(min(self.image_count, self.control_size)),
desc='Reading images'): # 最多10k张图片
image_name = self.train_data[i]
image_path = os.path.join(self.training_image_path, image_name)
image = cv2.imread(image_path) # shape [h,w,c] BGR
assert image is not None, 'Image Not Found' + image_path
self.imgs[i] = image
print('center:', "120,120")
# center = (256,256)
center = (120, 120)
# print('center:',"0,0")
# center = (0,0)
small = True
if small:
# 注意，这里x和y方向的偏移像素的话要考虑原图的大小，比如我的原图是240*240，ntg原论文是512*512，像素偏移的太大对结果影响很大
self.theta = generate_affine_param(scale=1.1,
degree=10,
translate_x=-10,
translate_y=10,
center=center)
else:
self.theta = generate_affine_param(scale=1.25,
degree=30,
translate_x=-20,
translate_y=20,
center=center)
# 将opencv的参数转换为pytorch的参数
self.theta = torch.from_numpy(self.theta.astype(np.float32)).expand(
1, 2, 3)
self.theta = param2theta(self.theta, 512, 512,
use_cuda=self.use_cuda)[0]
def generate_cave_pair():
image_dir = '/Users/zale/project/datasets/complete_ms_data/oil_painting_ms/oil_painting_ms/'
image_name = "oil_painting_ms_31.png"
image_np = cv2.imread(image_dir + image_name)
height, width, channel = image_np.shape
small = False
center = (height // 2, width // 2)
if small:
theta = generate_affine_param(scale=1.1,
degree=10,
translate_x=-10,
translate_y=10,
center=center)
else:
theta = generate_affine_param(scale=1.25,
degree=30,
translate_x=-20,
translate_y=20,
center=center)
warped_np = cv2.warpAffine(image_np,
theta, (width, height),
flags=cv2.INTER_CUBIC)
cv2.imwrite(image_name, warped_np)
def generate_matlab_pair():
data_dir = '/mnt/4T/zlk/datasets/mulitspectral/Harvard'
array_data = read_matlab_data(data_dir)
small = True
if small:
theta = generate_affine_param(scale=1.1,
degree=10,
translate_x=-10,
translate_y=10)
else:
theta = generate_affine_param(scale=1.25,
degree=30,
translate_x=-20,
translate_y=20)
theta = torch.from_numpy(theta).float()
a, b = theta.shape
theta = theta.expand(array_data.shape[2], a, b)
theta = param2theta(theta, 240, 240, use_cuda=False)
save_matlab_pic(array_data, theta)
# img2 = io.imread('../datasets/row_data/multispectral/mul_1s_s.png')
img1 = io.imread(
'../datasets/row_data/multispectral/fake_and_real_tomatoes_ms_17.png')
img2 = io.imread(
'../datasets/row_data/multispectral/fake_and_real_tomatoes_ms_28.png')
# img1 = io.imread('../datasets/row_data/texs1.jpeg')
# img2 = io.imread('../datasets/row_data/test2.jpeg')
fn_grid_loss = GridLoss(use_cuda=False, grid_size=512)
center = (256, 256)
theta_GT = generate_affine_param(scale=1.1,
degree=10,
translate_x=-10,
translate_y=10,
center=center)
theta_GT = torch.from_numpy(theta_GT).unsqueeze(0).float()
# 第一个是target，第二个是source
p = estimate_affine_param(img1, img2, itermax=1000)
im2warped = affine_transform(img2, p)
print(p)
print(theta_GT)
p = torch.from_numpy(p).unsqueeze(0).float()
grid_loss = fn_grid_loss.compute_grid_loss(p, theta_GT)
print(grid_loss)
small = True
# img1 = io.imread('../datasets/row_data/multispectral/fake_and_real_tomatoes_ms_31.png') * 1.0
# img2 = io.imread('../datasets/row_data/multispectral/fake_and_real_tomatoes_ms_17.png') * 1.0
# img1 = io.imread('../datasets/row_data/multispectral/mul_1s_s.png')
# img2 = io.imread('../datasets/row_data/multispectral/mul_1t_s.png')
img1 = io.imread('../datasets/row_data/texs1.jpeg')
img2 = io.imread('../datasets/row_data/test2.jpeg')
center = (img1.shape[0]/2,img1.shape[1]/2)
center = (0,0)
if small:
theta = generate_affine_param(scale=1.1, degree=10, translate_x=-10, translate_y=10,center=center)
else:
theta = generate_affine_param(scale=1.25, degree=30, translate_x=-20, translate_y=20,center=center)
use_cuda = torch.cuda.is_available()
env = "ntg_pytorch"
vis = VisdomHelper(env)
p = use_torch_ntg(img1,img2)
# p = use_cv2_ntg(img1,img2)
# im2warped = affine_transform(img2,p)
im2warped = affine_transform(img1,p)