metadata

library_name: transformers
language:
  - ko

Model Card for Model ID

Model Details

line 단위로 수식이 포함된 글자를 인식 모델입니다.
한국어 + latex 데이터셋으로 finetuning 했습니다.

Uses

Direct Use

from PIL import Image
import glob
from transformers import TrOCRProcessor, VisionEncoderDecoderModel
import torch
import IPython.display as ipd

## 이미지 준비
img_path_list = sorted(glob.glob('images/mathematical_expression_2-*.png'))
img_list = [Image.open(img_path).convert("RGB") for img_path in img_path_list]

## 모델 및 프로세서 준비
model_path = 'models/math_ocr'
processor = TrOCRProcessor.from_pretrained(model_path)
model = VisionEncoderDecoderModel.from_pretrained(model_path)
model.eval()

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

processor.feature_extractor.size = model.config.encoder.image_size

gc = model.generation_config
gc.max_length = 128
gc.early_stopping = True
gc.no_repeat_ngram_size = 3
gc.length_penalty = 2.0
gc.num_beams = 4
gc.eos_token_id = processor.tokenizer.sep_token_id

## TrOCR 추론
pixel_values = processor(img_list, return_tensors="pt").pixel_values
generated_ids = model.generate(pixel_values.to(model.device), pad_token_id=processor.tokenizer.eos_token_id)
generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)

for img,text in zip(img_list, generated_text):
    ipd.display(img)
    print(text)