Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import cv2 | |
| import numpy as np | |
| import tempfile | |
| from PIL import Image | |
| from ultralytics import YOLO | |
| def process_lines(image_path): | |
| thickness = 3 | |
| image = cv2.imread(image_path) | |
| result = image.copy() | |
| gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) | |
| edges = cv2.Canny(gray, 50, 150, apertureSize=3) | |
| lines = cv2.HoughLinesP(edges, 1, np.pi / 180, threshold=80, minLineLength=40, maxLineGap=40) | |
| line_mask = np.zeros_like(gray) | |
| if lines is not None: | |
| for line in lines: | |
| x1, y1, x2, y2 = line[0] | |
| cv2.line(line_mask, (x1, y1), (x2, y2), (255, 255, 255), thickness=3) | |
| return line_mask | |
| def detect_text(image_path): | |
| model = YOLO("best 3.pt") | |
| results = model.predict(image_path) | |
| annotated_image = results[0].plot() | |
| return annotated_image | |
| st.title("Line and Text Extraction") | |
| st.sidebar.header("Upload an Image") | |
| uploaded_file = st.sidebar.file_uploader("Choose an image file", type=["png", "jpg", "jpeg", "tif"]) | |
| if st.sidebar.button("Process Image"): | |
| if uploaded_file is not None: | |
| with tempfile.NamedTemporaryFile(delete=False, suffix=uploaded_file.name) as temp_file: | |
| temp_file.write(uploaded_file.read()) | |
| temp_file_path = temp_file.name | |
| line_mask = process_lines(temp_file_path) | |
| text_extracted=detect_text(temp_file_path) | |
| st.subheader("Original image") | |
| st.image(uploaded_file) | |
| st.subheader("Lines Extracted") | |
| st.image(line_mask, channels="GRAY") | |
| st.subheader("Text Detected") | |
| st.image(cv2.cvtColor(text_extracted, cv2.COLOR_BGR2RGB)) | |
| else: | |
| st.sidebar.error("Please upload an image file before clicking 'Process Image'.") | |