radames commited on
Commit
ee9318a
1 Parent(s): ad1e21b
Files changed (2) hide show
  1. pipeline.py +26 -0
  2. requirements.txt +4 -0
pipeline.py ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ 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
+
9
+ class PreTrainedPipeline():
10
+ def __init__(self, path=""):
11
+
12
+ self.device = torch.device(
13
+ "cuda" if torch.cuda.is_available() else "cpu")
14
+ self.processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
15
+ self.model = SamModel.from_pretrained(
16
+ "facebook/sam-vit-base").to(self.device)
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(
24
+ inputs["pixel_values"])
25
+
26
+ return feature_vector.tolist()
requirements.txt ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ git+https://github.com/huggingface/transformers.git
2
+ torch
3
+ numpy
4
+ Pillow