Jasper Lu commited on
Commit
b34fba5
1 Parent(s): 3455ede

Use captioning

Browse files
Files changed (1) hide show
  1. handler.py +5 -2
handler.py CHANGED
@@ -5,12 +5,15 @@ from PIL import Image
5
  import pdb
6
  import requests
7
 
 
 
8
  class EndpointHandler():
9
  def __init__(self, path=""):
10
  #self.processor = AutoProcessor.from_pretrained("jasper-lu/pix2struct_embedding")
11
  #self.model = MarkupLMModel.from_pretrained("jasper-lu/pix2struct_embedding")
12
- self.processor = AutoProcessor.from_pretrained("google/pix2struct-textcaps-base")
13
- self.model = Pix2StructVisionModel.from_pretrained("google/pix2struct-textcaps-base")
 
14
 
15
  def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
16
  url = data.pop("inputs", data)
 
5
  import pdb
6
  import requests
7
 
8
+ MODEL = "google/pix2struct-screen2words-large"
9
+
10
  class EndpointHandler():
11
  def __init__(self, path=""):
12
  #self.processor = AutoProcessor.from_pretrained("jasper-lu/pix2struct_embedding")
13
  #self.model = MarkupLMModel.from_pretrained("jasper-lu/pix2struct_embedding")
14
+ self.processor = AutoProcessor.from_pretrained(MODEL)
15
+ self.processor.image_processor.is_vqa = False
16
+ self.model = Pix2StructVisionModel.from_pretrained(MODEL)
17
 
18
  def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
19
  url = data.pop("inputs", data)