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