mp-02 commited on
Commit
d023795
·
verified ·
1 Parent(s): d0ef4b5

Update inference.py

Browse files
Files changed (1) hide show
  1. inference.py +2 -1
inference.py CHANGED
@@ -15,7 +15,8 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
15
  model.to(device)
16
 
17
 
18
- def prediction(Image.open(image)):
 
19
  boxes, words = OCR(image)
20
  encoding = processor(image, words, boxes=boxes, return_offsets_mapping=True, return_tensors="pt", truncation=True)
21
  offset_mapping = encoding.pop('offset_mapping')
 
15
  model.to(device)
16
 
17
 
18
+ def prediction(image):
19
+ image = Image.open(image)
20
  boxes, words = OCR(image)
21
  encoding = processor(image, words, boxes=boxes, return_offsets_mapping=True, return_tensors="pt", truncation=True)
22
  offset_mapping = encoding.pop('offset_mapping')