github/privateGPT
1
0
Fork 1
mirror of https://github.com/imartinez/privateGPT.git synced 2025-08-08 10:57:48 +00:00
Code Issues Packages Projects Releases Wiki Activity

updated gpu for ocr

Browse Source
This commit is contained in:
quick-karsth 2024-02-27 10:30:26 +05:45
parent b9556626b1
commit a8a10e412e
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
private_gpt/components/ocr_components/table_ocr.py
View File

@ -1,14 +1,16 @@
# from paddleocr import PaddleOCR # from paddleocr import PaddleOCR
import cv2 import cv2
import torch
from doctr.models import ocr_predictor from doctr.models import ocr_predictor
from doctr.io import DocumentFile from doctr.io import DocumentFile
device = "cuda" if torch.cuda.is_available() else "cpu"
class GetOCRText: class GetOCRText:
def __init__(self) -> None: def __init__(self) -> None:
self._image = None self._image = None
# self.ocr = PaddleOCR(use_angle_cls=True, lang='en') # self.ocr = PaddleOCR(use_angle_cls=True, lang='en')
self.doctr = ocr_predictor(det_arch='db_resnet50', reco_arch='crnn_vgg16_bn', pretrained=True) self.doctr = ocr_predictor(det_arch='db_resnet50', reco_arch='crnn_vgg16_bn', pretrained=True).to(device)
def _preprocess_image(self, img): def _preprocess_image(self, img):
resized_image = cv2.resize(img, None, fx=1.6, fy=1.6, interpolation=cv2.INTER_CUBIC) resized_image = cv2.resize(img, None, fx=1.6, fy=1.6, interpolation=cv2.INTER_CUBIC)