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'.") | |