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.infer3(image1, image2)
 
 
 # with gr.Blocks() as demo:

image_feature.py

@@ -51,6 +51,8 @@ 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("chanhua/autotrain-izefx-v3qh0")
+model = AutoModel.from_pretrained("chanhua/autotrain-izefx-v3qh0").to(DEVICE)
 
 
 # tensor([0.6061], device='cuda:0', grad_fn=<SumBackward1>)
@@ -58,14 +60,18 @@ DEVICE = torch.device('cpu')
 # 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(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(requests.get(url2, stream=True).raw).convert("RGB")
+ image_gen = Image.open(url2).convert('RGB')
 
  print("利用模型获取图片特征向量")
  outputs = pipe([image_real, image_gen])
@@ -89,31 +95,33 @@ def infer3(url1, url2):
  # 无论是否发生异常，都会执行此代码块
  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")
-# inputs = processor(image, return_tensors="pt").to(DEVICE)
-# 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(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}")
-# finally:
-# # 无论是否发生异常，都会执行此代码块
-# print("这是finally块")
+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块")