Upload 2 files

image_feature.py

@@ -49,8 +49,10 @@ 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")
+# 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("chanhua/autotrain-izefx-v3qh0")
 # model = AutoModel.from_pretrained("chanhua/autotrain-izefx-v3qh0").to(DEVICE)
 
@@ -59,7 +61,8 @@ model = AutoModel.from_pretrained("google/vit-base-patch16-224").to(DEVICE)
 
 # 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)
-pipe = pipeline(task="image-feature-extraction", model_name="google/vit-base-patch16-224", device=DEVICE, pool=True, revision="29e7a1e183")
+# pipe = pipeline(task="image-feature-extraction", model_name="google/vit-base-patch16-224", device=DEVICE, pool=True, revision="29e7a1e183")
+pipe = pipeline(task="image-feature-extraction", model_name="google/vit-base-patch16-224-in21k", device=DEVICE, pool=True)
 
 
 # 推理