huaapp / image_feature.py
chanhua's picture
Upload 2 files
612a9ab verified
raw
history blame
4.32 kB
# 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块")