image_feature.py

# 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 ViTImageProcessor, ViTModel
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("google/vit-base-patch16-224-in21k")
# model = AutoModel.from_pretrained("google/vit-base-patch16-224-in21k").to(DEVICE)
# processor = AutoImageProcessor.from_pretrained("chanhua/autotrain-izefx-v3qh0")
# model = AutoModel.from_pretrained("chanhua/autotrain-izefx-v3qh0").to(DEVICE)
# processor = ViTImageProcessor.from_pretrained('google/vit-large-patch16-224-in21k')
# model = ViTModel.from_pretrained('google/vit-large-patch16-224-in21k')
# processor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224-in21k')
# model = ViTModel.from_pretrained('google/vit-base-patch16-224-in21k')


# 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)
# pipe = pipeline(task="image-feature-extraction", model_name="google/vit-base-patch16-224", device=DEVICE, pool=True, revision="29e7a1e183")


#  推理
def infer4(url1, url2):
    try:
        pipe = pipeline(task="image-feature-extraction", model_name="google/vit-base-patch16-224-in21k", device=DEVICE, pool=True)

        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):
    processor = AutoImageProcessor.from_pretrained('google/vit-large-patch16-224-in21k')
    model = ViTModel.from_pretrained('google/vit-large-patch16-224-in21k')
    # 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(images=image, return_tensors="pt").to(DEVICE)

    outputs = model(**inputs)
    # last_hidden_states = outputs.last_hidden_state

    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块")

#  输出图片向量
def similarity_cpu(image1, image2):
    try:
        embed_real = xl_infer(image1)
        embed_gen = xl_infer(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 '异常'+ str(e)
    finally:
        # 无论是否发生异常，都会执行此代码块
        print("这是finally块")


#  推理
def xl_infer(url):
    image = url.convert('RGB')

    # processor = AutoImageProcessor.from_pretrained('facebook/dinov2-giant')
    # model = AutoModel.from_pretrained('facebook/dinov2-giant')

    processor = AutoImageProcessor.from_pretrained('google/vit-large-patch16-224-in21k')
    # model = ViTModel.from_pretrained('google/vit-large-patch16-224-in21k')
    model = AutoModel.from_pretrained('google/vit-large-patch16-224-in21k')

    inputs = processor(images=image, return_tensors="pt").to(DEVICE)
    outputs = model(**inputs)
    return outputs.pooler_output