Divyansh12 commited on
Commit
6eb9b18
·
verified ·
1 Parent(s): 01eaeb8

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -9
app.py CHANGED
@@ -1,21 +1,16 @@
1
  import streamlit as st
2
-
3
  from PIL import Image
4
  import re
5
- from transformers import AutoModel, AutoTokenizer
6
 
7
  st.set_page_config(page_title="OCR Application", page_icon="🖼️", layout="wide")
8
  device = "cpu"
9
 
10
  @st.cache_resource
11
- #def load_model():
12
- #processor = TrOCRProcessor.from_pretrained('microsoft/trocr-base-handwritten')
13
- #model = VisionEncoderDecoderModel.from_pretrained('microsoft/trocr-base-handwritten', device_map='cpu')
14
- #@st.cache_resource
15
  def load_model():
16
- tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, device_map='cpu')
17
- model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, low_cpu_mem_usage=True, device_map='cpu', use_safetensors=True)
18
- processor=tokenizer
19
  return processor, model
20
 
21
  def extract_text(image, processor, model):
 
1
  import streamlit as st
 
2
  from PIL import Image
3
  import re
4
+ from transformers import VisionEncoderDecoderModel, AutoProcessor
5
 
6
  st.set_page_config(page_title="OCR Application", page_icon="🖼️", layout="wide")
7
  device = "cpu"
8
 
9
  @st.cache_resource
 
 
 
 
10
  def load_model():
11
+ # Load processor and model for OCR
12
+ processor = AutoProcessor.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
13
+ model = VisionEncoderDecoderModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, device_map='cpu', low_cpu_mem_usage=True)
14
  return processor, model
15
 
16
  def extract_text(image, processor, model):