mo-ocr / utils /model.py
mnhkahn's picture
feat: 添加表格检测与OCR功能
c2e1fed
Raw
History Blame Contribute Delete
540 Bytes
import torch
from transformers import AutoModelForObjectDetection, TableTransformerForObjectDetection
def load_detection_model():
model = AutoModelForObjectDetection.from_pretrained("microsoft/table-transformer-detection", revision="no_timm")
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
return model, device
def load_structure_model(device):
model = TableTransformerForObjectDetection.from_pretrained("microsoft/table-structure-recognition-v1.1-all")
model.to(device)
return model