176 lines
6.2 KiB
Python
176 lines
6.2 KiB
Python
from enum import Enum
|
|
import numpy as np
|
|
# 根据标签修改
|
|
# class tlabel(Enum):
|
|
# BBLOCK = 5 # 蓝色方块
|
|
# RBLOCK = 2 # 红色方块
|
|
# HOSPITAL = 3 # 医院
|
|
# BBALL = 4 # 蓝球
|
|
# YBALL = 5 # 黄球
|
|
# TOWER = 6 # 通信塔
|
|
# RBALL = 7 # 红球
|
|
# BASKET = 8 # 球筐
|
|
# MARKL = 9 # 指向标
|
|
# MARKR = 10 # 指向标
|
|
# SPILLAR = 11 # 小柱体 (红色)
|
|
# MPILLAR = 12 # 中柱体 (蓝色)
|
|
# LPILLAR = 13 # 大柱体 (红色)
|
|
# SIGN = 14 # 文字标牌
|
|
# TARGET = 15 # 目标靶
|
|
# SHELTER = 16 # 停车区
|
|
# BASE = 17 # 基地
|
|
class tlabel(Enum):
|
|
TOWER = 0
|
|
SIGN = 1
|
|
SHELTER = 2
|
|
HOSPITAL = 3
|
|
BASKET = 4
|
|
BASE = 5
|
|
YBALL = 6
|
|
SPILLER = 7
|
|
RMARK = 8
|
|
RBLOCK = 9
|
|
RBALL = 10
|
|
MPILLER = 11
|
|
LPILLER = 12
|
|
LMARK = 13
|
|
BBLOCK = 14
|
|
BBALL = 15
|
|
test_resp = {
|
|
'code': 0,
|
|
'data': np.array([
|
|
[4., 0.97192055, 26.64415, 228.26755, 170.16872, 357.6216],
|
|
[4., 0.97049206, 474.0152, 251.2854, 612.91644, 381.6831],
|
|
[5., 0.972649, 250.84174, 238.43622, 378.115, 367.34906]
|
|
])
|
|
}
|
|
test1_resp = {
|
|
'code': 0,
|
|
'data': np.array([])
|
|
}
|
|
class label_filter:
|
|
def __init__(self, socket, threshold=0.6):
|
|
self.num = 0
|
|
self.pos = []
|
|
self.socket = socket
|
|
self.threshold = threshold
|
|
|
|
self.img_size = (320, 240)
|
|
def get_resp(self):
|
|
self.socket.send_string('')
|
|
response = self.socket.recv_pyobj()
|
|
return response
|
|
def switch_camera(self,camera_id):
|
|
if camera_id == 1 or camera_id == 2:
|
|
self.socket.send_string(f'{camera_id}')
|
|
response = self.socket.recv_pyobj()
|
|
return response
|
|
def filter_box(self,data):
|
|
if len(data) > 0:
|
|
expect_boxes = (data[:, 1] > self.threshold) & (data[:, 0] > -1)
|
|
np_boxes = data[expect_boxes, :]
|
|
results = [
|
|
[
|
|
item[0], # 'label':
|
|
item[1], # 'score':
|
|
item[2], # 'xmin':
|
|
item[3], # 'ymin':
|
|
item[4], # 'xmax':
|
|
item[5] # 'ymax':
|
|
]
|
|
for item in np_boxes
|
|
]
|
|
if len(results) > 0:
|
|
return True, np.array(results)
|
|
return False, None
|
|
def get(self, tlabel):
|
|
# 循环查找匹配的标签值
|
|
# 返回对应标签的个数，以及坐标列表
|
|
# TODO self.filter_box none judge
|
|
response = self.get_resp()
|
|
if response['code'] == 0:
|
|
ret, results = self.filter_box(response['data'])
|
|
if ret:
|
|
expect_boxes = (results[:, 0] == tlabel.value)
|
|
boxes = results[expect_boxes, :]
|
|
self.num = len(boxes)
|
|
self.pos = boxes[:, 2:] # [[x1 y1 x2 y2]]
|
|
return self.num, self.pos
|
|
return 0, []
|
|
def find(self, tlabel):
|
|
# 遍历返回的列表，有对应标签则返回 True
|
|
response = self.get_resp()
|
|
if response['code'] == 0:
|
|
ret, results = self.filter_box(response['data'])
|
|
if ret:
|
|
expect_boxes = (results[:, 0] == tlabel.value)
|
|
boxes = results[expect_boxes, :]
|
|
if len(boxes) != 0:
|
|
return True
|
|
return False
|
|
def aim_left(self, tlabel):
|
|
# 如果标签存在，则返回列表中位置最靠左的目标框和中心的偏移值
|
|
response = self.get_resp()
|
|
if response['code'] == 0:
|
|
ret, results = self.filter_box(response['data'])
|
|
if ret:
|
|
expect_boxes = (results[:, 0] == tlabel.value)
|
|
boxes = results[expect_boxes, :]
|
|
if len(boxes) == 0:
|
|
return (False, )
|
|
xmin_values = boxes[:, 2] # xmin
|
|
xmin_index = np.argmin(xmin_values)
|
|
error = (boxes[xmin_index][4] + boxes[xmin_index][2] - self.img_size[0]) / 2
|
|
return (True, error)
|
|
return (False, )
|
|
def aim_right(self, tlabel):
|
|
# 如果标签存在，则返回列表中位置最靠右的目标框和中心的偏移值
|
|
response = self.get_resp()
|
|
if response['code'] == 0:
|
|
ret, results = self.filter_box(response['data'])
|
|
if ret:
|
|
expect_boxes = (results[:, 0] == tlabel.value)
|
|
boxes = results[expect_boxes, :]
|
|
if len(boxes) == 0:
|
|
return (False, )
|
|
xmax_values = boxes[:, 4] # xmax
|
|
xmax_index = np.argmax(xmax_values)
|
|
error = (boxes[xmax_index][4] + boxes[xmax_index][2] - self.img_size[0]) / 2
|
|
return (True, error)
|
|
return (False, )
|
|
def aim_near(self, tlabel):
|
|
# 如果标签存在，则返回列表中位置最近的目标框和中心的偏移值
|
|
response = self.get_resp()
|
|
if response['code'] == 0:
|
|
ret, results = self.filter_box(response['data'])
|
|
if ret:
|
|
expect_boxes = (results[:, 0] == tlabel.value)
|
|
boxes = results[expect_boxes, :]
|
|
if len(boxes) == 0:
|
|
return (False, 0)
|
|
center_x_values = np.abs(boxes[:, 2] + boxes[:, 4] - self.img_size[0])
|
|
center_x_index = np.argmin(center_x_values)
|
|
error = (boxes[center_x_index][4] + boxes[center_x_index][2] - self.img_size[0]) / 2
|
|
return (True, error+15)
|
|
return (False, 0)
|
|
|
|
# class Calibrate:
|
|
# def __init__(self,by_cmd):
|
|
# # 车控制对象初始化
|
|
# self.by_cmd = by_cmd
|
|
# def aim(self,error):
|
|
# self.by_cmd.send_distance_x(error,)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
obj = label_filter(None)
|
|
# results = obj.filter_box(resp['data'])
|
|
# expect_boxes = (results[:, 0] == tlabel.SPILLAR.value)
|
|
# np_boxes = results[expect_boxes, :]
|
|
# print(np_boxes[:, 2:])
|
|
# print(len(np_boxes))
|
|
print(obj.find(tlabel.BBALL))
|
|
print(obj.aim_left(tlabel.BBALL))
|
|
print(obj.aim_right(tlabel.BBALL))
|
|
print(obj.aim_near(tlabel.BBALL))
|
|
print(obj.get(tlabel.HOSPITAL)) |