Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import matplotlib.pyplot as plt | |
| import cv2 | |
| import easyocr | |
| import numpy as np | |
| from PIL import Image as PILImage | |
| from io import BytesIO | |
| import pdf2image | |
| # Function to handle file upload and OCR processing | |
| def process_image(image, language): | |
| reader = easyocr.Reader([language]) # Initialize OCR reader with chosen language | |
| # Convert the image to OpenCV format | |
| open_cv_image = np.array(image) | |
| open_cv_image = open_cv_image[:, :, ::-1].copy() # Convert RGB to BGR for OpenCV | |
| # Apply the scan effect | |
| processed_image = scan_effect(open_cv_image) | |
| # Perform OCR | |
| output = reader.readtext(processed_image) | |
| return processed_image, output | |
| # Function for image scan effect | |
| def scan_effect(img): | |
| blackPoint = 66 | |
| whitePoint = 130 | |
| image = highPassFilter(img, kSize=51) | |
| image_white = whitePointSelect(image, whitePoint) | |
| img_black = blackPointSelect(image_white, blackPoint) | |
| image = blackPointSelect(img, blackPoint) | |
| white = whitePointSelect(image, whitePoint) | |
| img_black = blackAndWhite(white) | |
| return img_black | |
| # Helper functions for scan effect | |
| def map(x, in_min, in_max, out_min, out_max): | |
| return (x - in_min) * (out_max - out_min) / (in_max - in_min) + out_min | |
| def highPassFilter(img, kSize): | |
| if not kSize % 2: | |
| kSize += 1 | |
| kernel = np.ones((kSize, kSize), np.float32) / (kSize * kSize) | |
| filtered = cv2.filter2D(img, -1, kernel) | |
| filtered = img.astype('float32') - filtered.astype('float32') | |
| filtered = filtered + 127 * np.ones(img.shape, np.uint8) | |
| filtered = filtered.astype('uint8') | |
| return filtered | |
| def blackPointSelect(img, blackPoint): | |
| img = img.astype('int32') | |
| img = map(img, blackPoint, 255, 0, 255) | |
| _, img = cv2.threshold(img, 0, 255, cv2.THRESH_TOZERO) | |
| img = img.astype('uint8') | |
| return img | |
| def whitePointSelect(img, whitePoint): | |
| _, img = cv2.threshold(img, whitePoint, 255, cv2.THRESH_TRUNC) | |
| img = img.astype('int32') | |
| img = map(img, 0, whitePoint, 0, 255) | |
| img = img.astype('uint8') | |
| return img | |
| def blackAndWhite(img): | |
| lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB) | |
| (l, a, b) = cv2.split(lab) | |
| img = cv2.add(cv2.subtract(l, b), cv2.subtract(l, a)) | |
| return img | |
| # Streamlit app layout | |
| st.title("OCR and Image Processing App") | |
| # Sidebar for file upload and language selection | |
| st.sidebar.title("Upload and Settings") | |
| uploaded_file = st.sidebar.file_uploader("Upload PDF or Image", type=["pdf", "png", "jpg", "jpeg"]) | |
| language_option = st.sidebar.radio("Select OCR Language", ('ne', 'en')) # 'ne' for Nepali, 'en' for English | |
| if uploaded_file: | |
| file_type = uploaded_file.type | |
| if file_type == "application/pdf": | |
| # Convert PDF to images | |
| images = pdf2image.convert_from_bytes(uploaded_file.read()) | |
| st.image(images, caption='Uploaded PDF as Images', use_column_width=True) | |
| if st.button("Process PDF"): | |
| # Process the first page of the PDF | |
| st.write("Processing PDF...") | |
| for page_image in images: | |
| with st.spinner('Processing...'): | |
| processed_image, ocr_output = process_image(page_image, language_option) | |
| st.image(processed_image, caption="Processed Image", use_column_width=True) | |
| # Display OCR output | |
| st.write("Extracted Text:") | |
| for item in ocr_output: | |
| st.write(item[1]) | |
| # Allow download of processed image | |
| img = PILImage.fromarray(processed_image) | |
| buf = BytesIO() | |
| img.save(buf, format="PNG") | |
| byte_im = buf.getvalue() | |
| st.download_button(label="Download Processed Image", data=byte_im, file_name="processed_image.png", mime="image/png") | |
| else: | |
| # Handle image files | |
| image = PILImage.open(uploaded_file) | |
| st.image(image, caption="Uploaded Image", use_column_width=True) | |
| if st.button("Process Image"): | |
| # Process the uploaded image | |
| with st.spinner('Processing...'): | |
| processed_image, ocr_output = process_image(image, language_option) | |
| st.image(processed_image, caption="Processed Image", use_column_width=True) | |
| # Display OCR output | |
| st.write("Extracted Text:") | |
| for item in ocr_output: | |
| st.write(item[1]) | |
| # Allow download of processed image | |
| img = PILImage.fromarray(processed_image) | |
| buf = BytesIO() | |
| img.save(buf, format="PNG") | |
| byte_im = buf.getvalue() | |
| st.download_button(label="Download Processed Image", data=byte_im, file_name="processed_image.png", mime="image/png") | |