radames commited on
Commit
afe4a3f
1 Parent(s): 550bf5c

Update pipeline.py

Browse files
Files changed (1) hide show
  1. pipeline.py +1 -2
pipeline.py CHANGED
@@ -2,7 +2,6 @@ from typing import List
2
  import torch
3
  from transformers import SamModel, SamProcessor
4
  from PIL import Image
5
- from io import BytesIO
6
  import numpy as np
7
 
8
 
@@ -17,7 +16,7 @@ class PreTrainedPipeline():
17
  self.model.eval()
18
  self.model = self.model.to(self.device)
19
 
20
- def __call__(self, inputs: "Image.Image") -> BytesIO:
21
  raw_image = inputs.convert("RGB")
22
  inputs = self.processor(raw_image, return_tensors="pt").to(self.device)
23
  feature_vector = self.model.get_image_embeddings(
 
2
  import torch
3
  from transformers import SamModel, SamProcessor
4
  from PIL import Image
 
5
  import numpy as np
6
 
7
 
 
16
  self.model.eval()
17
  self.model = self.model.to(self.device)
18
 
19
+ def __call__(self, inputs: "Image.Image") -> List[float]:
20
  raw_image = inputs.convert("RGB")
21
  inputs = self.processor(raw_image, return_tensors="pt").to(self.device)
22
  feature_vector = self.model.get_image_embeddings(