|
|
import cv2
|
|
|
import numpy as np
|
|
|
class ball_tracker:
|
|
|
"""
|
|
|
A basic color tracker, it will look for colors in a range and
|
|
|
create an x and y offset valuefrom the midpoint
|
|
|
"""
|
|
|
|
|
|
def __init__(self, height, width, color_lower, color_upper):
|
|
|
self.color_lower = color_lower
|
|
|
self.color_upper = color_upper
|
|
|
self.midx = int(width / 2)
|
|
|
self.midy = int(height / 2)
|
|
|
self.midz = 50
|
|
|
self.xoffset = 0
|
|
|
self.yoffset = 0
|
|
|
self.zoffset = 0.0
|
|
|
self.knownWidth = 3
|
|
|
self.focalLength = 500
|
|
|
self.kernel = np.ones((5, 5), np.uint8) # 卷积核
|
|
|
|
|
|
def draw_arrows(self, frame):
|
|
|
"""Show the direction vector output in the cv2 window"""
|
|
|
#cv2.putText(frame,"Color:", (0, 35), cv2.FONT_HERSHEY_SIMPLEX, 1, 255, thickness=2)
|
|
|
cv2.arrowedLine(frame, (self.midx, self.midy),
|
|
|
(self.midx + self.xoffset, self.midy - self.yoffset),
|
|
|
(0, 0, 255), 5)
|
|
|
return frame
|
|
|
def find_circle(self,frame):
|
|
|
|
|
|
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) # 转换为灰色通道
|
|
|
hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV) # 转换为HSV空间
|
|
|
# 消除噪声
|
|
|
mask = cv2.inRange(hsv,self.color_lower, self.color_upper) # 设定掩膜取值范围
|
|
|
opening = cv2.morphologyEx(mask, cv2.MORPH_OPEN, self.kernel) # 形态学开运算
|
|
|
#bila = cv2.bilateralFilter(mask, 10, 200, 200) # 双边滤波消除噪声
|
|
|
edges = cv2.Canny(opening, 50, 100) # 边缘识别
|
|
|
circles = cv2.HoughCircles(
|
|
|
edges, cv2.HOUGH_GRADIENT, 1, 100, param1=100, param2=10, minRadius=10, maxRadius=100)
|
|
|
if circles is not None: # 如果识别出圆
|
|
|
for circle in circles[0]:
|
|
|
x = int(circle[0])
|
|
|
y = int(circle[1])
|
|
|
radius = int(circle[2])
|
|
|
center = (x, y)
|
|
|
return x,y,radius
|
|
|
else:
|
|
|
return 0,0,0
|
|
|
|
|
|
def track(self, frame):
|
|
|
x,y,radius=self.find_circle(frame)
|
|
|
center=(x,y)
|
|
|
if radius > 10:
|
|
|
z = self.distance_to_camera(radius)
|
|
|
print("distance={}".format(z))
|
|
|
# draw the circle and centroid on the frame,
|
|
|
# then update the list of tracked points
|
|
|
cv2.circle(frame, (int(x), int(y)), int(radius),
|
|
|
(0, 255, 255), 2)
|
|
|
cv2.circle(frame, center, 5, (0, 0, 255), -1)
|
|
|
|
|
|
self.xoffset = int(center[0] - self.midx)
|
|
|
self.yoffset = int(self.midy - center[1])
|
|
|
self.zoffset = z - self.midz
|
|
|
print("zdistance{}".format(self.zoffset))
|
|
|
else:
|
|
|
self.xoffset = 0
|
|
|
self.yoffset = 0
|
|
|
self.zoffset = 0.0
|
|
|
return self.xoffset, self.yoffset, self.zoffset
|
|
|
|
|
|
def distance_to_camera(self, perWidth):
|
|
|
"""
|
|
|
knownWidth：知道的目标宽度 厘米
|
|
|
focalLength：摄像头焦距
|
|
|
perWidth：检测框宽度 像素
|
|
|
|
|
|
#焦距1.98mm 球的半径，球的实际半径"""
|
|
|
return (self.knownWidth * self.focalLength) / perWidth
|
