Nikhitha2310's picture
Update app.py
f5c0547 verified
raw
history blame
1.75 kB
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'.")