File size: 4,599 Bytes
7eebac0 42a8761 63c0b4b 42a8761 9f95f84 97f0d9d 9f95f84 9c14df6 9f95f84 324f00f 9f95f84 870f731 7272dc2 9f95f84 7f3b84b 9f95f84 bbe7475 42a8761 8f0ccb2 68faeed 07938c8 b94470a 8f0ccb2 54b68aa 8f0ccb2 9897471 8fd94b4 2516acd 9897471 8f0ccb2 4aa12df 8f0ccb2 4aa12df bcbdd06 b94470a 68faeed b94470a cc41846 68faeed 42a8761 09c4a22 42a8761 7766475 dc75d3f e7a8bfc 68faeed 7766475 a9852e7 f43bbe1 42a8761 b7aeaf8 a105c4f dc75d3f 42a8761 9897471 42a8761 e7a8bfc 42a8761 68faeed 32998d8 ada9763 42a8761 931bb21 42a8761 3ea24a6 ba24468 7e6a83d e7a8bfc 7e6a83d 931bb21 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 |
import gradio as gr
import cv2
import numpy as np
import json
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=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 get_length(p1, p2):
line_length = ((p1[0] - p2[0]) ** 2 + (p1[1] - p2[1]) ** 2) ** 0.5
return line_length
def get_angle(tip, tail):
return (np.degrees(np.arctan2(tip[1] - tail[1], tip[0] - tail[0]))*-1)
def get_max_distance_point(cnt, tip):
max_distance = 0
max_point = None
for [x, y] in cnt:
distance = get_length((x, y), tip)
if distance > max_distance:
max_distance = distance
max_point = (x, y)
return max_point
def find_tail(points, tip):
tip = np.array(tip, dtype=np.float32)
points = np.array(points, dtype=np.int32)
# Find the tail point based on maximum distance
tail_point = get_max_distance_point(points, tip)
# Ensure that the tail point is sufficiently different from the tip
if np.linalg.norm(np.array(tail_point) - np.array(tip)) > 5:
return tail_point
return None
def draw_arrow(img_result, tip, tail, length, angle):
# Draw arrow on the blank image with inverted tip and tail
cv2.arrowedLine(img_result, tuple(tail), tuple(tip), (0, 255, 0), 3)
"""
# Add length and angle as text next to the tip point
text_length = f"Length: {length:.2f}"
text_angle = f"Angle: {angle:.2f}"
cv2.putText(img_result, text_length, (tip[0] + 10, tip[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1, cv2.LINE_AA)
cv2.putText(img_result, text_angle, (tip[0] + 10, tip[1] - 30), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1, cv2.LINE_AA)
"""
def infer(image_in):
img = cv2.imread(image_in)
contours, hierarchy = cv2.findContours(preprocess(img), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
# Create a blank image to accumulate arrows
img_result = np.zeros_like(img)
arrows_coordinates = []
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)
arrow_tail = find_tail(approx[:,0,:], arrow_tip)
if arrow_tail :
arrows_coordinates.append([list(arrow_tail), list(arrow_tip)])
cv2.circle(img, arrow_tail, 3, (255, 0, 0), cv2.FILLED)
# Calculate length and angle
arrow_length = get_length(arrow_tip, arrow_tail)
arrow_angle = get_angle(arrow_tip, arrow_tail)
# Draw arrow on the same blank image
draw_arrow(img_result, arrow_tip, arrow_tail, arrow_length, arrow_angle)
cv2.imwrite("Image_result.png", img)
cv2.imwrite("Arrows_on_same_blank.png", img_result)
print(f"arrows coordinates: {arrows_coordinates}")
result_list = json.loads(f"{arrows_coordinates}")
print(result_list)
return "Image_result.png", "Arrows_on_same_blank.png", f"{arrows_coordinates}"
gr.Interface(
fn=infer,
inputs=gr.Image(
sources=["upload"],
type="filepath"
),
outputs=[
gr.Image(label="detected arrows"),
gr.Image(label="vectors"),
gr.Textbox(label="arrows coordinates")
],
examples = ["example.png"]
).launch() |