You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
346 lines
12 KiB
Python
346 lines
12 KiB
Python
import socket
|
|
from math_ops.Math_Ops import Math_Ops as M
|
|
import numpy as np
|
|
|
|
class Draw():
|
|
_socket = None
|
|
|
|
def __init__(self, is_enabled:bool, unum:int, host:str, port:int) -> None:
|
|
self.enabled = is_enabled
|
|
self._is_team_right = None
|
|
self._unum = unum
|
|
self._prefix = f'?{unum}_'.encode() # temporary prefix that should never be used in normal circumstances
|
|
|
|
#Create one socket for all instances
|
|
if Draw._socket is None:
|
|
Draw._socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM )
|
|
Draw._socket.connect((host, port))
|
|
Draw.clear_all()
|
|
|
|
|
|
def set_team_side(self, is_right):
|
|
''' Called by world parser to switch side '''
|
|
'''
|
|
Generate an appropriate player ID
|
|
RoboViz has a bug/feature: we send "swap buffers for player: 'l_1' and RoboViz
|
|
will swap every buffer that contains 'l_1' in the name, including
|
|
'l_10' and 'l_11'. To avoid that, we swap the separator to 'l-10', 'l-11'
|
|
'''
|
|
self._is_team_right = is_right
|
|
self._prefix = f"{'r' if is_right else 'l'}{'_' if self._unum < 10 else '-'}{self._unum}_".encode() #e.g. b'l_5', b'l-10'
|
|
|
|
|
|
@staticmethod
|
|
def _send(msg, id, flush):
|
|
''' Private method to send message if RoboViz is accessible '''
|
|
try:
|
|
if flush:
|
|
Draw._socket.send(msg + id + b'\x00\x00\x00' + id + b'\x00')
|
|
else:
|
|
Draw._socket.send(msg + id + b'\x00')
|
|
except ConnectionRefusedError:
|
|
pass
|
|
|
|
|
|
def circle(self, pos2d, radius, thickness, color:bytes, id:str, flush=True):
|
|
'''
|
|
Draw circle
|
|
|
|
Examples
|
|
----------
|
|
Circle in 2D (z=0): circle((-1,2), 3, 2, Draw.Color.red, "my_circle")
|
|
'''
|
|
if not self.enabled: return
|
|
assert type(color)==bytes, "The RGB color must be a bytes object, e.g. red: b'\xFF\x00\x00'"
|
|
assert not np.isnan(pos2d).any(), "Argument 'pos2d' contains 'nan' values"
|
|
|
|
if self._is_team_right:
|
|
pos2d = (-pos2d[0],-pos2d[1])
|
|
|
|
msg = b'\x01\x00' + (
|
|
f'{f"{pos2d[0] :.4f}":.6s}'
|
|
f'{f"{pos2d[1] :.4f}":.6s}'
|
|
f'{f"{radius :.4f}":.6s}'
|
|
f'{f"{thickness :.4f}":.6s}').encode() + color
|
|
|
|
Draw._send(msg, self._prefix + id.encode(), flush)
|
|
|
|
|
|
def line(self, p1, p2, thickness, color:bytes, id:str, flush=True):
|
|
'''
|
|
Draw line
|
|
|
|
Examples
|
|
----------
|
|
Line in 3D: line((0,0,0), (0,0,2), 3, Draw.Color.red, "my_line")
|
|
Line in 2D (z=0): line((0,0), (0,1), 3, Draw.Color.red, "my_line")
|
|
'''
|
|
if not self.enabled: return
|
|
assert type(color)==bytes, "The RGB color must be a bytes object, e.g. red: b'\xFF\x00\x00'"
|
|
assert not np.isnan(p1).any(), "Argument 'p1' contains 'nan' values"
|
|
assert not np.isnan(p2).any(), "Argument 'p2' contains 'nan' values"
|
|
|
|
z1 = p1[2] if len(p1)==3 else 0
|
|
z2 = p2[2] if len(p2)==3 else 0
|
|
|
|
if self._is_team_right:
|
|
p1 = (-p1[0],-p1[1],p1[2]) if len(p1)==3 else (-p1[0],-p1[1])
|
|
p2 = (-p2[0],-p2[1],p2[2]) if len(p2)==3 else (-p2[0],-p2[1])
|
|
|
|
msg = b'\x01\x01' + (
|
|
f'{f"{p1[0] :.4f}":.6s}'
|
|
f'{f"{p1[1] :.4f}":.6s}'
|
|
f'{f"{z1 :.4f}":.6s}'
|
|
f'{f"{p2[0] :.4f}":.6s}'
|
|
f'{f"{p2[1] :.4f}":.6s}'
|
|
f'{f"{z2 :.4f}":.6s}'
|
|
f'{f"{thickness :.4f}":.6s}').encode() + color
|
|
|
|
Draw._send(msg, self._prefix + id.encode(), flush)
|
|
|
|
|
|
def point(self, pos, size, color:bytes, id:str, flush=True):
|
|
'''
|
|
Draw point
|
|
|
|
Examples
|
|
----------
|
|
Point in 3D: point((1,1,1), 3, Draw.Color.red, "my_point")
|
|
Point in 2D (z=0): point((1,1), 3, Draw.Color.red, "my_point")
|
|
'''
|
|
if not self.enabled: return
|
|
assert type(color)==bytes, "The RGB color must be a bytes object, e.g. red: b'\xFF\x00\x00'"
|
|
assert not np.isnan(pos).any(), "Argument 'pos' contains 'nan' values"
|
|
|
|
z = pos[2] if len(pos)==3 else 0
|
|
|
|
if self._is_team_right:
|
|
pos = (-pos[0],-pos[1],pos[2]) if len(pos)==3 else (-pos[0],-pos[1])
|
|
|
|
msg = b'\x01\x02' + (
|
|
f'{f"{pos[0] :.4f}":.6s}'
|
|
f'{f"{pos[1] :.4f}":.6s}'
|
|
f'{f"{z :.4f}":.6s}'
|
|
f'{f"{size :.4f}":.6s}').encode() + color
|
|
|
|
Draw._send(msg, self._prefix + id.encode(), flush)
|
|
|
|
|
|
def sphere(self, pos, radius, color:bytes, id:str, flush=True):
|
|
'''
|
|
Draw sphere
|
|
|
|
Examples
|
|
----------
|
|
Sphere in 3D: sphere((1,1,1), 3, Draw.Color.red, "my_sphere")
|
|
Sphere in 2D (z=0): sphere((1,1), 3, Draw.Color.red, "my_sphere")
|
|
'''
|
|
if not self.enabled: return
|
|
assert type(color)==bytes, "The RGB color must be a bytes object, e.g. red: b'\xFF\x00\x00'"
|
|
assert not np.isnan(pos).any(), "Argument 'pos' contains 'nan' values"
|
|
|
|
z = pos[2] if len(pos)==3 else 0
|
|
|
|
if self._is_team_right:
|
|
pos = (-pos[0],-pos[1],pos[2]) if len(pos)==3 else (-pos[0],-pos[1])
|
|
|
|
msg = b'\x01\x03' + (
|
|
f'{f"{pos[0] :.4f}":.6s}'
|
|
f'{f"{pos[1] :.4f}":.6s}'
|
|
f'{f"{z :.4f}":.6s}'
|
|
f'{f"{radius :.4f}":.6s}').encode() + color
|
|
|
|
Draw._send(msg, self._prefix + id.encode(), flush)
|
|
|
|
|
|
def polygon(self, vertices, color:bytes, alpha:int, id:str, flush=True):
|
|
'''
|
|
Draw polygon
|
|
|
|
Examples
|
|
----------
|
|
Polygon in 3D: polygon(((0,0,0),(1,0,0),(0,1,0)), Draw.Color.red, 255, "my_polygon")
|
|
'''
|
|
if not self.enabled: return
|
|
assert type(color)==bytes, "The RGB color must be a bytes object, e.g. red: b'\xFF\x00\x00'"
|
|
assert 0<=alpha<=255, "The alpha channel (degree of opacity) must be in range [0,255]"
|
|
|
|
if self._is_team_right:
|
|
vertices = [(-v[0],-v[1],v[2]) for v in vertices]
|
|
|
|
msg = b'\x01\x04' + bytes([len(vertices)]) + color + alpha.to_bytes(1,'big')
|
|
|
|
for v in vertices:
|
|
msg += (
|
|
f'{f"{v[0] :.4f}":.6s}'
|
|
f'{f"{v[1] :.4f}":.6s}'
|
|
f'{f"{v[2] :.4f}":.6s}').encode()
|
|
|
|
Draw._send(msg, self._prefix + id.encode(), flush)
|
|
|
|
|
|
def annotation(self, pos, text, color:bytes, id:str, flush=True):
|
|
'''
|
|
Draw annotation
|
|
|
|
Examples
|
|
----------
|
|
Annotation in 3D: annotation((1,1,1), "SOMEtext!", Draw.Color.red, "my_annotation")
|
|
Annotation in 2D (z=0): annotation((1,1), "SOMEtext!", Draw.Color.red, "my_annotation")
|
|
'''
|
|
if not self.enabled: return
|
|
if type(text) != bytes: text = str(text).encode()
|
|
assert type(color)==bytes, "The RGB color must be a bytes object, e.g. red: b'\xFF\x00\x00'"
|
|
z = pos[2] if len(pos)==3 else 0
|
|
|
|
if self._is_team_right:
|
|
pos = (-pos[0],-pos[1],pos[2]) if len(pos)==3 else (-pos[0],-pos[1])
|
|
|
|
msg = b'\x02\x00' + (
|
|
f'{f"{pos[0] :.4f}":.6s}'
|
|
f'{f"{pos[1] :.4f}":.6s}'
|
|
f'{f"{z :.4f}":.6s}').encode() + color + text + b'\x00'
|
|
|
|
Draw._send(msg, self._prefix + id.encode(), flush)
|
|
|
|
|
|
def arrow(self, p1, p2, arrowhead_size, thickness, color:bytes, id:str, flush=True):
|
|
'''
|
|
Draw arrow
|
|
|
|
Examples
|
|
----------
|
|
Arrow in 3D: arrow((0,0,0), (0,0,2), 0.1, 3, Draw.Color.red, "my_arrow")
|
|
Arrow in 2D (z=0): arrow((0,0), (0,1), 0.1, 3, Draw.Color.red, "my_arrow")
|
|
'''
|
|
if not self.enabled: return
|
|
assert type(color)==bytes, "The RGB color must be a bytes object, e.g. red: b'\xFF\x00\x00'"
|
|
|
|
# No need to invert sides, the called shapes will handle that
|
|
if len(p1)==2: p1 = M.to_3d(p1)
|
|
else: p1 = np.asarray(p1)
|
|
if len(p2)==2: p2 = M.to_3d(p2)
|
|
else: p2 = np.asarray(p2)
|
|
|
|
vec = p2-p1
|
|
vec_size = np.linalg.norm(vec)
|
|
if vec_size == 0: return #return without warning/error
|
|
if arrowhead_size > vec_size: arrowhead_size = vec_size
|
|
|
|
ground_proj_perpendicular = np.array([ vec[1], -vec[0], 0 ])
|
|
|
|
if np.all(ground_proj_perpendicular == 0): #vertical arrow
|
|
ground_proj_perpendicular = np.array([ arrowhead_size/2, 0, 0 ])
|
|
else:
|
|
ground_proj_perpendicular *= arrowhead_size/2 / np.linalg.norm(ground_proj_perpendicular)
|
|
|
|
head_start = p2 - vec * (arrowhead_size/vec_size)
|
|
head_pt1 = head_start + ground_proj_perpendicular
|
|
head_pt2 = head_start - ground_proj_perpendicular
|
|
|
|
self.line(p1,p2,thickness,color,id,False)
|
|
self.line(p2,head_pt1,thickness,color,id,False)
|
|
self.line(p2,head_pt2,thickness,color,id,flush)
|
|
|
|
|
|
def flush(self, id):
|
|
''' Flush specific drawing by ID '''
|
|
if not self.enabled: return
|
|
|
|
Draw._send(b'\x00\x00', self._prefix + id.encode(), False)
|
|
|
|
def clear(self, id):
|
|
''' Clear specific drawing by ID '''
|
|
if not self.enabled: return
|
|
|
|
Draw._send(b'\x00\x00', self._prefix + id.encode(), True) #swap buffer twice
|
|
|
|
|
|
def clear_player(self):
|
|
''' Clear all drawings made by this player '''
|
|
if not self.enabled: return
|
|
|
|
Draw._send(b'\x00\x00', self._prefix, True) #swap buffer twice
|
|
|
|
|
|
@staticmethod
|
|
def clear_all():
|
|
''' Clear all drawings of all players '''
|
|
if Draw._socket is not None:
|
|
Draw._send(b'\x00\x00\x00\x00\x00',b'',False) #swap buffer twice using no id
|
|
|
|
|
|
class Color():
|
|
'''
|
|
Based on X11 colors
|
|
The names are restructured to make better suggestions
|
|
'''
|
|
pink_violet = b'\xC7\x15\x85'
|
|
pink_hot = b'\xFF\x14\x93'
|
|
pink_violet_pale = b'\xDB\x70\x93'
|
|
pink = b'\xFF\x69\xB4'
|
|
pink_pale = b'\xFF\xB6\xC1'
|
|
|
|
red_dark = b'\x8B\x00\x00'
|
|
red = b'\xFF\x00\x00'
|
|
red_brick = b'\xB2\x22\x22'
|
|
red_crimson = b'\xDC\x14\x3C'
|
|
red_indian = b'\xCD\x5C\x5C'
|
|
red_salmon = b'\xFA\x80\x72'
|
|
|
|
orange_red = b'\xFF\x45\x00'
|
|
orange = b'\xFF\x8C\x00'
|
|
orange_ligth = b'\xFF\xA5\x00'
|
|
|
|
yellow_gold = b'\xFF\xD7\x00'
|
|
yellow = b'\xFF\xFF\x00'
|
|
yellow_light = b'\xBD\xB7\x6B'
|
|
|
|
brown_maroon =b'\x80\x00\x00'
|
|
brown_dark = b'\x8B\x45\x13'
|
|
brown = b'\xA0\x52\x2D'
|
|
brown_gold = b'\xB8\x86\x0B'
|
|
brown_light = b'\xCD\x85\x3F'
|
|
brown_pale = b'\xDE\xB8\x87'
|
|
|
|
green_dark = b'\x00\x64\x00'
|
|
green = b'\x00\x80\x00'
|
|
green_lime = b'\x32\xCD\x32'
|
|
green_light = b'\x00\xFF\x00'
|
|
green_lawn = b'\x7C\xFC\x00'
|
|
green_pale = b'\x90\xEE\x90'
|
|
|
|
cyan_dark = b'\x00\x80\x80'
|
|
cyan_medium = b'\x00\xCE\xD1'
|
|
cyan = b'\x00\xFF\xFF'
|
|
cyan_light = b'\xAF\xEE\xEE'
|
|
|
|
blue_dark = b'\x00\x00\x8B'
|
|
blue = b'\x00\x00\xFF'
|
|
blue_royal = b'\x41\x69\xE1'
|
|
blue_medium = b'\x1E\x90\xFF'
|
|
blue_light = b'\x00\xBF\xFF'
|
|
blue_pale = b'\x87\xCE\xEB'
|
|
|
|
purple_violet = b'\x94\x00\xD3'
|
|
purple_magenta = b'\xFF\x00\xFF'
|
|
purple_light = b'\xBA\x55\xD3'
|
|
purple_pale = b'\xDD\xA0\xDD'
|
|
|
|
white = b'\xFF\xFF\xFF'
|
|
gray_10 = b'\xE6\xE6\xE6'
|
|
gray_20 = b'\xCC\xCC\xCC'
|
|
gray_30 = b'\xB2\xB2\xB2'
|
|
gray_40 = b'\x99\x99\x99'
|
|
gray_50 = b'\x80\x80\x80'
|
|
gray_60 = b'\x66\x66\x66'
|
|
gray_70 = b'\x4C\x4C\x4C'
|
|
gray_80 = b'\x33\x33\x33'
|
|
gray_90 = b'\x1A\x1A\x1A'
|
|
black = b'\x00\x00\x00'
|
|
|
|
@staticmethod
|
|
def get(r,g,b):
|
|
''' Get RGB color (0-255) '''
|
|
return bytes([int(r),int(g),int(b)])
|