Upload 2 files

app.py

@@ -3,8 +3,8 @@ import image_feature as func
 
 
 def work11(image1, image2):
- # return func.infer1(image1, image2)
- return func.infer3(image1, image2)
+ return func.infer1(image1, image2)
+ # return func.infer4(image1, image2)
 
 
 # with gr.Blocks() as demo:

image_feature.py

@@ -4,7 +4,7 @@ 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
@@ -51,11 +51,14 @@ DEVICE = torch.device('cpu')
 # 第二种方式推理图片相似度
 # 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("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-base-patch16-224-in21k')
+model = ViTModel.from_pretrained('google/vit-base-patch16-224-in21k')
+
 
 # tensor([0.6061], device='cuda:0', grad_fn=<SumBackward1>)
 
@@ -66,7 +69,7 @@ pipe = pipeline(task="image-feature-extraction", model_name="google/vit-base-pat
 
 
 # 推理
-def infer3(url1, url2):
+def infer4(url1, url2):
  try:
  print("进入推理")
  print("打开图片1")
@@ -100,7 +103,7 @@ def infer3(url1, url2):
  print("这是finally块")
 
 # 推理
-def infer2(url):
+def infer3(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')
@@ -109,8 +112,21 @@ def infer2(url):
  outputs = model(**inputs)
  return outputs.pooler_output
 
-
 # 推理
+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(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)
@@ -129,3 +145,5 @@ def infer1(image1, image2):
  finally:
  # 无论是否发生异常，都会执行此代码块
  print("这是finally块")
+
+