feat: use cuda

2024-10-26 10:34:32 -04:00 · 2024-10-26 10:34:32 -04:00 · a98190895a
commit a98190895a
parent fb20c58fe4
1 changed files with 3 additions and 3 deletions
--- a/src/main.py
+++ b/src/main.py
@ -26,7 +26,7 @@ def root():
        pretrained=True,
        assume_straight_pages=False,
        preserve_aspect_ratio=True,
-    )
+    ).cuda()

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