515 lines
18 KiB
Python
515 lines
18 KiB
Python
#!/usr/bin/env python
|
|
# coding: utf-8
|
|
|
|
# In[1]:
|
|
|
|
|
|
#Tous les codes sont basés sur l'environnement suivant
|
|
#python 3.7
|
|
#opencv 3.1.0
|
|
#pytorch 1.4.0
|
|
|
|
import torch
|
|
from torch.autograd import Variable
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
import cv2
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
import random
|
|
import math
|
|
import pickle
|
|
import random
|
|
from PIL import Image
|
|
import sys
|
|
|
|
|
|
# In[2]:
|
|
|
|
|
|
# Les fonctions dans ce bloc ne sont pas utilisées par le réseau, mais certaines fonctions d'outils
|
|
|
|
# Les fonctions de ce bloc se trouvent dans le programme d'apprentissage
|
|
# “Apprentissage_MSELoss_avec_GPU“
|
|
# et les commentaires détaillés se trouvent dans le programme d'apprentissage
|
|
|
|
def tensor_imshow(im_tensor,cannel):
|
|
b,c,h,w=im_tensor.shape
|
|
if c==1:
|
|
plt.imshow(im_tensor.squeeze().detach().numpy())
|
|
else:
|
|
plt.imshow(im_tensor.squeeze().detach().numpy()[cannel,:])
|
|
|
|
def get_training_fragment(frag_size,im):
|
|
h,w,c=im.shape
|
|
n=random.randint(0,int(h/frag_size)-1)
|
|
m=random.randint(0,int(w/frag_size)-1)
|
|
|
|
shape=frag_size/4
|
|
vt_h=math.ceil((h+1)/shape)
|
|
vt_w=math.ceil((w+1)/shape)
|
|
vt=np.zeros([vt_h,vt_w])
|
|
vt_h_po=round((vt_h-1)*(n*frag_size/(h-1)+(n+1)*frag_size/(h-1))/2)
|
|
vt_w_po=round((vt_w-1)*(m*frag_size/(w-1)+(m+1)*frag_size/(w-1))/2)
|
|
vt[vt_h_po,vt_w_po]=1
|
|
vt = np.float32(vt)
|
|
vt=torch.from_numpy(vt.reshape(1,1,vt_h,vt_w))
|
|
|
|
return im[n*frag_size:(n+1)*frag_size,m*frag_size:(m+1)*frag_size,:],vt
|
|
|
|
def write_result_in_file(result,file_name):
|
|
n=0
|
|
with open(file_name,'w') as file:
|
|
for i in range(len(result)):
|
|
while n<result[i][0]:
|
|
s=str(n)
|
|
n=n+1
|
|
s=s+"\n"
|
|
file.write(s)
|
|
s=str(result[i][0])+" "+str(result[i][1])+" "+str(result[i][2])+" "+str(result[i][3])
|
|
s=s+"\n"
|
|
n=n+1
|
|
file.write(s)
|
|
file.close()
|
|
|
|
|
|
def img2tensor(im):
|
|
im=np.array(im,dtype="float32")
|
|
tensor_cv = torch.from_numpy(np.transpose(im, (2, 0, 1)))
|
|
im_tensor=tensor_cv.unsqueeze(0)
|
|
return im_tensor
|
|
|
|
def show_coordonnee(position_pred):
|
|
map_corre=position_pred.squeeze().detach().numpy()
|
|
score=sum(sum(map_corre))
|
|
h,w=map_corre.shape
|
|
max_value=map_corre.max()
|
|
coordonnee=np.where(map_corre==max_value)
|
|
return score,coordonnee[0].mean()/h,coordonnee[1].mean()/w
|
|
|
|
def test_fragment32_32(frag,seuillage):
|
|
a=frag[:,:,0]+frag[:,:,1]+frag[:,:,2]
|
|
mask = (a == 0)
|
|
arr_new = a[mask]
|
|
if arr_new.size/a.size<=(1-seuillage):
|
|
return True
|
|
else:
|
|
return False
|
|
|
|
def save_net(file_path,net):
|
|
pkl_file = open(file_path, 'wb')
|
|
pickle.dump(net,pkl_file)
|
|
pkl_file.close()
|
|
|
|
def load_net(file_path):
|
|
pkl_file = open(file_path, 'rb')
|
|
net= pickle.load(pkl_file)
|
|
pkl_file.close()
|
|
return net
|
|
|
|
|
|
# In[3]:
|
|
|
|
|
|
# Les fonctions de ce bloc sont utilisées pour construire le réseau
|
|
|
|
# Les fonctions de ce bloc se trouvent dans le programme d'apprentissage
|
|
# “Apprentissage_MSELoss_avec_GPU“
|
|
# et les commentaires détaillés se trouvent dans le programme d'apprentissage
|
|
|
|
def ini():
|
|
kernel=torch.zeros([8,3,3,3])
|
|
array_0=np.array([[1,2,1],[0,0,0],[-1,-2,-1]],dtype='float32')
|
|
array_1=np.array([[2,1,0],[1,0,-1],[0,-1,-2]],dtype='float32')
|
|
array_2=np.array([[1,0,-1],[2,0,-2],[1,0,-1]],dtype='float32')
|
|
array_3=np.array([[0,-1,-2],[1,0,-1],[2,1,0]],dtype='float32')
|
|
array_4=np.array([[-1,-2,-1],[0,0,0],[1,2,1]],dtype='float32')
|
|
array_5=np.array([[-2,-1,0],[-1,0,1],[0,1,2]],dtype='float32')
|
|
array_6=np.array([[-1,0,1],[-2,0,2],[-1,0,1]],dtype='float32')
|
|
array_7=np.array([[0,1,2],[-1,0,1],[-2,-1,0]],dtype='float32')
|
|
for i in range(3):
|
|
kernel[0,i,:]=torch.from_numpy(array_0)
|
|
kernel[1,i,:]=torch.from_numpy(array_1)
|
|
kernel[2,i,:]=torch.from_numpy(array_2)
|
|
kernel[3,i,:]=torch.from_numpy(array_3)
|
|
kernel[4,i,:]=torch.from_numpy(array_4)
|
|
kernel[5,i,:]=torch.from_numpy(array_5)
|
|
kernel[6,i,:]=torch.from_numpy(array_6)
|
|
kernel[7,i,:]=torch.from_numpy(array_7)
|
|
return torch.nn.Parameter(kernel,requires_grad=True)
|
|
|
|
def kernel_add_ini(n,m):
|
|
input_canal=int(n*m)
|
|
output_canal=int(n/2)*int(m/2)
|
|
for i in range(int(n/2)):
|
|
for j in range(int(m/2)):
|
|
kernel_add=np.zeros([1,input_canal],dtype='float32')
|
|
kernel_add[0,i*2*m+j*2]=1
|
|
kernel_add[0,i*2*m+j*2+1]=1
|
|
kernel_add[0,(i*2+1)*m+j*2]=1
|
|
kernel_add[0,(i*2+1)*m+j*2+1]=1
|
|
if i==0 and j==0:
|
|
add=torch.from_numpy(kernel_add.reshape(1,input_canal,1,1))
|
|
else:
|
|
add_=torch.from_numpy(kernel_add.reshape(1,input_canal,1,1))
|
|
add=torch.cat((add,add_),0)
|
|
return torch.nn.Parameter(add,requires_grad=False)
|
|
|
|
def kernel_shift_ini(n,m):
|
|
input_canal=int(n*m)
|
|
output_canal=int(n*m)
|
|
|
|
kernel_shift=torch.zeros([output_canal,input_canal,3,3])
|
|
|
|
array_0=np.array([[1,0,0],[0,0,0],[0,0,0]],dtype='float32')
|
|
array_1=np.array([[0,0,1],[0,0,0],[0,0,0]],dtype='float32')
|
|
array_2=np.array([[0,0,0],[0,0,0],[1,0,0]],dtype='float32')
|
|
array_3=np.array([[0,0,0],[0,0,0],[0,0,1]],dtype='float32')
|
|
|
|
kernel_shift_0=torch.from_numpy(array_0)
|
|
kernel_shift_1=torch.from_numpy(array_1)
|
|
kernel_shift_2=torch.from_numpy(array_2)
|
|
kernel_shift_3=torch.from_numpy(array_3)
|
|
|
|
|
|
for i in range(n):
|
|
for j in range(m):
|
|
if i==0 and j==0:
|
|
kernel_shift[0,0,:]=kernel_shift_0
|
|
else:
|
|
if i%2==0 and j%2==0:
|
|
kernel_shift[i*m+j,i*m+j,:]=kernel_shift_0
|
|
if i%2==0 and j%2==1:
|
|
kernel_shift[i*m+j,i*m+j,:]=kernel_shift_1
|
|
if i%2==1 and j%2==0:
|
|
kernel_shift[i*m+j,i*m+j,:]=kernel_shift_2
|
|
if i%2==1 and j%2==1:
|
|
kernel_shift[i*m+j,i*m+j,:]=kernel_shift_3
|
|
|
|
return torch.nn.Parameter(kernel_shift,requires_grad=False)
|
|
|
|
def get_patch(fragment,psize,n,m):
|
|
return fragment[:,:,n*psize:(n+1)*psize,m*psize:(m+1)*psize]
|
|
|
|
class Net(nn.Module):
|
|
def __init__(self,frag_size,psize):
|
|
super(Net, self).__init__()
|
|
|
|
h_fr=frag_size
|
|
w_fr=frag_size
|
|
|
|
n=int(h_fr/psize) #n*m patches
|
|
m=int(w_fr/psize)
|
|
|
|
self.conv1 = nn.Conv2d(3,8,kernel_size=3,stride=1,padding=1)
|
|
#self.conv1.weight=ini()
|
|
self.Relu = nn.ReLU(inplace=True)
|
|
self.maxpooling=nn.MaxPool2d(3,stride=2, padding=1)
|
|
|
|
self.shift1=nn.Conv2d(n*m,n*m,kernel_size=3,stride=1,padding=1)
|
|
self.shift1.weight=kernel_shift_ini(n,m)
|
|
self.add1 = nn.Conv2d(n*m,int(n/2)*int(m/2),kernel_size=1,stride=1,padding=0)
|
|
self.add1.weight=kernel_add_ini(n,m)
|
|
|
|
n=int(n/2)
|
|
m=int(m/2)
|
|
if n>=2 and m>=2:
|
|
self.shift2=nn.Conv2d(n*m,n*m,kernel_size=3,stride=1,padding=1)
|
|
self.shift2.weight=kernel_shift_ini(n,m)
|
|
self.add2 = nn.Conv2d(n*m,int(n/2)*int(m/2),kernel_size=1,stride=1,padding=0)
|
|
self.add2.weight=kernel_add_ini(n,m)
|
|
|
|
n=int(n/2)
|
|
m=int(m/2)
|
|
if n>=2 and m>=2:
|
|
self.shift3=nn.Conv2d(n*m,n*m,kernel_size=3,stride=1,padding=1)
|
|
self.shift3.weight=kernel_shift_ini(n,m)
|
|
self.add3 = nn.Conv2d(n*m,int(n/2)*int(m/2),kernel_size=1,stride=1,padding=0)
|
|
self.add3.weight=kernel_add_ini(n,m)
|
|
|
|
|
|
def get_descripteur(self,img,using_cuda):
|
|
descripteur_img=self.Relu(self.conv1(img))
|
|
b,c,h,w=descripteur_img.shape
|
|
couche_constante=0.5*torch.ones([1,1,h,w])
|
|
if using_cuda:
|
|
couche_constante=couche_constante.cuda()
|
|
descripteur_img=torch.cat((descripteur_img,couche_constante),1)
|
|
descripteur_img_norm=descripteur_img/torch.norm(descripteur_img,dim=1)
|
|
return descripteur_img_norm
|
|
|
|
def forward(self,img,frag,using_cuda):
|
|
psize=4
|
|
|
|
descripteur_input1=self.get_descripteur(img,using_cuda)
|
|
descripteur_input2=self.get_descripteur(frag,using_cuda)
|
|
|
|
b,c,h,w=frag.shape
|
|
n=int(h/psize)
|
|
m=int(w/psize)
|
|
|
|
for i in range(n):
|
|
for j in range(m):
|
|
if i==0 and j==0:
|
|
map_corre=F.conv2d(descripteur_input1,get_patch(descripteur_input2,psize,i,j),padding=2)
|
|
else:
|
|
a=F.conv2d(descripteur_input1,get_patch(descripteur_input2,psize,i,j),padding=2)
|
|
map_corre=torch.cat((map_corre,a),1)
|
|
#shift
|
|
map_corre=self.maxpooling(map_corre)
|
|
map_corre=self.shift1(map_corre)
|
|
map_corre=self.add1(map_corre)
|
|
|
|
|
|
n=int(n/2)
|
|
m=int(m/2)
|
|
if n>=2 and m>=2:
|
|
map_corre=self.maxpooling(map_corre)
|
|
map_corre=self.shift2(map_corre)
|
|
map_corre=self.add2(map_corre)
|
|
|
|
|
|
n=int(n/2)
|
|
m=int(m/2)
|
|
if n>=2 and m>=2:
|
|
map_corre=self.maxpooling(map_corre)
|
|
map_corre=self.shift3(map_corre)
|
|
map_corre=self.add3(map_corre)
|
|
|
|
|
|
b,c,h,w=map_corre.shape
|
|
map_corre=map_corre/(map_corre.max())
|
|
#map_corre=(F.softmax(map_corre.reshape(1,1,h*w,1),dim=2)).reshape(b,c,h,w)
|
|
return map_corre
|
|
|
|
|
|
# In[4]:
|
|
|
|
|
|
# Les fonctions de ce bloc sont utilisées pour appliquer le réseau à des fragments (pas à des patchs carrés)
|
|
|
|
|
|
# Cette fonction permet de sélectionner un ensemble de patchs carrés à partir d'un fragment
|
|
# Le paramètre “frag_size” fait ici référence à la taille du patch d'entrée carré (16 * 16)
|
|
# Le paramètre “seuillage” limite la proportion de pixels non noirs dans chaque patch
|
|
# Le paramètre “limite” peut limiter le nombre de correctifs trouvés dans chaque fragment
|
|
def get_patch_list(frag,frag_size,limite,seuillage):
|
|
n=0
|
|
m=0
|
|
h,w,c=frag.shape
|
|
patch_list=[]
|
|
position_list=[]
|
|
for i in range(4):
|
|
if len(patch_list)>limite and limite!=0:
|
|
break
|
|
for j in range(4):
|
|
if len(patch_list)>limite and limite!=0:
|
|
break
|
|
n_offset=i*4 # n offset
|
|
m_offset=j*4 # m offset
|
|
n=0
|
|
while n+frag_size+n_offset<h:
|
|
m=0
|
|
while m+frag_size+m_offset<w:
|
|
patch=frag[n+n_offset:n+frag_size+n_offset,m+m_offset:m+frag_size+m_offset,:]
|
|
if test_fragment32_32(patch,seuillage):
|
|
patch_list.append(patch)
|
|
position_list.append([int((n+frag_size/2)+n_offset),int((m+frag_size/2)+m_offset)])
|
|
m=m+frag_size
|
|
n=n+frag_size
|
|
return patch_list,position_list
|
|
|
|
# Entrez du fragment et de la fresque, exécutez le réseau
|
|
def run_net_v3(net,img,frag,frag_size,limite,seuillage,using_cuda,rotation):
|
|
Img=Image.fromarray(frag)
|
|
frag=np.array(Img.rotate(rotation))
|
|
img_tensor=img2tensor(img)
|
|
|
|
# la collection de patchs carrée dans le fragement "sont frag_list[]"
|
|
# La position de leur centre dans la fragment sont "position_frag[]"
|
|
frag_list,position_frag=get_patch_list(frag,frag_size,limite,seuillage)
|
|
if using_cuda:
|
|
img_tensor=img_tensor.cuda()
|
|
|
|
score_list=[]
|
|
coordonnee_list=[]
|
|
|
|
# Pour chaque patch carré dans la collection, effectuez un calcul en réseau de leur position
|
|
# Le résultat est placé en "coordonnee_list[]"
|
|
# "score_list[]" pas utile dans notre programme
|
|
for i in range(len(frag_list)):
|
|
frag_tensor=img2tensor(frag_list[i])
|
|
if using_cuda:
|
|
frag_tensor=frag_tensor.cuda()
|
|
res=net.forward(img_tensor,frag_tensor,using_cuda)
|
|
if using_cuda:
|
|
res=res.cpu()
|
|
score,po_h,po_w=show_coordonnee(res)
|
|
coordonnee_list.append([po_h,po_w])
|
|
score_list.append(score)
|
|
h_img,w_img,c=img.shape
|
|
position=[]
|
|
|
|
# Mettez les paires correspondante en "position[]"
|
|
# [x,y,x',y']
|
|
# La position (x,y) dans le fragment correspond à la position (x',y') dans la fresque
|
|
for i in range(len(coordonnee_list)):
|
|
x0=position_frag[i][0]
|
|
y0=position_frag[i][1]
|
|
x1=int(round(h_img*coordonnee_list[i][0]))
|
|
y1=int(round(w_img*coordonnee_list[i][1]))
|
|
position.append([x0,y0,x1,y1])
|
|
return score_list,position
|
|
|
|
|
|
# In[12]:
|
|
|
|
|
|
# Cette partie du code consiste à implémenter l'algorithme RANSAC amélioré
|
|
|
|
# Ecrire le point sous forme [x,y,1]T,
|
|
# Utilisé pour construire l'équation de la matrice de transformation
|
|
def creer_point(x,y):
|
|
p=np.zeros((3,1))
|
|
p[0][0]=x
|
|
p[1][0]=y
|
|
p[2][0]=1
|
|
return p
|
|
|
|
# Sélectionnez aléatoirement n points sans duplication à partir de M points
|
|
def selectionner_points(n,M):
|
|
table=[]
|
|
for i in range(M):
|
|
table.append(i)
|
|
result=[]
|
|
for i in range(n):
|
|
index=random.randint(0,M-i-1)
|
|
result.append(table[index])
|
|
table[index]=table[M-1-i]
|
|
return result
|
|
|
|
# Selon la matrice de transformation affine, calculer la position centrale transformée et l'angle de rotation
|
|
def position_rotation(h,centre_frag):
|
|
centre=h@centre_frag
|
|
cos_rot=(h[0][0]+h[1][1])/2
|
|
sin_rot=(h[1][0]-h[0][1])/2
|
|
tan_rot=sin_rot/(cos_rot+0.0000001)
|
|
if cos_rot>0:
|
|
rot_frag=math.atan(tan_rot)*(180/pi)
|
|
else:
|
|
rot_frag=math.atan(tan_rot)*(180/pi)+180
|
|
rot_frag=-rot_frag
|
|
if rot_frag>0:
|
|
rot_frag-=360
|
|
return centre[0][0],centre[1][0],rot_frag
|
|
|
|
# Vérifiez les résultats de Ransac en avec des changements de distance euclidienne
|
|
def test_frag(inline,frag,fres):
|
|
itera=10
|
|
frag_inline=[]
|
|
fres_inline=[]
|
|
# Metter les coordonnées du point inline dans "frag_inline[]",et "fres_inline[]"
|
|
for i in range(np.size(inline,0)):
|
|
if inline[i]==1:
|
|
frag_inline.append([frag[i][0],frag[i][1]])
|
|
fres_inline.append([fres[i][0],fres[i][1]])
|
|
p=[]
|
|
|
|
# Faites une boucle dix fois,
|
|
# sélectionnez à chaque fois deux paires correspondantes inline
|
|
# calculer le changement de leur distance euclidienne
|
|
for i in range(itera):
|
|
point_test=selectionner_points(2,np.size(frag_inline,0))
|
|
diff_x_frag=frag_inline[point_test[1]][0]-frag_inline[point_test[0]][0]
|
|
diff_y_frag=frag_inline[point_test[1]][1]-frag_inline[point_test[0]][1]
|
|
diff_frag=sqrt(pow(diff_x_frag,2)+pow(diff_y_frag,2))
|
|
|
|
diff_x_fres=fres_inline[point_test[1]][0]-fres_inline[point_test[0]][0]
|
|
diff_y_fres=fres_inline[point_test[1]][1]-fres_inline[point_test[0]][1]
|
|
diff_fres=sqrt(pow(diff_x_fres,2)+pow(diff_y_fres,2))
|
|
if diff_frag !=0:
|
|
fsf=diff_fres/diff_frag
|
|
p.append([fsf])
|
|
result=np.mean(p)
|
|
return result
|
|
|
|
def frag_match(frag,img,position):
|
|
|
|
frag_size=frag.shape
|
|
centre_frag=creer_point(frag_size[0]/2,frag_size[1]/2)
|
|
|
|
retained_matches = []
|
|
frag=[]
|
|
fres=[]
|
|
|
|
for i in range(len(position)):
|
|
frag.append([float(position[i][0]),float(position[i][1])])
|
|
fres.append([float(position[i][2]),float(position[i][3])])
|
|
|
|
if np.size(frag)>0:
|
|
# Calculer la matrice de transformation affine à l'aide de la méthode Ransac
|
|
h,inline=cv2.estimateAffinePartial2D(np.array(frag),np.array(fres))
|
|
# Si “h” n'est pas sous la forme de matrice 2 * 3, la matrice de transformation affine n'est pas trouvée
|
|
if np.size(h)!=6:
|
|
return ([-1])
|
|
else:
|
|
x,y,rot=position_rotation(h,centre_frag)
|
|
pourcenttage=sum(inline)/np.size(frag,0)
|
|
# Le nombre de points inline doit être supérieur à un certain nombre
|
|
if sum(inline)>3:
|
|
p=test_frag(inline,frag,fres)
|
|
# La distance euclidienne entre les points correspondants ne doit pas trop changer,
|
|
# sinon cela prouve que le résultat de Ransac est incorrect
|
|
# ici,le changement de la distance euclidienne sont entre 0.7 et 1.3
|
|
if abs(p-1)<0.3:
|
|
# Ce n'est qu'alors que Ransac renvoie le résultat correct
|
|
return([round(y),round(x),round(rot,3)])
|
|
else:
|
|
return ([-2])
|
|
else:
|
|
return ([-3])
|
|
else:
|
|
return ([-4])
|
|
|
|
|
|
# In[14]:
|
|
|
|
|
|
if __name__=="__main__":
|
|
|
|
frag_size=16
|
|
using_cuda=True
|
|
net=load_net("./net_trainned6000")
|
|
img_test=cv2.imread("./fresque0.ppm")
|
|
|
|
result=[]
|
|
for n in range(315):
|
|
if n<10:
|
|
frag_test=cv2.imread("./frag_eroded0/frag_eroded_000"+str(n)+".ppm")
|
|
elif n<100:
|
|
frag_test=cv2.imread("./frag_eroded0/frag_eroded_00"+str(n)+".ppm")
|
|
else:
|
|
frag_test=cv2.imread("./frag_eroded0/frag_eroded_0"+str(n)+".ppm")
|
|
|
|
# Faites pivoter les pièces de 20 degrés à chaque fois pour correspondre, répétez 18 fois
|
|
for i in range(18):
|
|
rotation=20*i
|
|
score_list,position=run_net_v3(net,img_test,frag_test,frag_size,60,0.7,using_cuda,rotation)
|
|
frag_position=frag_match(frag_test,img_test,position)
|
|
# Lorsque Ransac obtient le bon résultat, sortez de la boucle
|
|
if len(frag_position)==3:
|
|
rotation_base=i*20
|
|
break
|
|
# Enregistrez les fragments correctement localisés dans "result[]"
|
|
if len(frag_position)==3:
|
|
frag_position[2]=rotation_base-360-frag_position[2]
|
|
if frag_position[2]>0:
|
|
frag_position[2]=frag_position[2]-360
|
|
result.append([n,frag_position[0],frag_position[1],round(frag_position[2],3)])
|
|
|
|
|
|
# In[15]:
|
|
|
|
|
|
result
|
|
|