sitloboi2012 commited on
Commit
4799f35
1 Parent(s): 96ac1a0

update handler

Browse files
Files changed (1) hide show
  1. handler.py +3 -2
handler.py CHANGED
@@ -21,8 +21,9 @@ class EndpointHandler:
21
  self.PROCESSOR = LayoutLMv3Processor(self.FEATURE_EXTRACTOR, self.TOKENIZER)
22
  self.MODEL = LayoutLMv3ForSequenceClassification.from_pretrained("OtraBoi/document_classifier_testing").to(device)
23
 
24
- def __call__(self, data: Dict):
25
- image = Image.open(io.BytesIO(data["inputs"])).convert("RGB")
 
26
  encoding = self.PROCESSOR(image, return_tensors="pt", padding="max_length", truncation=True)
27
 
28
  for k,v in encoding.items():
 
21
  self.PROCESSOR = LayoutLMv3Processor(self.FEATURE_EXTRACTOR, self.TOKENIZER)
22
  self.MODEL = LayoutLMv3ForSequenceClassification.from_pretrained("OtraBoi/document_classifier_testing").to(device)
23
 
24
+ def __call__(self, data: Dict[str, bytes]):
25
+ #image = Image.open(io.BytesIO(data["inputs"])).convert("RGB")
26
+ image = data.pop("inputs", data)
27
  encoding = self.PROCESSOR(image, return_tensors="pt", padding="max_length", truncation=True)
28
 
29
  for k,v in encoding.items():