cerbervs/doctr-api
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: use cuda

Browse Source
This commit is contained in:
Mark Bailey 2024-10-26 10:34:32 -04:00
parent fb20c58fe4
commit a98190895a
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
src/main.py
View File

@ -26,7 +26,7 @@ def root():
pretrained=True, pretrained=True,
assume_straight_pages=False, assume_straight_pages=False,
preserve_aspect_ratio=True, preserve_aspect_ratio=True,
) ).cuda()
result = {} result = {}
i = 0 i = 0
@ -67,9 +67,9 @@ def file(file: Image):
"db_resnet50", "db_resnet50",
"vitstr_base", "vitstr_base",
pretrained=True, pretrained=True,
assume_straight_pages=False, straighten_pages=True,
preserve_aspect_ratio=True, preserve_aspect_ratio=True,
) ).cuda()
doc = DocumentFile.from_images("data/" + file_name) doc = DocumentFile.from_images("data/" + file_name)
pred_res = predictor(doc) pred_res = predictor(doc)
json_res = pred_res.export() json_res = pred_res.export()