Spaces:
Sleeping
Sleeping
import gradio as gr | |
import cv2 | |
import numpy as np | |
def preprocess(img): | |
img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) | |
_, img_thresh = cv2.threshold(img_gray, 127, 255, cv2.THRESH_BINARY) | |
img_dilate = cv2.dilate(img_thresh, np.ones((3, 3)), iterations=2) | |
img_erode = cv2.erode(img_dilate, np.ones((3, 3)), iterations=1) | |
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, _ = cv2.findContours(preprocess(img), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) | |
for cnt in contours: | |
approx = cv2.approxPolyDP(cnt, 0.02 * cv2.arcLength(cnt, True), 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() |