# import numpy as np | |
import requests | |
import torch | |
from PIL import Image | |
from torch.nn.functional import cosine_similarity | |
from transformers import AutoImageProcessor, AutoModel | |
from transformers import pipeline | |
# import transformers | |
# | |
# print(transformers.__version__) | |
# | |
# img_urls = ["https://img0.baidu.com/it/u=3704428154,2884159591&fm=253&fmt=auto&app=138&f=JPEG?w=889&h=500", | |
# "https://img0.baidu.com/it/u=3704428154,2884159591&fm=253&fmt=auto&app=138&f=JPEG?w=889&h=500"] | |
# | |
# image_real = Image.open(requests.get(img_urls[0], stream=True).raw).convert("RGB") | |
# image_gen = Image.open(requests.get(img_urls[1], stream=True).raw).convert("RGB") | |
# | |
# # DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
DEVICE = torch.device('cpu') | |
# pipe = pipeline(task="image-feature-extraction", model_name="google/vit-base-patch16-384", device=DEVICE, pool=True) | |
# | |
# # 1提取图片特征向量 | |
# outputs = pipe([image_real, image_gen]) | |
# | |
# # get the length of a single output | |
# print(len(outputs[0][0])) | |
# # show outputs | |
# print(outputs) | |
# | |
# # 768 | |
# # [[[-0.03909236937761307, 0.43381670117378235, -0.06913255900144577, | |
# | |
# # 2计算图片相似度 | |
# similarity_score = cosine_similarity(torch.Tensor(outputs[0]), | |
# torch.Tensor(outputs[1]), dim=1) | |
# | |
# print(similarity_score) | |
# tensor([0.6043]) | |
# pipe = pipeline(task="image-feature-extraction", model_name="google/vit-base-patch16-224", device=DEVICE) | |
# output = pipe(image_real) | |
# | |
# # 其中第一个维度是批量大小,最后两个维度是嵌入形状。 | |
# print(np.array(outputs).shape) | |
# # (1, 197, 768) | |
# 第二种方式推理图片相似度 | |
# processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224") | |
# model = AutoModel.from_pretrained("google/vit-base-patch16-224").to(DEVICE) | |
processor = AutoImageProcessor.from_pretrained("chanhua/autotrain-izefx-v3qh0") | |
model = AutoModel.from_pretrained("chanhua/autotrain-izefx-v3qh0").to(DEVICE) | |
# tensor([0.6061], device='cuda:0', grad_fn=<SumBackward1>) | |
# pipe = pipeline(task="image-feature-extraction", model_name="google/vit-base-patch16-384", device=DEVICE, pool=True) | |
pipe = pipeline(task="image-feature-extraction", model_name="chanhua/autotrain-izefx-v3qh0", device=DEVICE, pool=True) | |
# 推理 | |
def infer3(url1, url2): | |
try: | |
print("进入推理") | |
print("打开图片1") | |
# image_real = Image.open(requests.get(url1, stream=True).raw).convert("RGB") | |
image_real = Image.open(url1).convert('RGB') | |
print("打开图片2") | |
# image_gen = Image.open(requests.get(url2, stream=True).raw).convert("RGB") | |
image_gen = Image.open(url2).convert('RGB') | |
print("利用模型获取图片特征向量") | |
outputs = pipe([image_real, image_gen]) | |
print(f"得到图片特征向量计算相似度: {outputs}") | |
similarity_score = cosine_similarity(torch.Tensor(outputs[0]), torch.Tensor(outputs[1]), dim=1) | |
print(f"得到图片相似度: {similarity_score}") | |
t_cpu = similarity_score.cpu() | |
# 然后提取这个值 | |
return t_cpu.item() | |
except Exception as e: | |
print(f"发生了一个错误: {e}") | |
return 0.0 | |
finally: | |
# 无论是否发生异常,都会执行此代码块 | |
print("这是finally块") | |
# 推理 | |
def infer2(url): | |
# image_real = Image.open(requests.get(img_urls[0], stream=True).raw).convert("RGB") | |
# image = Image.open(requests.get(url, stream=True).raw).convert("RGB") | |
image = Image.open(url).convert('RGB') | |
inputs = processor(image, return_tensors="pt").to(DEVICE) | |
outputs = model(**inputs) | |
return outputs.pooler_output | |
# 推理 | |
def infer1(image1, image2): | |
try: | |
embed_real = infer2(image1) | |
embed_gen = infer2(image2) | |
similarity_score = cosine_similarity(embed_real, embed_gen, dim=1) | |
print(similarity_score) | |
# 如果你想在CPU上操作这个值,你需要先将tensor移动到CPU | |
t_cpu = similarity_score.cpu() | |
# 然后提取这个值 | |
return t_cpu.item() | |
except Exception as e: | |
print(f"发生了一个错误: {e}") | |
return 0.0 | |
finally: | |
# 无论是否发生异常,都会执行此代码块 | |
print("这是finally块") | |