Arch10 commited on
Commit
28fd050
·
verified ·
1 Parent(s): aca710a

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +41 -0
app.py ADDED
@@ -0,0 +1,41 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Streamlit app for extracting text from an image using the General OCR Theory (GOT) 2.0 model
2
+ import streamlit as st
3
+ from transformers import AutoTokenizer, AutoModel
4
+ import torch
5
+ from PIL import Image
6
+ import requests
7
+
8
+ # Load the pre-trained GOT OCR 2.0 model and tokenizer
9
+ @st.cache_resource(show_spinner=True)
10
+ def load_model():
11
+ tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
12
+ model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, low_cpu_mem_usage=True, device_map='cuda', use_safetensors=True, pad_token_id=tokenizer.eos_token_id)
13
+ return tokenizer, model.eval().cuda()
14
+
15
+ # Streamlit interface
16
+ st.title("OCR Application using General OCR Theory (GOT) 2.0")
17
+ st.write("Upload an image to extract text using the GOT OCR 2.0 model.")
18
+
19
+ # File upload handler
20
+ uploaded_file = st.file_uploader("Choose an image...", type=["png", "jpg", "jpeg"])
21
+
22
+ if uploaded_file is not None:
23
+ # Display the uploaded image
24
+ st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)
25
+
26
+ # Load model
27
+ tokenizer, model = load_model()
28
+
29
+ # Load the image into the model
30
+ with open(uploaded_file.name, 'wb') as f:
31
+ f.write(uploaded_file.getbuffer())
32
+
33
+ image_file = uploaded_file.name
34
+
35
+ # Perform OCR
36
+ with st.spinner("Extracting text..."):
37
+ res = model.chat(tokenizer, image_file, ocr_type='ocr')
38
+
39
+ # Display the result
40
+ st.write("Extracted Text:")
41
+ st.text(res)