Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

For Python 3 and recent PyTorch #1

Open
wants to merge 4 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 57 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]

# C extensions
*.so

# Distribution / packaging
.Python
env/
build/
develop-eggs/
dist/
downloads/
eggs/
lib/
lib64/
parts/
sdist/
var/
*.egg-info/
.installed.cfg
*.egg

# PyInstaller
# Usually these files are written by a python script from a template
# before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec

# Installer logs
pip-log.txt
pip-delete-this-directory.txt

# Unit test / coverage reports
htmlcov/
.tox/
.coverage
.cache
nosetests.xml
coverage.xml

# Translations
*.mo
*.pot

# Django stuff:
*.log

# Sphinx documentation
docs/_build/

# PyBuilder
target/

# ignore saved weights files
weights/
121 changes: 61 additions & 60 deletions model/PieAPPv0pt1_PT.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,64 +8,65 @@
import skimage.io as io

class PieAPP(nn.Module): # How to ensure that everything goes on a GPU? do I need to fetch?
def __init__(self,batch_size,num_patches):
super(PieAPP, self).__init__()
self.conv1 = nn.Conv2d(3,64,3,padding=1)
self.conv2 = nn.Conv2d(64,64,3,padding=1)
self.pool2 = nn.MaxPool2d(2,2)
self.conv3 = nn.Conv2d(64,64,3,padding=1)
self.conv4 = nn.Conv2d(64,128,3,padding=1)
self.pool4 = nn.MaxPool2d(2,2)
self.conv5 = nn.Conv2d(128,128,3,padding=1)
self.conv6 = nn.Conv2d(128,128,3,padding=1)
self.pool6 = nn.MaxPool2d(2,2)
self.conv7 = nn.Conv2d(128,256,3,padding=1)
self.conv8 = nn.Conv2d(256,256,3,padding=1)
self.pool8 = nn.MaxPool2d(2,2)
self.conv9 = nn.Conv2d(256,256,3,padding=1)
self.conv10 = nn.Conv2d(256,512,3,padding=1)
self.pool10 = nn.MaxPool2d(2,2)
self.conv11 = nn.Conv2d(512,512,3,padding=1)
self.fc1_score = nn.Linear(120832, 512)
self.fc2_score = nn.Linear(512,1)
self.fc1_weight = nn.Linear(2048,512)
self.fc2_weight = nn.Linear(512,1)
self.ref_score_subtract = nn.Linear(1,1)
self.batch_size = batch_size
self.num_patches = num_patches
def __init__(self,batch_size,num_patches):
super(PieAPP, self).__init__()
self.conv1 = nn.Conv2d(3,64,3,padding=1)
self.conv2 = nn.Conv2d(64,64,3,padding=1)
self.pool2 = nn.MaxPool2d(2,2)
self.conv3 = nn.Conv2d(64,64,3,padding=1)
self.conv4 = nn.Conv2d(64,128,3,padding=1)
self.pool4 = nn.MaxPool2d(2,2)
self.conv5 = nn.Conv2d(128,128,3,padding=1)
self.conv6 = nn.Conv2d(128,128,3,padding=1)
self.pool6 = nn.MaxPool2d(2,2)
self.conv7 = nn.Conv2d(128,256,3,padding=1)
self.conv8 = nn.Conv2d(256,256,3,padding=1)
self.pool8 = nn.MaxPool2d(2,2)
self.conv9 = nn.Conv2d(256,256,3,padding=1)
self.conv10 = nn.Conv2d(256,512,3,padding=1)
self.pool10 = nn.MaxPool2d(2,2)
self.conv11 = nn.Conv2d(512,512,3,padding=1)
self.fc1_score = nn.Linear(120832, 512)
self.fc2_score = nn.Linear(512,1)
self.fc1_weight = nn.Linear(2048,512)
self.fc2_weight = nn.Linear(512,1)
self.ref_score_subtract = nn.Linear(1,1)
self.batch_size = int(batch_size)
self.num_patches = int(num_patches)

def flatten(self,matrix): # takes NxCxHxW input and outputs NxHWC
return matrix.view((self.batch_size*self.num_patches,-1))
Copy link

