From 5ab4ce587da900fc1491f5509ddde1f6f7ac74ca Mon Sep 17 00:00:00 2001 From: mark Date: Sun, 15 Sep 2024 21:29:19 +0800 Subject: [PATCH] test messi.py --- Client/src/field.cpp | 35 +++-- ZBin/py_playground/rocos/algm/messi.py | 142 ++++++++++++++++++++- ZBin/py_playground/rocos/algm/test_cuda.py | 29 +++++ share/proto/zss_debug.proto | 15 ++- 4 files changed, 197 insertions(+), 24 deletions(-) create mode 100644 ZBin/py_playground/rocos/algm/test_cuda.py diff --git a/Client/src/field.cpp b/Client/src/field.cpp index 8b4aa10..e1899be 100644 --- a/Client/src/field.cpp +++ b/Client/src/field.cpp @@ -185,7 +185,7 @@ Field::Field(QQuickItem *parent) resetAfterMouseEvent(); // draw Score - score_thread = new std::thread([ = ] {receiveScore();}); + score_thread = new std::thread([&] {receiveScore();}); //record ZRecRecorder::instance()->init(); @@ -854,29 +854,38 @@ void Field::parseScores(QUdpSocket* const socket) { datagram.resize(socket->pendingDatagramSize()); socket->readDatagram(datagram.data(), datagram.size()); scores.ParseFromArray(datagram.data(), datagram.size()); - auto size = scores.points_size(); + auto heat_size = scores.heat_size(); + auto cm = QString::fromStdString(scores.cmap()); + auto shape = scores.shape(); score_mutex.lock(); - for(int i = 0; i < size; i++) { - auto score = scores.points(i); - auto c = limitRange(score.value(), 0.0f, 1.0f); + for(int i = 0; i < heat_size; i++) { + auto score = scores.heat(i); auto size_x = score.x_size(); auto size_y = score.y_size(); - auto width = score.size(); - if(size_x != size_y) { - std::cerr << "DEBUG_SCORE : not correct size : " << size_x << " " << size_y << std::endl; + auto size_v = score.value_size(); + auto size_s = score.size_size(); + if(size_x != size_y || (size_x != size_v && size_v != 1) || (size_x != size_s && size_s != 1)) { + std::cerr << "DEBUG_SCORE : not correct size : " << size_x << " " << size_y << " " << size_v << std::endl; continue; } - scorePainter.setBrush(cmap(QString::fromStdString(scores.cmap()),c)); - QVector rects; for(int k = 0; k < size_x; k++) { auto x = score.x(k); auto y = score.y(k); - rects.push_back(QRectF(::x(x-width/2), ::y(y+width/2), ::w(width), ::h(-width))); + auto v = size_v == 1 ? score.value(0) : score.value(k); + auto s = size_s == 1 ? score.size(0) : score.size(k); + scorePainter.setBrush(cmap(cm,v)); + QRectF rect = QRectF(::x(x-s/2), ::y(y+s/2), ::w(s), ::h(-s)); + switch (shape) { + case Debug_Heatmap_Shape_SQUARE: + scorePainter.drawRect(rect); + break; + case Debug_Heatmap_Shape_CIRCLE: + scorePainter.drawEllipse(rect); + break; + } } - scorePainter.drawRects(rects); } score_mutex.unlock(); - score_buffer_mutex.lock(); score_pixmap_buffer->fill(COLOR_DARKGREEN); scorebufferPainter.drawPixmap(0, 0, *score_pixmap); diff --git a/ZBin/py_playground/rocos/algm/messi.py b/ZBin/py_playground/rocos/algm/messi.py index a096d03..72f3daf 100644 --- a/ZBin/py_playground/rocos/algm/messi.py +++ b/ZBin/py_playground/rocos/algm/messi.py @@ -1,4 +1,5 @@ import numpy as np +from scipy.spatial import distance_matrix from tbkpy.socket.udp import UDPMultiCastReceiver, UDPSender from tbkpy.socket.plugins import ProtobufParser from tzcp.ssl.rocos.zss_vision_detection_pb2 import Vision_DetectionFrame @@ -8,6 +9,84 @@ HEATMAP_COLORS = ["gray", "rainbow", "jet", "PiYG", "cool", "coolwarm", "seismic", "default"] +class DEF: + HEATMAP = "rainbow" + FLX = 9000 + FLY = 6000 + PLX = 1000 + PLY = 2000 + ROBOT_RADIUS = 90 + GL = 1000 + STEP = 100 + GOAL = np.array((FLX/2, 0)) + + MAX_ACC = 4000 + MAX_VEL = 3500 + MAX_BALL_VEL = 6000 + +def get_points_and_sizes(robot): + points = np.empty((0,2)) + sizes = np.empty(0) + # resolution of heatmap + res = DEF.STEP + R = DEF.ROBOT_RADIUS + + # represent points from unimportant to important + # points in back field + res = 3*DEF.STEP + p = np.mgrid[-DEF.FLX/2-R:0:res, -DEF.FLY/2:DEF.FLY/2:res].reshape(2, -1).T + points, sizes = np.concatenate((points, p)), np.concatenate((sizes, np.ones(len(p))*res)) + # points in front field + res = DEF.STEP + p = np.mgrid[0:DEF.FLX/2:res, -DEF.FLY/2:DEF.FLY/2:res].reshape(2, -1).T + points, sizes = np.concatenate((points, p)), np.concatenate((sizes, np.ones(len(p))*res)) + # points around robot + res = DEF.STEP*0.8 + dl = DEF.FLX/10 + p = np.mgrid[-dl:dl:res, -dl:dl:res].reshape(2, -1).T + circle = np.linalg.norm(p, axis=1) < 1.0*dl + p = p[circle] + for r in robot: + # points, sizes = np.concatenate((points, p+r)), np.concatenate((sizes, np.ones(len(p))*res)) + pass + + in_their_penalty = np.logical_and(points[:,0] > DEF.FLX/2 - DEF.PLX - R, np.abs(points[:,1]) < DEF.PLY/2+R) + in_our_penalty = np.logical_and(points[:,0] < -DEF.FLX/2 + DEF.PLX + R, np.abs(points[:,1]) < DEF.PLY/2+R) + out_of_field = np.logical_or(np.abs(points[:,0]) > DEF.FLX/2-R, np.abs(points[:,1]) > DEF.FLY/2-R) + ban = np.logical_or(np.logical_or(in_their_penalty, in_our_penalty), out_of_field) + points = points[~ban] + sizes = sizes[~ban] + print("points", len(points)) + return points, sizes + +def dist(pos:np.ndarray, target:np.ndarray): + return np.linalg.norm(pos - target, axis=1) + +def norm(v:np.ndarray): + return (v-np.min(v))/(np.max(v)-np.min(v)) + +def max_run_dist(t): + h = np.minimum(t/2*DEF.MAX_ACC, DEF.MAX_VEL) + w1 = np.maximum(t - 2*h/DEF.MAX_ACC, 0) + return 0.5*h*(w1 + t) + +def calculate_interception(points, ball, enemy): + lines = points - ball + enemy_relative = enemy - ball + angles = np.arctan2(lines[:,1], lines[:,0]) + matrixs = np.array([[np.cos(angles), -np.sin(angles)], [np.sin(angles), np.cos(angles)]]).transpose(2,0,1) + enemy_rotate = np.dot(enemy_relative, matrixs) + projection_x = enemy_rotate[...,0] + projection_y = np.abs(enemy_rotate[...,1]) + dist_mx = distance_matrix(enemy_relative, lines) + dist = np.linalg.norm(lines, axis=1) + ban1 = projection_x < 0 + ban2 = projection_x > dist + need_compare = np.logical_or(ban1, ban2) + projection_y[need_compare] = dist_mx[need_compare] + value = -(1/np.clip(projection_y/max_run_dist(dist / DEF.MAX_BALL_VEL), 0.25, 1.25)) + return value.min(axis=0) + class Messi: def __init__(self): self.signal = Event() @@ -15,8 +94,8 @@ def __init__(self): self.sender = UDPSender(plugin=ProtobufParser(Debug_Heatmap)) self.heatmap_endpoint = ("127.0.0.1", 20003) self.debug_endpoint = ("127.0.0.1", 20001) - - self.heatmap_index = 0 + self.heatmap_name = DEF.HEATMAP + self.step = 1 def callback(self, recv): self.vision = recv[0] self.signal.set() @@ -27,13 +106,13 @@ def test_heatmap(self): heatmap.cmap = HEATMAP_COLORS[self.heatmap_index] self.heatmap_index = (self.heatmap_index + 1) % len(HEATMAP_COLORS) for i in range(91): - heat = Debug_Heatmap.Heat() + heat = Debug_Heatmap.HeatDiscrete() y = np.linspace(-3000, 3000, 61) heat.x.extend([-4500+i*100]*len(y)) heat.y.extend(y) heat.value = float(i/90) heat.size = 100 ## mm - heatmap.points.append(heat) + heatmap.discrete.append(heat) print("pb_size", heatmap.ByteSize()) self.sender.send(heatmap, self.heatmap_endpoint) @@ -51,13 +130,64 @@ def test_heatmap(self): self.sender.send(debug, self.debug_endpoint) self.signal.clear() + def calculate(self): + self.signal.wait() + + robot = np.array([(robot.x, robot.y) for robot in self.vision.robots_blue]) + enemy = np.array([(robot.x, robot.y) for robot in self.vision.robots_yellow]) + ball = np.array([self.vision.balls.x, self.vision.balls.y]) + + points, sizes = get_points_and_sizes(robot) + + value = np.zeros(len(points)) + + # near to goal + value += 0.7*-np.clip(dist(points, DEF.GOAL),2000, 5000) / 3000 + # near to robot + value += -0.5*np.clip(distance_matrix(points, robot).min(axis=1) / 3000, 0.3, 1.0) + # far from enemy + value += 2*(np.clip(distance_matrix(points, enemy).min(axis=1) / 3000, 0.0, 0.3)) + # dist to ball + value += 1*np.clip(dist(points, ball) / 3000, 0.0, 0.5) + # intercept by enemy + value += 0.7*calculate_interception(points, ball, enemy) + + self.send_heatmap(points, value, sizes) + self.signal.clear() + + def histogram_equalization(self, values): + hist, bins = np.histogram(values, bins=256, range=(0,1)) + cdf = hist.cumsum() + cdf = (cdf - cdf.min()) / (cdf.max() - cdf.min()) + values = np.interp(values, bins[:-1], cdf) + return values + def send_heatmap(self, points, values, size=[DEF.STEP]): + values = norm(values) + values = self.histogram_equalization(values) + heatmap = Debug_Heatmap() + heatmap.cmap = self.heatmap_name + heat = Debug_Heatmap.Heat() + heat.x.extend(points[:,0]) + heat.y.extend(points[:,1]) + heat.value.extend(values) + heat.size.extend(size) + heatmap.heat.append(heat) + # heatmap.shape = Debug_Heatmap.Shape.CIRCLE + self.sender.send(heatmap, self.heatmap_endpoint) def main(): import time messi = Messi() + def changeCMAP(): + while True: + time.sleep(1) + messi.heatmap_name = HEATMAP_COLORS[np.random.randint(0, len(HEATMAP_COLORS))] + import threading + # threading.Thread(target=changeCMAP).start() while True: - time.sleep(1) - messi.test_heatmap() + time.sleep(0.01) + messi.calculate() + # messi.test_heatmap() def get_cmap(cmap_name): import matplotlib.cm as cm diff --git a/ZBin/py_playground/rocos/algm/test_cuda.py b/ZBin/py_playground/rocos/algm/test_cuda.py new file mode 100644 index 0000000..4401da3 --- /dev/null +++ b/ZBin/py_playground/rocos/algm/test_cuda.py @@ -0,0 +1,29 @@ +from numba import cuda +import numpy as np +import math + +# @cuda.jit +# def my_kernel(io_array): +# pos = cuda.grid(1) +# if pos < io_array.size: +# io_array[pos] *= 2 # do the computation + +# # Host code +# data = np.ones(256) +# threadsperblock = 256 +# blockspergrid = math.ceil(data.shape[0] / threadsperblock) +# my_kernel[blockspergrid, threadsperblock](data) +# print(data) + +# ------------------- + +@cuda.jit +def matmul(A, B, C): + """Perform matrix multiplication of C = A * B + """ + row, col = cuda.grid(2) + if row < C.shape[0] and col < C.shape[1]: + tmp = 0. + for k in range(A.shape[1]): + tmp += A[row, k] * B[k, col] + C[row, col] = tmp \ No newline at end of file diff --git a/share/proto/zss_debug.proto b/share/proto/zss_debug.proto index f5211de..e1d39c8 100644 --- a/share/proto/zss_debug.proto +++ b/share/proto/zss_debug.proto @@ -75,11 +75,16 @@ message Debug_Msgs{ } message Debug_Heatmap{ message Heat{ - float value = 1; // 0 - 1 - repeated float x = 2; - repeated float y = 3; - float size = 4; // mm + repeated float x = 1; + repeated float y = 2; + repeated float size = 3; // only one size if all size are the same, otherwise size.size == x.size == y.size + repeated float value = 4; // only one value if all value are the same, otherwise value.size == x.size == y.size } - repeated Heat points = 1; + repeated Heat heat = 1; string cmap = 2; + enum Shape{ + SQUARE = 0; + CIRCLE = 1; + } + Shape shape = 3; } \ No newline at end of file