def __init__(self,
csv_file,
training_image_path,
output_size=(480, 640),
geometric_model='affine',
transform=None,
random_sample=False,
random_t=0.5,
random_s=0.5,
random_alpha=1 / 6,
random_t_tps=0.4):
self.random_sample = random_sample
self.random_t = random_t
self.random_t_tps = random_t_tps
self.random_alpha = random_alpha
self.random_s = random_s
self.out_h, self.out_w = output_size
# read csv file
self.train_data = pd.read_csv(csv_file)
self.img_names = self.train_data.iloc[:, 0]
self.theta_array = self.train_data.iloc[:, 1:].values.astype('float')
# copy arguments
self.training_image_path = training_image_path
self.transform = transform
self.geometric_model = geometric_model
self.affineTnf = GeometricTnf(out_h=self.out_h,
out_w=self.out_w,
use_cuda=False)
def __init__(self,
csv_file=config.PF_WILLOW_TEST_DATA,
dataset_path=config.PF_WILLOW_DIR,
output_size=(240, 240),
transform=None,
category=None,
pck_procedure='scnet'):
self.category_names = [
'car(G)', 'car(M)', 'car(S)', 'duck(S)', 'motorbike(G)',
'motorbike(M)', 'motorbike(S)', 'winebottle(M)', 'winebottle(wC)',
'winebottle(woC)'
]
self.out_h, self.out_w = output_size
self.pairs = pd.read_csv(csv_file)
self.category = self.pairs.iloc[:, 2].as_matrix().astype('int')
if category is not None:
cat_idx = np.nonzero(self.category == category)[0]
self.category = self.category[cat_idx]
self.pairs = self.pairs.iloc[cat_idx, :]
self.img_A_names = self.pairs.iloc[:, 0]
self.img_B_names = self.pairs.iloc[:, 1]
self.point_A_coords = self.pairs.iloc[:, 3:5]
self.point_B_coords = self.pairs.iloc[:, 5:]
self.dataset_path = dataset_path
self.transform = transform
# no cuda as dataset is called from CPU threads in dataloader and produces confilct
self.affineTnf = GeometricTnf(out_h=self.out_h,
out_w=self.out_w,
use_cuda=False)
self.pck_procedure = pck_procedure
def __init__(self,
csv_file=config.PF_PASCAL_TRAIN_DATA,
training_image_path=config.PF_PASCAL_DIR,
dataset_size=None,
output_size=(240, 240),
transform=None,
random_crop=False):
self.random_crop = random_crop
self.out_h, self.out_w = output_size
self.train_data = pd.read_csv(csv_file)
if dataset_size is not None:
dataset_size = min((dataset_size, len(self.train_data)))
self.train_data = self.train_data.iloc[0:dataset_size, :]
self.img_A_names = self.train_data.iloc[:, 0]
self.img_B_names = self.train_data.iloc[:, 1]
self.set = self.train_data.iloc[:, 2].as_matrix()
self.flip = self.train_data.iloc[:, 3].as_matrix().astype('int')
self.training_image_path = training_image_path
self.transform = transform
# no cuda as dataset is called from CPU threads in dataloader and produces confilct
self.affineTnf = GeometricTnf(out_h=self.out_h,
out_w=self.out_w,
use_cuda=False)
def test(model,loss_fn,dataloader,pair_generation_tnf,use_cuda=True,geometric_model='affine'):
model.eval()
test_loss = 0
dice = 0
for batch_idx, batch in enumerate(dataloader):
tnf_batch = pair_generation_tnf(batch)
theta = model(tnf_batch)
loss = loss_fn(theta,tnf_batch['theta_GT'],tnf_batch)
test_loss += loss.data.cpu().numpy()
I = tnf_batch['target_mask']
geometricTnf = GeometricTnf(geometric_model, 240, 240, use_cuda = use_cuda)
if geometric_model == 'affine':
theta = theta.view(-1,2,3)
J = geometricTnf(tnf_batch['source_mask'],theta)
if use_cuda:
I = I.cuda()
J = J.cuda()
numerator = 2 * torch.sum(torch.sum(torch.sum(I * J,dim=3),dim=2),dim=1)
denominator = torch.sum(torch.sum(torch.sum(I + J,dim=3),dim=2),dim=1)
dice = dice + torch.sum(numerator/(denominator + 0.00001))/I.shape[0]
test_loss /= len(dataloader)
dice /=len(dataloader)
print('Test set: Average loss: {:.6f}'.format(test_loss))
print('Test set: Dice: {:.6f}'.format(dice))
return test_loss
def __init__(self, geometric_model='affine', tps_grid_size=3, tps_reg_factor=0, h_matches=15, w_matches=15, use_conv_filter=False, dilation_filter=None, use_cuda=True, normalize_inlier_count=False, offset_factor=227/210):
super(WeakInlierCount, self).__init__()
self.normalize=normalize_inlier_count
self.geometric_model = geometric_model
self.geometricTnf = GeometricTnf(geometric_model='affine',
tps_grid_size=tps_grid_size,
tps_reg_factor=tps_reg_factor,
out_h=h_matches, out_w=w_matches,
offset_factor = offset_factor,
use_cuda=use_cuda)
# define dilation filter
if dilation_filter is None:
dilation_filter = generate_binary_structure(2, 2)
# define identity mask tensor (w,h are switched and will be permuted back later)
mask_id = np.zeros((w_matches,h_matches,w_matches*h_matches))
idx_list = list(range(0, mask_id.size, mask_id.shape[2]+1))
mask_id.reshape((-1))[idx_list]=1
mask_id = mask_id.swapaxes(0,1)
# perform 2D dilation to each channel
if not use_conv_filter:
if not (isinstance(dilation_filter,int) and dilation_filter==0):
for i in range(mask_id.shape[2]):
mask_id[:,:,i] = binary_dilation(mask_id[:,:,i],structure=dilation_filter).astype(mask_id.dtype)
else:
for i in range(mask_id.shape[2]):
flt=np.array([[1/16,1/8,1/16],
[1/8, 1/4, 1/8],
[1/16,1/8,1/16]])
mask_id[:,:,i] = scipy.signal.convolve2d(mask_id[:,:,i], flt, mode='same', boundary='fill', fillvalue=0)
# convert to PyTorch variable
mask_id = Variable(torch.FloatTensor(mask_id).transpose(1,2).transpose(0,1).unsqueeze(0),requires_grad=False)
self.mask_id = mask_id
if use_cuda:
self.mask_id = self.mask_id.cuda();
def __init__(self,
csv_file,
dataset_path,
output_size=(240, 240),
transform=None,
category=None):
self.category_names = [
'aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car',
'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse', 'motorbike',
'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor'
]
self.out_h, self.out_w = output_size
self.pairs = pd.read_csv(csv_file)
self.category = self.pairs.iloc[:, 2].as_matrix().astype('float')
if category is not None:
cat_idx = np.nonzero(self.category == category)[0]
self.category = self.category[cat_idx]
self.pairs = self.pairs.iloc[cat_idx, :]
self.img_A_names = self.pairs.iloc[:, 0]
self.img_B_names = self.pairs.iloc[:, 1]
self.point_A_coords = self.pairs.iloc[:, 3:5]
self.point_B_coords = self.pairs.iloc[:, 5:]
self.dataset_path = dataset_path
self.transform = transform
# no cuda as dataset is called from CPU threads in dataloader and produces confilct
self.affineTnf = GeometricTnf(out_h=self.out_h,
out_w=self.out_w,
use_cuda=False)
def __init__(self, csv_file, dataset_path, output_size=(240, 240), transform=None):
self.category_names = [
'Faces', 'Faces_easy', 'Leopards', 'Motorbikes', 'accordion', 'airplanes', 'anchor',
'ant', 'barrel', 'bass', 'beaver', 'binocular', 'bonsai', 'brain', 'brontosaurus',
'buddha', 'butterfly', 'camera', 'cannon', 'car_side', 'ceiling_fan', 'cellphone',
'chair', 'chandelier', 'cougar_body', 'cougar_face', 'crab', 'crayfish', 'crocodile',
'crocodile_head', 'cup', 'dalmatian', 'dollar_bill', 'dolphin', 'dragonfly',
'electric_guitar', 'elephant', 'emu', 'euphonium', 'ewer', 'ferry', 'flamingo',
'flamingo_head', 'garfield', 'gerenuk', 'gramophone', 'grand_piano', 'hawksbill',
'headphone', 'hedgehog', 'helicopter', 'ibis', 'inline_skate', 'joshua_tree',
'kangaroo', 'ketch', 'lamp', 'laptop', 'llama', 'lobster', 'lotus', 'mandolin',
'mayfly', 'menorah', 'metronome', 'minaret', 'nautilus', 'octopus', 'okapi', 'pagoda',
'panda', 'pigeon', 'pizza', 'platypus', 'pyramid', 'revolver', 'rhino', 'rooster',
'saxophone', 'schooner', 'scissors', 'scorpion', 'sea_horse', 'snoopy', 'soccer_ball',
'stapler', 'starfish', 'stegosaurus', 'stop_sign', 'strawberry', 'sunflower', 'tick',
'trilobite', 'umbrella', 'watch', 'water_lilly', 'wheelchair', 'wild_cat',
'windsor_chair', 'wrench', 'yin_yang'
]
self.out_h, self.out_w = output_size
self.pairs = pd.read_csv(csv_file)
self.img_A_names = self.pairs.iloc[:, 0]
self.img_B_names = self.pairs.iloc[:, 1]
self.category = self.pairs.iloc[:, 2].as_matrix().astype('float')
self.annot_A_str = self.pairs.iloc[:, 3:5]
self.annot_B_str = self.pairs.iloc[:, 5:]
self.dataset_path = dataset_path
self.transform = transform
# no cuda as dataset is called from CPU threads in dataloader and produces confilct
self.affineTnf = GeometricTnf(out_h=self.out_h, out_w=self.out_w, use_cuda=False)
def __init__(self, params, is_train=True, frame_gap=1, augment=['crop', 'flip', 'frame_gap']):
self.filelist = params['filelist']
self.batchSize = params['batchSize']
self.imgSize = params['imgSize']
self.cropSize = params['cropSize']
self.cropSize2 = params['cropSize2']
self.videoLen = params['videoLen']
# prediction distance, how many frames far away
self.predDistance = params['predDistance']
# offset x,y parameters
self.offset = params['offset']
# gridSize = 3
self.gridSize = params['gridSize']
self.is_train = is_train
self.frame_gap = frame_gap
self.augment = augment
f = open(self.filelist, 'r')
self.jpgfiles = []
self.fnums = []
for line in f:
rows = line.split()
jpgfile = rows[0]
fnum = int(rows[1])
self.jpgfiles.append(jpgfile)
self.fnums.append(fnum)
f.close()
self.geometricTnf = GeometricTnf('affine', out_h=params['cropSize2'], out_w=params['cropSize2'], use_cuda = False)
def __init__(self,
fr_feature_size=15,
fr_kernel_sizes=[7, 5],
fr_channels=[128, 64],
feature_extraction_cnn='vgg',
feature_extraction_last_layer='',
return_correlation=False,
normalize_features=True,
normalize_matches=True,
batch_normalization=True,
train_fe=False,
use_cuda=True,
s1_output_dim=6,
s2_output_dim=18):
super(TwoStageCNNGeometric, self).__init__(
output_dim=s1_output_dim,
fr_feature_size=fr_feature_size,
fr_kernel_sizes=fr_kernel_sizes,
fr_channels=fr_channels,
feature_extraction_cnn=feature_extraction_cnn,
feature_extraction_last_layer=feature_extraction_last_layer,
return_correlation=return_correlation,
normalize_features=normalize_features,
normalize_matches=normalize_matches,
batch_normalization=batch_normalization,
train_fe=train_fe,
use_cuda=use_cuda)
if s1_output_dim == 6:
self.geoTnf = GeometricTnf(geometric_model='affine',
use_cuda=use_cuda)
else:
tps_grid_size = np.sqrt(s2_output_dim / 2)
self.geoTnf = GeometricTnf(geometric_model='tps',
tps_grid_size=tps_grid_size,
use_cuda=use_cuda)
self.FeatureRegression2 = PPOFeatureRegression(
output_dim=s2_output_dim,
use_cuda=use_cuda,
feature_size=fr_feature_size,
kernel_sizes=fr_kernel_sizes,
channels=fr_channels,
batch_normalization=batch_normalization)
self.critic_linear = nn.Linear(self.output_size, 1)
def __init__(self, params, is_train=True, frame_gap=1, augment=['crop', 'flip', 'frame_gap'], task="actions_present", mode="train"):
self.filelist = params['filelist']
self.batchSize = params['batchSize']
self.imgSize = params['imgSize']
self.cropSize = params['cropSize']
self.cropSize2 = params['cropSize2']
self.videoLen = params['videoLen']
# prediction distance, how many frames far away
self.predDistance = params['predDistance']
# offset x,y parameters
self.offset = params['offset']
# gridSize = 3
self.gridSize = params['gridSize']
self.is_train = is_train
self.frame_gap = frame_gap
self.augment = augment
"""
f = open(self.filelist, 'r')
self.jpgfiles = []
self.fnums = []
for line in f:
rows = line.split()
jpgfile = rows[0]
fnum = int(rows[1])
self.jpgfiles.append(jpgfile)
self.fnums.append(fnum)
f.close()
"""
self.geometricTnf = GeometricTnf('affine', out_h=params['cropSize2'], out_w=params['cropSize2'], use_cuda = False)
# splits
if mode == 'train':
path_split = "lists/{}/train_subsetT.txt"
elif mode == 'val':
path_split = "lists/{}/train_subsetV.txt"
elif mode == 'test':
path_split = "lists/{}/val.txt"
else: raise ValueError('wrong mode')
path_split = path_split.format(task)
self.jpgfiles = []
self.fnums = []
black_list = [75, 76, 385, 4798, 4803, 4820, 6531, 6532, 6536]
with open(os.path.join(self.filelist, path_split),"r") as f:
for s in f.readlines():
path = s.split(' ')[0].split('.')[0]
if int(path.split('_')[-1]) in black_list:
print("Skipping video {}".format(path))
continue
path = os.path.join(self.filelist, 'frame/videos', path)
self.jpgfiles.append(path)
self.fnums.append(300)
def __init__(self,
fr_feature_size_two=15,
fr_kernel_sizes_two=[7,5],
fr_channels_two=[128,64],
fr_feature_size_four=15,
fr_kernel_sizes_four=[3,3,3],
fr_channels_four=[10,10,1],
feature_extraction_cnn='vgg',
feature_extraction_last_layer='',
return_correlation=False,
normalize_features=True,
normalize_matches=True,
batch_normalization=True,
train_fe=False,
use_cuda=True,
s1_output_dim=6,
s2_output_dim=256):
super(TwoStageCNNGeometric, self).__init__(output_dim=s1_output_dim,
fr_feature_size_two=fr_feature_size_two,
fr_kernel_sizes_two=fr_kernel_sizes_two,
fr_channels_two=fr_channels_two,
feature_extraction_cnn=feature_extraction_cnn,
feature_extraction_last_layer=feature_extraction_last_layer,
return_correlation=return_correlation,
normalize_features=normalize_features,
normalize_matches=normalize_matches,
batch_normalization=batch_normalization,
train_fe=train_fe,
use_cuda=use_cuda)
if s1_output_dim == 6:
self.geoTnf = GeometricTnf(geometric_model='affine', use_cuda=use_cuda)
else:
self.geoTnf = GeometricTnf(geometric_model='def', use_cuda=use_cuda)
if s2_output_dim < 256:
self.FeatureCorrelation2 = FeatureCorrelation(shape='3D',normalization=True)
self.FeatureRegression2 = FeatureRegression3D(s2_output_dim,use_cuda=self.use_cuda,
feature_size=fr_feature_size_two,kernel_sizes=fr_kernel_sizes_two,
channels=fr_channels_two)
else:
self.FeatureCorrelation2 = FeatureCorrelation(shape='4D',normalization=False)
self.FeatureRegression2 = FeatureRegression4D(s2_output_dim,use_cuda=use_cuda,
feature_size=fr_feature_size_four,kernel_sizes=fr_kernel_sizes_four,
channels=fr_channels_four)
self.GridDeformation2 = GridDeformation(s2_output_dim,use_cuda=self.use_cuda)
def eval_model_multistage(model,geometric_model,num_of_iters,source_image,target_image):
if geometric_model == 'affine_simple' or geometric_model == 'affine_simple_4':
geoTnf = GeometricTnf(geometric_model='affine', use_cuda=torch.cuda.is_available())
else:
geoTnf = GeometricTnf(geometric_model=geometric_model, use_cuda=torch.cuda.is_available())
for it in range(num_of_iters):
# First iteration
if it==0:
theta = model({'source_image':source_image,'target_image':target_image})
if geometric_model=='hom':
theta = homography_mat_from_4_pts(theta)
elif geometric_model == 'affine_simple' or geometric_model == 'affine_simple_4':
theta_simple = theta
theta = affine_mat_from_simple(theta_simple)
continue
# Subsequent iterations
# Compute warped image
warped_image = None
warped_image = geoTnf(source_image,theta)
# Re-estimate tranformation
theta_iter = model({'source_image':warped_image,'target_image':target_image})
# update accumultated transformation
if geometric_model=='hom':
theta = compose_H_matrices(theta,homography_mat_from_4_pts(theta_iter))
elif geometric_model == 'affine_simple' or geometric_model == 'affine_simple_4':
theta_simple = compose_simple(theta_simple, theta_iter)
theta = affine_mat_from_simple(theta_simple)
elif geometric_model=='affine':
theta = compose_aff_matrices(theta,theta_iter)
elif geometric_model=='tps':
theta = compose_tps(theta,theta_iter)
# warp one last time using final transformation
warped_image = None
warped_image = geoTnf(source_image,theta)
if geometric_model == 'affine_simple' or geometric_model == 'affine_simple_4':
theta = theta_simple
return (theta,warped_image)
def __init__(self,
out_h=240,
out_w=240,
use_cuda=True):
super(TaskLoss, self).__init__()
self.affTnf = GeometricTnf(geometric_model='affine',
out_h=out_h,
out_w=out_w,
use_cuda=use_cuda)
def __init__(self, csv_file, training_image_path, geometric_model='pose',output_size=(480, 640), transform=None):
self.out_h, self.out_w = output_size
# read csv file
self.train_data = pd.read_csv(csv_file)
self.img_src_names = self.train_data.iloc[:, 0]
self.img_target_names = self.train_data.iloc[:, 1]
self.theta_array = self.train_data.iloc[:, 2:].as_matrix().astype('float')
# copy arguments
self.training_image_path = training_image_path
self.transform = transform
self.geometric_model = geometric_model
self.affineTnf = GeometricTnf(out_h=self.out_h, out_w=self.out_w, use_cuda=False)
def __init__(self, csv_file, dataset_path, output_size=(240,240), transform=None):
self.out_h, self.out_w = output_size
self.pairs = pd.read_csv(csv_file)
self.img_A_names = self.pairs.iloc[:,0]
self.img_B_names = self.pairs.iloc[:,1]
self.point_A_coords = self.pairs.iloc[:, 2:22].as_matrix().astype('float')
self.point_B_coords = self.pairs.iloc[:, 22:].as_matrix().astype('float')
self.dataset_path = dataset_path
self.transform = transform
# no cuda as dataset is called from CPU threads in dataloader and produces confilct
self.affineTnf = GeometricTnf(out_h=self.out_h, out_w=self.out_w, use_cuda = False)
def __init__(self, csv_file, training_image_path, output_size=(240, 240),num_of_points=12, transform=None):
self.out_h, self.out_w = output_size
self.train_data = pd.read_csv(csv_file)
self.img_A_names = self.train_data.iloc[:, 0]
self.img_B_names = self.train_data.iloc[:, 1]
self.point_A_coords = self.train_data.iloc[:, 2:(num_of_points*2)+2].as_matrix().astype('float')
self.point_B_coords = self.train_data.iloc[:, (num_of_points*2)+2:].as_matrix().astype('float')
self.training_image_path = training_image_path
self.transform = transform
self.num_of_points = num_of_points
# no cuda as dataset is called from CPU threads in dataloader and produces confilct
self.affineTnf = GeometricTnf(out_h=self.out_h, out_w=self.out_w, use_cuda=False)
def __init__(self, csv_file, dataset_path,output_size=(240,240),transform=None):
self.out_h, self.out_w = output_size
self.pairs = pd.read_csv(csv_file)
self.img_A_names = self.pairs.iloc[:,0]
self.img_B_names = self.pairs.iloc[:,1]
self.flow_direction = self.pairs.iloc[:, 2].values.astype('int')
self.flip_img_A = self.pairs.iloc[:, 3].values.astype('int')
self.pair_category = self.pairs.iloc[:, 4].values.astype('int')
self.dataset_path = dataset_path
self.transform = transform
# no cuda as dataset is called from CPU threads in dataloader and produces confilct
self.affineTnf = GeometricTnf(out_h=self.out_h, out_w=self.out_w, use_cuda = False)
def __init__(self, csv_file, training_image_path, geometric_model='pose',use_timestamp=False,output_size=(480, 640), transform=None):
self.out_h, self.out_w = output_size
# read csv file
self.train_data = pd.read_csv(csv_file)
self.img_src_names = self.train_data.iloc[:, 0]
self.img_target_names = self.train_data.iloc[:, 1]
self.use_timestamp = use_timestamp
if self.use_timestamp:
self.timestamps = self.train_data.iloc[:,2]
# copy arguments
self.training_image_path = training_image_path
self.transform = transform
self.geometric_model = geometric_model
self.affineTnf = GeometricTnf(out_h=self.out_h, out_w=self.out_w, use_cuda=False)
def preprocess_image(image):
resizeCNN = GeometricTnf(out_h=240, out_w=240, use_cuda=False)
# convert to torch Variable
image = np.expand_dims(image.transpose((2, 0, 1)), 0)
image = torch.Tensor(image.astype(np.float32) / 255.0)
image_var = Variable(image, requires_grad=False)
# Resize image using bilinear sampling with identity affine tnf
image_var = resizeCNN(image_var)
# Normalize image
image_var = normalize_image(image_var)
return image_var
def __init__(self,
csv_file,
training_image_path,
output_size=(240, 240),
transform=None):
self.out_h, self.out_w = output_size
self.train_data = pd.read_csv(csv_file)
self.source_image_names = self.train_data.iloc[:, 0]
self.target_image_names = self.train_data.iloc[:, 1]
self.training_image_path = training_image_path
self.transform = transform
# no cuda as dataset is called from CPU threads in dataloader and produces confilct
self.affineTnf = GeometricTnf(out_h=self.out_h,
out_w=self.out_w,
use_cuda=False)
def read_input(path_imgA, path_imgB, output_size=(240, 240)):
out_h, out_w = output_size
resizeTnf = GeometricTnf(out_h=out_h, out_w=out_w, use_cuda=False)
image_A, im_size_A = readImage(img_path=path_imgA, affineTnf=resizeTnf)
image_B, im_size_B = readImage(img_path=path_imgB, affineTnf=resizeTnf)
sample = {
'source_image': image_A,
'target_image': image_B,
'source_im_size': im_size_A,
'target_im_size': im_size_B
}
sample = __transform(sample)
batch = __batchTensorToVars(sample)
return batch
def affTpsTnf(source_image, theta_aff, theta_aff_tps, use_cuda=use_cuda):
tpstnf = GeometricTnf(geometric_model='tps', use_cuda=use_cuda)
sampling_grid = tpstnf(image_batch=source_image,
theta_batch=theta_aff_tps,
return_sampling_grid=True)[1]
X = sampling_grid[:, :, :, 0].unsqueeze(3)
Y = sampling_grid[:, :, :, 1].unsqueeze(3)
Xp = X * theta_aff[:, 0].unsqueeze(1).unsqueeze(
2) + Y * theta_aff[:, 1].unsqueeze(1).unsqueeze(
2) + theta_aff[:, 2].unsqueeze(1).unsqueeze(2)
Yp = X * theta_aff[:, 3].unsqueeze(1).unsqueeze(
2) + Y * theta_aff[:, 4].unsqueeze(1).unsqueeze(
2) + theta_aff[:, 5].unsqueeze(1).unsqueeze(2)
sg = torch.cat((Xp, Yp), 3)
warped_image_batch = F.grid_sample(source_image, sg)
return warped_image_batch
def __init__(self,
csv_file,
training_image_path,
output_size=(240, 240),
transform=None):
self.out_h, self.out_w = output_size
self.train_data = pd.read_csv(csv_file)
self.img_A_names = self.train_data.iloc[:, 0]
self.img_B_names = self.train_data.iloc[:, 1]
self.point_A_coords = self.train_data.iloc[:, 2:22].as_matrix().astype(
'float')
self.point_B_coords = self.train_data.iloc[:, 22:].as_matrix().astype(
'float')
self.training_image_path = training_image_path
self.transform = transform
self.affineTnf = GeometricTnf(out_h=self.out_h, out_w=self.out_w)
def forward(self, theta, theta_GT, tnf_batch):
### compute square root of ssd
A = tnf_batch['target_image']
geometricTnf = GeometricTnf(self.geometric_model,
240,
240,
use_cuda=self.use_cuda)
B = geometricTnf(tnf_batch['source_image'], theta)
ssd = torch.sum(torch.sum(torch.sum(torch.pow(A - B, 2), dim=3),
dim=2),
dim=1)
ssd = torch.sum(ssd) / (A.shape[0] * A.shape[1] * A.shape[2] *
A.shape[3])
ssd = torch.sqrt(ssd)
return ssd
def __init__(self,
csv_file,
test_image_path,
output_size=(540, 540),
transform=None):
self.out_h, self.out_w = output_size
self.test_data = pd.read_csv(csv_file)
self.src_names = self.test_data.iloc[:, 0]
self.trg_names = self.test_data.iloc[:, 1]
self.src_point_coords = self.test_data.iloc[:, 2:42].values.astype(
'float')
self.theta_GT = self.test_data.iloc[:, 42:].values.astype('float')
self.test_image_path = test_image_path
self.transform = transform
# no cuda as dataset is called from CPU threads in dataloader and produces confilct
self.affineTnf = GeometricTnf(out_h=self.out_h,
out_w=self.out_w,
use_cuda=False)
def __init__(
self,
csv_file,
training_image_path,
output_size=(540, 540),
geometric_model='affine',
transform=None,
# def __init__(self, csv_file, training_image_path, output_size=(2160,2160), geometric_model='affine', transform=None,
random_sample=False,
random_t=0.5,
random_s=0.5,
random_alpha=1 / 6,
random_t_tps=0.4): # Original random_s=0.5, random_alpha = 1/6
self.random_sample = random_sample
self.random_t = random_t
self.random_t_tps = random_t_tps
self.random_alpha = random_alpha
self.random_s = random_s
self.out_h, self.out_w = output_size
# read csv file
self.train_data = pd.read_csv(csv_file)
self.img_names = self.train_data.iloc[:, 0]
self.img_names2 = self.train_data.iloc[:, 1]
self.theta_array = self.train_data.iloc[:,
2:].as_matrix().astype('float')
# copy arguments
self.training_image_path = training_image_path
self.transform = transform
self.geometric_model = geometric_model
self.affineTnf = GeometricTnf(out_h=self.out_h,
out_w=self.out_w,
use_cuda=False)
# Ready for distance
grid_size = 20
axis_coords = np.linspace(-1, 1, grid_size)
self.N = grid_size * grid_size
X, Y = np.meshgrid(axis_coords, axis_coords)
X = np.reshape(X, (1, self.N))
Y = np.reshape(Y, (1, self.N))
self.P = np.concatenate((X, Y))
# Check tilt
self.cosVal = [1, 0, -1, 0]
self.sinVal = [0, 1, 0, -1]
self.tilt = [1, 1 / np.cos(7 / 18 * np.pi)]
def process(self, source_img_path, target_img_path):
""" Main process of model
Input:
source_img_path: path to image you want to transform or np.array (RGB)
target_img_path: path to target image you want to transform or np.array (RGB)
Output:
affine_image (np.ndarray same shape with target) Affine Transformation result
affine_tps_image (np.ndarray same shape with target): Affine TPS transformation result
"""
# Load data
src_img, src_shape = self.load_image(source_img_path)
target_img, target_shape = self.load_image(target_img_path)
batch = {'source_image': src_img, 'target_image': target_img}
self.model.eval()
with torch.no_grad():
theta_aff, theta_aff_tps = self.model(batch)
resizeTgt = GeometricTnf(out_h=target_shape[0],
out_w=target_shape[1],
use_cuda=self.use_cuda)
warped_image_aff = self.affTnf(batch['source_image'],
theta_aff.view(-1, 2, 3))
warped_image_aff_tps = affTpsTnf(batch['source_image'], theta_aff,
theta_aff_tps)
warped_image_aff_np = normalize_image(
resizeTgt(warped_image_aff),
forward=False).data.squeeze(0).transpose(0, 1).transpose(
1, 2).cpu().numpy()
warped_image_aff_tps_np = normalize_image(
resizeTgt(warped_image_aff_tps),
forward=False).data.squeeze(0).transpose(0, 1).transpose(
1, 2).cpu().numpy()
warped_image_aff_tps_np = warped_image_aff_tps_np.astype(np.uint8)
warped_image_aff_np = warped_image_aff_np.astype(np.uint8)
return warped_image_aff_np, warped_image_aff_tps_np
def __init__(self,
threshold=config.THRESHOLD,
use_cuda=True):
super(CosegLoss, self).__init__()
self.threshold = threshold
self.affTnf = GeometricTnf(geometric_model='affine',
out_h=240,
out_w=240,
use_cuda=use_cuda)
self.extractor = models.resnet50(pretrained=True)
self.extractor = nn.Sequential(*list(self.extractor.children())[:-1])
for name,param in self.extractor.named_parameters():
param.requires_grad = False
if use_cuda:
self.extractor = self.extractor.cuda()
def __init__(self,
csv_file,
dataset_path,
output_size=(240, 240),
transform=None,
dataset_size=0):
self.out_h, self.out_w = output_size
self.pairs = pd.read_csv(csv_file)
if dataset_size != 0:
dataset_size = min((dataset_size, len(self.pairs)))
self.pairs = self.pairs.iloc[0:dataset_size, :]
self.img_A_names = self.pairs.iloc[:, 0]
self.img_B_names = self.pairs.iloc[:, 1]
self.dataset_path = dataset_path
self.transform = transform
# no cuda as dataset is called from CPU threads in dataloader and produces confilct
self.affineTnf = GeometricTnf(out_h=self.out_h,
out_w=self.out_w,
use_cuda=False)
def __init__(self, csv_file, training_image_path, output_size=(240,240), geometric_model='affine', transform=None,
random_sample=False, random_t=0.5, random_s=0.5, random_alpha=1/6, random_t_tps=0.4):
# random_sample is used to indicate whether deformation coefficients are randomly generated?
self.random_sample = random_sample
self.random_t = random_t
self.random_t_tps = random_t_tps
self.random_alpha = random_alpha
self.random_s = random_s
self.out_h, self.out_w = output_size
# read csv file
self.train_data = pd.read_csv(csv_file)
self.img_A_names = self.train_data.iloc[:,0]
self.img_B_names = self.train_data.iloc[:,1]
self.theta_array = self.train_data.iloc[:, 2:].values.astype('float')
# copy arguments
self.training_image_path = training_image_path
self.transform = transform
self.geometric_model = geometric_model
# affine transform used to rescale images
self.affineTnf = GeometricTnf(out_h=self.out_h, out_w=self.out_w, use_cuda = False)