Tahir5 commited on
Commit
d47801d
·
verified ·
1 Parent(s): 70bb097

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +48 -0
app.py ADDED
@@ -0,0 +1,48 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ from transformers import TrOCRProcessor, VisionEncoderDecoderModel
3
+ from PIL import Image
4
+ import requests
5
+ import torch
6
+
7
+ def load_model():
8
+ processor = TrOCRProcessor.from_pretrained('microsoft/trocr-base-handwritten')
9
+ model = VisionEncoderDecoderModel.from_pretrained('microsoft/trocr-base-handwritten')
10
+ return processor, model
11
+
12
+ def recognize_text(image, processor, model):
13
+ pixel_values = processor(images=image, return_tensors="pt").pixel_values
14
+ generated_ids = model.generate(pixel_values)
15
+ generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
16
+ return generated_text
17
+
18
+ def main():
19
+ st.title("Handwritten Text Recognition with TrOCR")
20
+
21
+ processor, model = load_model()
22
+
23
+ uploaded_file = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
24
+ image_url = st.text_input("Or enter an image URL:", "https://fki.tic.heia-fr.ch/static/img/a01-122-02-00.jpg")
25
+
26
+ if uploaded_file is not None:
27
+ image = Image.open(uploaded_file).convert("RGB")
28
+ elif image_url:
29
+ try:
30
+ image = Image.open(requests.get(image_url, stream=True).raw).convert("RGB")
31
+ except Exception as e:
32
+ st.error(f"Error loading image: {e}")
33
+ return
34
+ else:
35
+ st.warning("Please upload an image or provide a URL.")
36
+ return
37
+
38
+ st.image(image, caption="Uploaded Image", use_column_width=True)
39
+
40
+ if st.button("Recognize Text"):
41
+ with st.spinner("Processing..."):
42
+ text = recognize_text(image, processor, model)
43
+
44
+ st.success("Recognized Text:")
45
+ st.write(text)
46
+
47
+ if __name__ == "__main__":
48
+ main()