使用pytorch实现在图像上单机鼠标截取两个相同尺寸的区域并利用孪生神经网络度量其相似性,并将截图相似的部分变成绿色,输出实现代码
时间: 2024-02-06 08:10:05 浏览: 123
以下是实现的代码:
```python
import cv2
import torch
import torchvision.transforms as transforms
from PIL import Image
# 加载孪生神经网络
class SiameseNetwork(torch.nn.Module):
def __init__(self):
super(SiameseNetwork, self).__init__()
self.cnn = torch.nn.Sequential(
torch.nn.Conv2d(3, 32, kernel_size=8),
torch.nn.ReLU(inplace=True),
torch.nn.MaxPool2d(kernel_size=2, stride=2),
torch.nn.Conv2d(32, 64, kernel_size=6),
torch.nn.ReLU(inplace=True),
torch.nn.MaxPool2d(kernel_size=2, stride=2),
torch.nn.Conv2d(64, 128, kernel_size=4),
torch.nn.ReLU(inplace=True),
torch.nn.MaxPool2d(kernel_size=2, stride=2),
torch.nn.Conv2d(128, 256, kernel_size=2),
torch.nn.ReLU(inplace=True),
torch.nn.Flatten(),
torch.nn.Linear(256*2*2, 512),
torch.nn.ReLU(inplace=True),
torch.nn.Linear(512, 2)
)
def forward_once(self, x):
output = self.cnn(x)
return output
def forward(self, input1, input2):
output1 = self.forward_once(input1)
output2 = self.forward_once(input2)
return output1, output2
# 加载模型
model = SiameseNetwork()
model.load_state_dict(torch.load('model.pth'))
# 定义图像变换
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor()
])
# 鼠标截取两个相似区域
def on_mouse(event, x, y, flags, param):
global img1, img2, point1, point2
if event == cv2.EVENT_LBUTTONDOWN:
if point1 is None:
point1 = (x, y)
elif point2 is None:
point2 = (x, y)
img1_crop = img1[point1[1]:point2[1], point1[0]:point2[0]]
img2_crop = img2[point1[1]:point2[1], point1[0]:point2[0]]
img1_crop_pil = Image.fromarray(cv2.cvtColor(img1_crop, cv2.COLOR_BGR2RGB))
img2_crop_pil = Image.fromarray(cv2.cvtColor(img2_crop, cv2.COLOR_BGR2RGB))
img1_crop_tensor = transform(img1_crop_pil).unsqueeze(0)
img2_crop_tensor = transform(img2_crop_pil).unsqueeze(0)
with torch.no_grad():
output1, output2 = model(img1_crop_tensor, img2_crop_tensor)
similarity = torch.nn.functional.cosine_similarity(output1, output2).item()
print('相似度:', similarity)
if similarity > 0.8:
img1[point1[1]:point2[1], point1[0]:point2[0]] = [0, 255, 0]
img2[point1[1]:point2[1], point1[0]:point2[0]] = [0, 255, 0]
else:
print('两个区域不相似')
point1 = None
point2 = None
# 加载图像
img1 = cv2.imread('img1.jpg')
img2 = cv2.imread('img2.jpg')
# 显示图像
cv2.namedWindow('Image1')
cv2.namedWindow('Image2')
cv2.setMouseCallback('Image1', on_mouse)
cv2.setMouseCallback('Image2', on_mouse)
point1 = None
point2 = None
while True:
cv2.imshow('Image1', img1)
cv2.imshow('Image2', img2)
if cv2.waitKey(1) == 27:
break
cv2.destroyAllWindows()
```
在代码中,我们首先加载了一个孪生神经网络模型,这个模型的作用是度量两个图像之间的相似性。然后定义了一个图像变换,用于将读取的图像转换成模型所需的输入格式。接着实现了鼠标事件的回调函数,当鼠标左键点击时,程序会记录下第一个点击的点,然后等待第二个点击的点,将两个点击点之间的区域截取下来,并将其转换成模型所需的输入格式。然后使用模型将这两个区域进行度量,并计算相似度。如果相似度大于0.8,则将这两个区域在原图像上标记为绿色。最后,在图像窗口中注册鼠标事件回调函数,并循环显示两个图像,等待用户操作。
阅读全文