@ShirAmir ShirAmir Feb 12, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should be:

return matrix.reshape((self.batch_size*self.num_patches,-1))

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Was it hitting an error about non-contiguous tensors?


def compute_features(self,input):
#conv1 -> relu -> conv2 -> relu -> pool2 -> conv3 -> relu
x3 = F.relu(self.conv3(self.pool2(F.relu(self.conv2(F.relu(self.conv1(input)))))))
# conv4 -> relu -> pool4 -> conv5 -> relu
x5 = F.relu(self.conv5(self.pool4(F.relu(self.conv4(x3)))))
# conv6 -> relu -> pool6 -> conv7 -> relu
x7 = F.relu(self.conv7(self.pool6(F.relu(self.conv6(x5)))))
# conv8 -> relu -> pool8 -> conv9 -> relu
x9 = F.relu(self.conv9(self.pool8(F.relu(self.conv8(x7)))))
# conv10 -> relu -> pool10 -> conv11 -> relU
x11 = self.flatten(F.relu(self.conv11(self.pool10(F.relu(self.conv10(x9))))))
# flatten and concatenate
feature_ms = torch.cat((self.flatten(x3),self.flatten(x5),self.flatten(x7),self.flatten(x9),x11),1)
return feature_ms, x11

def compute_score(self,image_A_patches, image_ref_patches):
A_multi_scale, A_coarse = self.compute_features(image_A_patches)
ref_multi_scale, ref_coarse = self.compute_features(image_ref_patches)
diff_ms = ref_multi_scale - A_multi_scale
diff_coarse = ref_coarse - A_coarse
# per patch score: fc1_score -> relu -> fc2_score
per_patch_score = self.ref_score_subtract(0.01*self.fc2_score(F.relu(self.fc1_score(diff_ms))))
per_patch_score.view((-1,self.num_patches))
# per patch weight: fc1_weight -> relu -> fc2_weight
const = Variable(torch.from_numpy(0.000001*np.ones((1,))).float(), requires_grad=False)
const_cuda = const.cuda()
per_patch_weight = self.fc2_weight(F.relu(self.fc1_weight(diff_coarse)))+const_cuda
per_patch_weight.view((-1,self.num_patches))
product_val = torch.mul(per_patch_weight,per_patch_score)
dot_product_val = torch.sum(product_val)
norm_factor = torch.sum(per_patch_weight)
return torch.div(dot_product_val, norm_factor), per_patch_score, per_patch_weight
def flatten(self,matrix): # takes NxCxHxW input and outputs NxHWC
return matrix.view((int(self.batch_size*self.num_patches),-1))

def compute_features(self,input):
#conv1 -> relu -> conv2 -> relu -> pool2 -> conv3 -> relu
x3 = F.relu(self.conv3(self.pool2(F.relu(self.conv2(F.relu(self.conv1(input)))))))
# conv4 -> relu -> pool4 -> conv5 -> relu
x5 = F.relu(self.conv5(self.pool4(F.relu(self.conv4(x3)))))
# conv6 -> relu -> pool6 -> conv7 -> relu
x7 = F.relu(self.conv7(self.pool6(F.relu(self.conv6(x5)))))
# conv8 -> relu -> pool8 -> conv9 -> relu
x9 = F.relu(self.conv9(self.pool8(F.relu(self.conv8(x7)))))
# conv10 -> relu -> pool10 -> conv11 -> relU
x11 = self.flatten(F.relu(self.conv11(self.pool10(F.relu(self.conv10(x9))))))
# flatten and concatenate
feature_ms = torch.cat((self.flatten(x3),self.flatten(x5),self.flatten(x7),self.flatten(x9),x11),1)
return feature_ms, x11

