Spaces:
Sleeping
Sleeping
| 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, (3, 3), 1) | |
| cv2.imwrite("2_blur.png", img_blur) | |
| # Detect edges with Canny | |
| img_canny = cv2.Canny(img_blur, 100, 200) | |
| cv2.imwrite("3_canny.png", img_canny) | |
| # Dilate the edges | |
| kernel = np.ones((3, 3)) | |
| img_dilate = cv2.dilate(img_canny, kernel, iterations=1) | |
| cv2.imwrite("4_dilate.png", img_dilate) | |
| # Erode the dilated edges | |
| img_erode = cv2.erode(img_dilate, kernel, iterations=0) | |
| 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 find_tail(points, tip): | |
| # Convert the tip and points to numpy arrays | |
| tip = np.array([tip], dtype=np.float32) | |
| points = np.array(points, dtype=np.int32) | |
| # Find the index of the farthest point | |
| distances = cv2.pointPolygonTest(points, tip, True) | |
| if distances is not None: | |
| farthest_index = np.argmax(distances) | |
| farthest_point = tuple(points[farthest_index][0]) | |
| return farthest_point | |
| else: | |
| return None | |
| def infer(image_in): | |
| img = cv2.imread(image_in) | |
| contours, hierarchy = cv2.findContours(preprocess(img), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) | |
| 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()) | |
| arrow_tail = find_tail(approx[:,0,:], arrow_tip) | |
| if arrow_tip and arrow_tail: | |
| cv2.drawContours(img, [cnt], -1, (0, 255, 0), 3) | |
| cv2.circle(img, arrow_tip, 3, (0, 0, 255), cv2.FILLED) | |
| cv2.circle(img, arrow_tail, 3, (255, 0, 0), cv2.FILLED) | |
| cv2.imwrite("Image_result.png", img) | |
| return "Image_result.png", "1_gray.png", "2_blur.png", "3_canny.png", "4_dilate.png", "5_erode.png" | |
| gr.Interface( | |
| fn=infer, | |
| inputs=gr.Image( | |
| sources=["upload"], | |
| type="filepath" | |
| ), | |
| outputs=[ | |
| gr.Image(), | |
| gr.Image(), | |
| gr.Image(), | |
| gr.Image(), | |
| gr.Image(), | |
| gr.Image(), | |
| ] | |
| ).launch() |