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