def compute_score(self,image_A_patches, image_ref_patches):
A_multi_scale, A_coarse = self.compute_features(image_A_patches)
ref_multi_scale, ref_coarse = self.compute_features(image_ref_patches)
diff_ms = ref_multi_scale - A_multi_scale
diff_coarse = ref_coarse - A_coarse
# per patch score: fc1_score -> relu -> fc2_score
per_patch_score = self.ref_score_subtract(0.01*self.fc2_score(F.relu(self.fc1_score(diff_ms))))
per_patch_score.view((-1,int(self.num_patches)))
# per patch weight: fc1_weight -> relu -> fc2_weight
const = Variable(torch.from_numpy(0.000001*np.ones((1,))).float(), requires_grad=False)
if next(self.fc1_weight.parameters()).is_cuda:
const = const.cuda()
per_patch_weight = self.fc2_weight(F.relu(self.fc1_weight(diff_coarse)))+const
per_patch_weight.view((-1,int(self.num_patches)))
product_val = torch.mul(per_patch_weight,per_patch_score)
dot_product_val = torch.sum(product_val)
norm_factor = torch.sum(per_patch_weight)
return torch.div(dot_product_val, norm_factor), per_patch_score, per_patch_weight
91 changes: 48 additions & 43 deletions test_PieAPP_PT.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
from __future__ import print_function

import numpy as np
import cv2
import sys
Expand All @@ -10,13 +12,14 @@
import argparse
import os


######## check for model and download if not present
if not os.path.isfile('weights/PieAPPv0.1.pth'):
print "downloading dataset"
os.system("bash scripts/download_PieAPPv0.1_PT_weights.sh")
if not os.path.isfile('weights/PieAPPv0.1.pth'):
print "PieAPPv0.1.pth not downloaded"
sys.exit()
print("downloading dataset")
os.system("bash scripts/download_PieAPPv0.1_PT_weights.sh")
if not os.path.isfile('weights/PieAPPv0.1.pth'):
print("PieAPPv0.1.pth not downloaded")
sys.exit()

######## variables
patch_size = 64
Expand All @@ -38,9 +41,9 @@
_,rows,cols,ch = imagesRef.shape

if args.sampling_mode == 'sparse':
stride_val = 27
stride_val = 27
else:
stride_val = 6
stride_val = 6

try:
gpu_num = float(args.gpu_id)
Expand All @@ -58,46 +61,48 @@

######## initialize the model
PieAPP_net = PieAPP(batch_size,num_patches_per_dim)
PieAPP_net.load_state_dict(torch.load('weights/PieAPPv0.1.pth'))
state_dict = torch.load('weights/PieAPPv0.1.pth')
state_dict['ref_score_subtract.weight'] = state_dict['ref_score_subtract.weight'].unsqueeze(1)
PieAPP_net.load_state_dict(state_dict)

if use_gpu == 1:
PieAPP_net.cuda()
PieAPP_net.cuda()

score_accum = 0.0
weight_accum = 0.0

# iterate through smaller size sub-images (to prevent memory overload)
for x_iter in range(0,num_x/num_patches_per_dim+1):
for y_iter in range(0,num_y/num_patches_per_dim+1):
# compute the size of the subimage
if (num_patches_per_dim*(x_iter + 1) > num_x):
size_slice_cols = cols - x_loc[num_patches_per_dim*x_iter]
else:
size_slice_cols = x_loc[num_patches_per_dim*(x_iter + 1)] - x_loc[num_patches_per_dim*x_iter] + patch_size - stride_val
if (num_patches_per_dim*(y_iter + 1) > num_y):
size_slice_rows = cols - y_loc[num_patches_per_dim*y_iter]
else:
size_slice_rows = y_loc[num_patches_per_dim*(y_iter + 1)] - y_loc[num_patches_per_dim*y_iter] + patch_size - stride_val
# obtain the subimage and samples patches
A_sub_im = imagesA[:, y_loc[num_patches_per_dim*y_iter]:y_loc[num_patches_per_dim*y_iter]+size_slice_rows, x_loc[num_patches_per_dim*x_iter]:x_loc[num_patches_per_dim*x_iter]+size_slice_cols,:]
ref_sub_im = imagesRef[:, y_loc[num_patches_per_dim*y_iter]:y_loc[num_patches_per_dim*y_iter]+size_slice_rows, x_loc[num_patches_per_dim*x_iter]:x_loc[num_patches_per_dim*x_iter]+size_slice_cols,:]
A_patches, ref_patches = sample_patches(A_sub_im, ref_sub_im, patch_size=64, strideval=stride_val, random_selection=False, uniform_grid_mode = 'strided')
num_patches_curr = A_patches.shape[0]/batch_size
PieAPP_net.num_patches = num_patches_curr
# initialize variable to be fed to PieAPP_net
A_patches_var = Variable(torch.from_numpy(np.transpose(A_patches,(0,3,1,2))), requires_grad=False)
ref_patches_var = Variable(torch.from_numpy(np.transpose(ref_patches,(0,3,1,2))), requires_grad=False)
if use_gpu == 1:
A_patches_var = A_patches_var.cuda()
ref_patches_var = ref_patches_var.cuda()

