arrow-detek / app.py
fffiloni's picture
Update app.py
9f95f84 verified
raw
history blame
1.92 kB
import gradio as gr
import cv2
import numpy as np
def preprocess(img):
# Convert to grayscale
img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
cv2.imwrite("1_gray.png", img_gray)
# Blur the image
img_blur = cv2.GaussianBlur(img_gray, (5, 5), 1)
cv2.imwrite("2_blur.png", img_blur)
# Detect edges with Canny
img_canny = cv2.Canny(img_blur, 50, 50)
cv2.imwrite("3_canny.png", img_canny)
# Dilate the edges
kernel = np.ones((3, 3))
img_dilate = cv2.dilate(img_canny, kernel, iterations=2)
cv2.imwrite("4_dilate.png", img_dilate)
# Erode the dilated edges
img_erode = cv2.erode(img_dilate, kernel, iterations=1)
cv2.imwrite("5_erode.png", img_erode)
return img_erode
def find_tip(points, convex_hull):
length = len(points)
indices = np.setdiff1d(range(length), convex_hull)
for i in range(2):
j = indices[i] + 2
if j > length - 1:
j = length - j
if np.all(points[j] == points[indices[i - 1] - 2]):
return tuple(points[j])
def infer(image_in):
img = cv2.imread(image_in)
contours, hierarchy = cv2.findContours(preprocess(img), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
for cnt in contours:
peri = cv2.arcLength(cnt, True)
approx = cv2.approxPolyDP(cnt, 0.025 * peri, True)
hull = cv2.convexHull(approx, returnPoints=False)
sides = len(hull)
if 6 > sides > 3 and sides + 2 == len(approx):
arrow_tip = find_tip(approx[:,0,:], hull.squeeze())
if arrow_tip:
cv2.drawContours(img, [cnt], -1, (0, 255, 0), 3)
cv2.circle(img, arrow_tip, 3, (0, 0, 255), cv2.FILLED)
cv2.imwrite("Image_result.png", img)
return "Image_result.png"
gr.Interface(
fn=infer,
inputs=gr.Image(
sources=["upload"],
type="filepath"
),
outputs=gr.Image()
).launch()