Upload image_feature.py

image_feature.py

@@ -1,7 +1,7 @@
 # import numpy as np
-# import requests
+import requests
 import torch
-# from PIL import Image
+from PIL import Image
 from torch.nn.functional import cosine_similarity
 from transformers import AutoImageProcessor, AutoModel
 
@@ -54,7 +54,9 @@ model = AutoModel.from_pretrained("google/vit-base-patch16-224").to(DEVICE)
 
 
 #  推理
-def infer(image):
+def infer(img):
+    # image_real = Image.open(requests.get(img_urls[0], stream=True).raw).convert("RGB")
+    image = Image.open(requests.get(img, stream=True).raw).convert("RGB")
     inputs = processor(image, return_tensors="pt").to(DEVICE)
     outputs = model(**inputs)
     return outputs.pooler_output