# forward pass
_, PieAPP_patchwise_errors, PieAPP_patchwise_weights = PieAPP_net.compute_score(A_patches_var.float(), ref_patches_var.float())
curr_err = PieAPP_patchwise_errors.cpu().data.numpy()
curr_weights = PieAPP_patchwise_weights.cpu().data.numpy()
score_accum += np.sum(np.multiply(curr_err, curr_weights))
weight_accum += np.sum(curr_weights)

print 'PieAPP value of '+args.A_path+ ' with respect to: '+str(score_accum/weight_accum)
for x_iter in range(0,num_x//num_patches_per_dim+1):
for y_iter in range(0,num_y//num_patches_per_dim+1):
# compute the size of the subimage
if (num_patches_per_dim*(x_iter + 1) > num_x):
size_slice_cols = cols - x_loc[num_patches_per_dim*x_iter]
else:
size_slice_cols = x_loc[num_patches_per_dim*(x_iter + 1)] - x_loc[num_patches_per_dim*x_iter] + patch_size - stride_val
if (num_patches_per_dim*(y_iter + 1) > num_y):
size_slice_rows = cols - y_loc[num_patches_per_dim*y_iter]
else:
size_slice_rows = y_loc[num_patches_per_dim*(y_iter + 1)] - y_loc[num_patches_per_dim*y_iter] + patch_size - stride_val
# obtain the subimage and samples patches
A_sub_im = imagesA[:, y_loc[num_patches_per_dim*y_iter]:y_loc[num_patches_per_dim*y_iter]+size_slice_rows, x_loc[num_patches_per_dim*x_iter]:x_loc[num_patches_per_dim*x_iter]+size_slice_cols,:]
ref_sub_im = imagesRef[:, y_loc[num_patches_per_dim*y_iter]:y_loc[num_patches_per_dim*y_iter]+size_slice_rows, x_loc[num_patches_per_dim*x_iter]:x_loc[num_patches_per_dim*x_iter]+size_slice_cols,:]
A_patches, ref_patches = sample_patches(A_sub_im, ref_sub_im, patch_size=64, strideval=stride_val, random_selection=False, uniform_grid_mode = 'strided')
num_patches_curr = A_patches.shape[0]/batch_size
PieAPP_net.num_patches = num_patches_curr
# initialize variable to be fed to PieAPP_net
A_patches_var = Variable(torch.from_numpy(np.transpose(A_patches,(0,3,1,2))), requires_grad=False)
ref_patches_var = Variable(torch.from_numpy(np.transpose(ref_patches,(0,3,1,2))), requires_grad=False)
if use_gpu == 1:
A_patches_var = A_patches_var.cuda()
ref_patches_var = ref_patches_var.cuda()

# forward pass
_, PieAPP_patchwise_errors, PieAPP_patchwise_weights = PieAPP_net.compute_score(A_patches_var.float(), ref_patches_var.float())
curr_err = PieAPP_patchwise_errors.cpu().data.numpy()
curr_weights = PieAPP_patchwise_weights.cpu().data.numpy()
score_accum += np.sum(np.multiply(curr_err, curr_weights))
weight_accum += np.sum(curr_weights)

print('PieAPP value of '+args.A_path+ ' with respect to: '+str(score_accum/weight_accum))