radames commited on
Commit
da570e8
1 Parent(s): 79a26b6

Update pipeline.py

Browse files
Files changed (1) hide show
  1. pipeline.py +8 -7
pipeline.py CHANGED
@@ -34,7 +34,7 @@ class PreTrainedPipeline():
34
  (0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))
35
  ])
36
 
37
- def __call__(self, inputs: Union[str, "Image.Image"]) -> List[float]:
38
  """
39
  Args:
40
  data (:obj:):
@@ -44,12 +44,13 @@ class PreTrainedPipeline():
44
  - "feature_vector": A list of floats corresponding to the image embedding.
45
  """
46
  parameters = {"mode": "image"}
47
- if isinstance(inputs, str):
48
- # decode base64 image to PIL
49
- image = Image.open(
50
- BytesIO(base64.b64decode(inputs))).convert("RGB")
51
- elif isinstance(inputs, "Image.Image"):
52
- image = Image.open(inputs).convert("RGB")
 
53
 
54
  image = self.transform(image).unsqueeze(0).to(device)
55
 
 
34
  (0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))
35
  ])
36
 
37
+ def __call__(self, inputs: "Image.Image") -> List[float]:
38
  """
39
  Args:
40
  data (:obj:):
 
44
  - "feature_vector": A list of floats corresponding to the image embedding.
45
  """
46
  parameters = {"mode": "image"}
47
+ image = inputs.convert("RGB")
48
+ # if isinstance(inputs, str):
49
+ # # decode base64 image to PIL
50
+ # image = Image.open(
51
+ # BytesIO(base64.b64decode(inputs))).convert("RGB")
52
+ # elif isinstance(inputs, "Image.Image"):
53
+ # image = Image.open(inputs).convert("RGB")
54
 
55
  image = self.transform(image).unsqueeze(0).to(device)
56