tstone / app.py
Stereo0001's picture
Update app.py
7a24815 verified
'''
from cv2 import threshold
import keras_ocr
import cv2
import math
import numpy as np
import gradio as gr
def midpoint(x1, y1, x2, y2):
x_mid = int((x1 + x2)/2)
y_mid = int((y1 + y2)/2)
return (x_mid, y_mid)
def segment_img(img):
hsv=cv2.cvtColor(img,cv2.COLOR_RGB2HSV)
#mask
mask=cv2.inRange(hsv,(40,25,25),(70,255,255))
imask=mask>0
threshold=np.zeros_like(img,np.uint8)
threshold[imask]=img[imask]
return threshold
#Main function that detects text and inpaints.
#Inputs are the image path and kreas_ocr pipeline
def inpaint_text(img_path, pipeline):
# read the image
img = keras_ocr.tools.read(img_path)
#img=segment_img(img)
# Recogize text (and corresponding regions)
# Each list of predictions in prediction_groups is a list of
# (word, box) tuples.
prediction_groups = pipeline.recognize([img])
#Define the mask for inpainting
mask = np.zeros(img.shape[:2], dtype="uint8")
for box in prediction_groups[0]:
x0, y0 = box[1][0]
x1, y1 = box[1][1]
x2, y2 = box[1][2]
x3, y3 = box[1][3]
x_mid0, y_mid0 = midpoint(x1, y1, x2, y2)
x_mid1, y_mi1 = midpoint(x0, y0, x3, y3)
#For the line thickness, we will calculate the length of the line between
#the top-left corner and the bottom-left corner.
thickness = int(math.sqrt( (x2 - x1)**2 + (y2 - y1)**2 ))
#Define the line and inpaint
cv2.line(mask, (x_mid0, y_mid0), (x_mid1, y_mi1), 255,
thickness)
inpainted_img = cv2.inpaint(img, mask, 7, cv2.INPAINT_NS)
return (inpainted_img)
# keras-ocr will automatically download pretrained
# weights for the detector and recognizer.
pipeline = keras_ocr.pipeline.Pipeline()
def RemoveText(image):
img_text_removed = inpaint_text(image, pipeline)
return cv2.cvtColor(img_text_removed, cv2.COLOR_BGR2RGB)
iface = gr.Interface(fn=RemoveText,
inputs=gr.inputs.Image(label="Image to Remove Text From", type="numpy"),
outputs="image",
examples=[["1.jpg"]],
title="Remove Text for Image")
iface.launch(debug=True)
'''
import cv2
import numpy as np
import tensorflow as tf
from tensorflow import keras
from keras_ocr import pipeline # 使用tf.keras_ocr
def detect_and_draw_lines(img, rho=1, theta=np.pi/180):
"""Detect lines in the image and draw them."""
edges = cv2.Canny(img, 50, 150, apertureSize=3)
lines = cv2.HoughLines(edges, rho, theta, 100)
if lines is not None:
for line in lines:
for x1, y1, x2, y2 in line:
cv2.line(img, (x1, y1), (x2, y2), (0, 0, 255), 2)
return img
def inpaint_lines(img):
"""Inpaint the detected lines."""
mask = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
mask[mask > 0] = 255
inpainted_img = cv2.inpaint(img, mask, 3, cv2.INPAINT_TELEA)
return inpainted_img
def inpaint_text(img_path, pipeline):
img = tf.keras.preprocessing.image.load_img(img_path, color_mode='rgb', target_size=(1024, 1024))
img_array = tf.keras.preprocessing.image.img_to_array(img)
img_array = np.expand_dims(img_array, axis=0)
# Detect and draw lines
img_with_lines = detect_and_draw_lines(img_array[0])
# Inpaint the lines
inpainted_img = inpaint_lines(img_with_lines)
predictions = pipeline.recognize([inpainted_img])
mask = np.zeros(inpainted_img.shape[:2], dtype="uint8")
for box in predictions[0]:
x0, y0, x1, y1, x2, y2, x3, y3 = box[1]
cv2.line(mask, (x0, y0), (x1, y1), 255, 10)
cv2.line(mask, (x2, y2), (x3, y3), 255, 10)
inpainted_img = cv2.inpaint(inpainted_img, mask, 3, cv2.INPAINT_TELEA)
return inpainted_img
pipeline = pipeline.Pipeline()
def RemoveTextAndLines(image):
img_text_and_lines_removed = inpaint_text(image, pipeline)
return cv2.cvtColor(img_text_and_lines_removed, cv2.COLOR_BGR2RGB)
# Example usage with Gradio interface (assuming you have Gradio installed)
import gradio as gr
iface = gr.Interface(
fn=RemoveTextAndLines,
inputs=gr.inputs.Image(label="Image to Remove Text and Lines From", type="numpy"),
outputs="image",
examples=[["1.jpg"]],
title="Remove Text and Lines for Image"
)
iface.launch(debug=True)