PD0AUTOMATIONAL commited on
Commit
d24c2b5
1 Parent(s): 39d7e8c

Upload handler.py

Browse files
Files changed (1) hide show
  1. handler.py +50 -0
handler.py ADDED
@@ -0,0 +1,50 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Dict, List, Any
2
+ from PIL import Image
3
+ import torch
4
+ import base64
5
+ from io import BytesIO
6
+ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
7
+ from torchvision import transforms
8
+ from torchvision.transforms.functional import InterpolationMode
9
+
10
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
11
+
12
+
13
+ class EndpointHandler():
14
+ def __init__(self, path="Salesforce/blip2-opt-6.7b-coco"):
15
+ # load the tokenizer and model
16
+ tokenizer = AutoTokenizer.from_pretrained(path)
17
+ model = AutoModelForSeq2SeqLM.from_pretrained(path)
18
+
19
+ self.image_to_text_pipeline = pipeline('image-to-text', model=model, tokenizer=tokenizer)
20
+
21
+ image_size = 384
22
+ self.transform = transforms.Compose([
23
+ transforms.Resize((image_size, image_size), interpolation=InterpolationMode.BICUBIC),
24
+ transforms.ToTensor(),
25
+ transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))
26
+ ])
27
+
28
+ def __call__(self, data: Dict[str, Any]) -> Dict[str, List[Any]]:
29
+ """
30
+ data args:
31
+ inputs (:obj: `str` | `PIL.Image` | `np.array`)
32
+ kwargs
33
+ Return:
34
+ A :obj:`dict`: will be serialized and returned
35
+ """
36
+ # Extract inputs and kwargs from the data
37
+ inputs = data["inputs"]
38
+ parameters = data.pop("parameters", None)
39
+
40
+ # Decode base64 image to PIL
41
+ image = Image.open(BytesIO(base64.b64decode(inputs['image'])))
42
+ image = self.transform(image).unsqueeze(0).to(device)
43
+
44
+ # Run the model for prediction
45
+ if parameters is not None:
46
+ predictions = self.image_to_text_pipeline(image, **parameters)
47
+ else:
48
+ predictions = self.image_to_text_pipeline(image)
49
+
50
+ return predictions