Upload 2 files

app.py

@@ -3,7 +3,8 @@ import image_feature as func
 
 
 def work11(image1, image2):
-    return func.infer1(image1, image2)
+    # return func.infer1(image1, image2)
+    return func.infer3(image1, image2)
 
 
 # with gr.Blocks() as demo:
@@ -15,11 +16,21 @@ def work11(image1, image2):
 # demo.launch()
 
 # 定义你的界面
-with gr.Interface(fn=work11,
-                  inputs=[gr.Textbox(label='图片1', lines=1), gr.Textbox(label='图片2', lines=1)],  # 两个文本输入框
-                  outputs=[gr.Textbox(lines=3, label="推理结果")],  # 输出为文本
-                  title="图片相似度推理",  # 界面标题
-                  description="输入两张图片链接进行相似度推理",  # 界面描述
-                  examples=[["https://example.com", "https://google.com"],  # 示例输入
-                            ["https://github.com", "https://twitter.com"]]) as demo:  # 更多示例输入
-    demo.launch()  # 启动界面
+# with gr.Interface(fn=work11,
+#                   inputs=[gr.Textbox(label='图片1', lines=1), gr.Textbox(label='图片2', lines=1)],  # 两个文本输入框
+#                   outputs=[gr.Textbox(lines=3, label="推理结果")],  # 输出为文本
+#                   title="图片相似度推理",  # 界面标题
+#                   description="输入两张图片链接进行相似度推理",  # 界面描述
+#                   examples=[["https://example.com", "https://google.com"],  # 示例输入
+#                             ["https://github.com", "https://twitter.com"]]) as demo:  # 更多示例输入
+#     demo.launch()  # 启动界面
+
+demo = gr.Interface(title="图片相似度推理",
+                    css="",
+                    fn=work11,
+                    inputs=[gr.Image(type="filepath", label="图片1"), gr.Image(type="filepath", label="图片2")],
+                    outputs=[gr.Textbox(lines=3, label="推理结果")])
+#
+# # demo = gr.Interface(fn=work, inputs="image,text", outputs="text")
+#
+demo.launch()

image_feature.py

@@ -5,7 +5,7 @@ from PIL import Image
 from torch.nn.functional import cosine_similarity
 from transformers import AutoImageProcessor, AutoModel
 
-# from transformers import pipeline
+from transformers import pipeline
 
 # import transformers
 #
@@ -53,6 +53,34 @@ processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224")
 model = AutoModel.from_pretrained("google/vit-base-patch16-224").to(DEVICE)
 
 
+# tensor([0.6061], device='cuda:0', grad_fn=<SumBackward1>)
+
+
+#  推理
+def infer3(url1, url2):
+    try:
+        image_real = Image.open(requests.get(url1, stream=True).raw).convert("RGB")
+        image_gen = Image.open(requests.get(url2, stream=True).raw).convert("RGB")
+
+        pipe = pipeline(task="image-feature-extraction", model_name="google/vit-base-patch16-384", device=DEVICE,
+                        pool=True)
+
+        outputs = pipe([image_real, image_gen])
+
+        similarity_score = cosine_similarity(torch.Tensor(outputs[0]), torch.Tensor(outputs[1]), dim=1)
+
+        t_cpu = similarity_score.cpu()
+
+        # 然后提取这个值
+        return t_cpu.item()
+
+    except Exception as e:
+        print(f"发生了一个错误: {e}")
+    finally:
+        # 无论是否发生异常，都会执行此代码块
+        print("这是finally块")
+
+
 #  推理
 def infer2(url):
     # image_real = Image.open(requests.get(img_urls[0], stream=True).raw).convert("RGB")
@@ -80,5 +108,3 @@ def infer1(image1, image2):
     finally:
         # 无论是否发生异常，都会执行此代码块
         print("这是finally块")
-
-# tensor([0.6061], device='cuda:0', grad_fn=<SumBackward1>)