sooh-j commited on
Commit
4e535be
1 Parent(s): 3607fbc

Delete pipeline.py

Browse files
Files changed (1) hide show
  1. pipeline.py +0 -56
pipeline.py DELETED
@@ -1,56 +0,0 @@
1
- # https://huggingface.co/florentgbelidji/blip_image_embeddings/raw/main/pipeline.py
2
- from typing import Dict, List, Any
3
- from PIL import Image
4
- import requests
5
- import torch
6
- import base64
7
- import os
8
- from io import BytesIO
9
- from models.blip_feature_extractor import blip_feature_extractor
10
- from torchvision import transforms
11
- from torchvision.transforms.functional import InterpolationMode
12
-
13
- device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
14
-
15
- class PreTrainedPipeline():
16
- def __init__(self, path=""):
17
- # load the optimized model
18
- self.model_path = os.path.join(path,'model_large_retrieval_coco.pth')
19
- self.model = blip_feature_extractor(
20
- pretrained=self.model_path,
21
- image_size=384,
22
- vit='large',
23
- med_config=os.path.join(path, 'configs/med_config.json')
24
- )
25
- self.model.eval()
26
- self.model = self.model.to(device)
27
-
28
- image_size = 384
29
- self.transform = transforms.Compose([
30
- transforms.Resize((image_size,image_size),interpolation=InterpolationMode.BICUBIC),
31
- transforms.ToTensor(),
32
- transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))
33
- ])
34
-
35
-
36
-
37
- def __call__(self, data: Any) -> Dict[str, List[float]]:
38
- """
39
- Args:
40
- data (:obj:):
41
- includes the input data and the parameters for the inference.
42
- Return:
43
- A :obj:`dict`:. The object returned should be a dict like {"feature_vector": [0.6331314444541931,0.8802216053009033,...,-0.7866355180740356,]} containing :
44
- - "feature_vector": A list of floats corresponding to the image embedding.
45
- """
46
- inputs = data.pop("inputs", data)
47
- parameters = data.pop("parameters", {"mode": "image"})
48
-
49
- # decode base64 image to PIL
50
- image = Image.open(BytesIO(base64.b64decode(inputs['image'])))
51
- image = self.transform(image).unsqueeze(0).to(device)
52
- text=""
53
- with torch.no_grad():
54
- feature_vector = self.model(image, text, mode=parameters["mode"])[0,0].tolist()
55
- # postprocess the prediction
56
- return {"feature_vector": feature_vector}