Upload 2 files

app.py

@@ -3,7 +3,7 @@ import image_feature as func
 
 
 def work11(image1, image2):
-    return func.xiang(image1)
+    return func.similarity_cpu(image1, image2)
     # return func.infer1(image1, image2)
     # return func.infer4(image1, image2)
 

image_feature.py

@@ -141,23 +141,17 @@ def infer1(image1, image2):
         print("这是finally块")
 
 #  输出图片向量
-def xiang(upload_image):
+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()
 
-        image = upload_image.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')
-
-        inputs = processor(images=image, return_tensors="pt").to(DEVICE)
-        outputs = model(**inputs)
-        # last_hidden_states = outputs.last_hidden_state
-
-        return outputs.last_hidden_state
-
+        # 然后提取这个值
+        return t_cpu.item()
     except Exception as e:
         print(f"发生了一个错误: {e}")
         return '异常'+ str(e)
@@ -166,3 +160,17 @@ def xiang(upload_image):
         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