Skip to content
Permalink
Browse files

Added row configuration and calculation of rotation (from 0-100%)

Updated Environment drawing to use Pillow for rendering (now environment view requires the Roboto font)
  • Loading branch information...
andreasschmidtjensen committed Mar 15, 2019
1 parent 726e929 commit 1b1f723695221e5afc19f41fd1b431a500447bc4
Showing with 101 additions and 17 deletions.
  1. +1 −1 tablesoccer/Ball.py
  2. +5 −4 tablesoccer/Detector.py
  3. +76 −6 tablesoccer/Players.py
  4. +19 −6 tablesoccer/SoccerField.py
@@ -7,7 +7,7 @@

class Ball:
def __init__(self):
self.position = Position(10)
self.position = Position(5)
self.history = []
self.direction = ""

@@ -77,6 +77,7 @@ def calculate_field(self, frame):
raw_image = detection_result["image"]
detection_map = self.create_detection_map(detection_result["detections"])

center = None
players = None

if 'field_center' in detection_map:
@@ -100,15 +101,15 @@ def calculate_field(self, frame):

if players is None:
# initialize players with the raw calculation of board location
players = Players(raw_top_left[0], raw_top_right[0], rows=4)
players = Players(raw_top_left[0], raw_top_right[0], row_config=(3, 3, 3, 3))

players.update(detection_map.get('player'))

if players is None or players.get_row(2) is None or len(players.get_row(2)) < 2:
if players is None or players.get_row(2) is None or len(players.get_row(2).get_players()) < 2:
# detection was not successful
return False

row = players.get_row(2)
row = players.get_row(2).get_players()

# reset image
self.calc_image = np.zeros((raw_image.shape[0], raw_image.shape[1], 3), np.uint8)
@@ -137,6 +138,6 @@ def detect(self, frame):
self.ball.update(detection_map.get('ball'))

if self.players is None:
self.players = Players(0, self.raw_image.shape[1], rows=4)
self.players = Players(0, self.raw_image.shape[1], row_config=(3, 3, 3, 3))

self.players.update(detection_map.get('player'))
@@ -1,15 +1,23 @@
import numpy as np

ROTATION_MIN = 0.5
ROTATION_MAX = 2.0


class Players:
def __init__(self, field_x_start, field_x_end, rows=4):
self.rows = rows
def __init__(self, field_x_start, field_x_end, row_config=(3, 3, 3, 3)):
self.rows = len(row_config)
self.field_x_start = field_x_start
self.field_x_end = field_x_end

self.row_size = (self.field_x_end - self.field_x_start) / rows
self.row_size = (self.field_x_end - self.field_x_start) / self.rows

self.players = [[] for _ in range(self.rows)]
self.players = [Row(i, num_players) for i, num_players in enumerate(row_config)]

def update(self, detections):
self.players = [[] for _ in range(self.rows)]
for r in self.players:
r.reset()

if detections is None:
return

@@ -20,7 +28,7 @@ def update(self, detections):

# put in players array
if self.rows > row:
self.players[row].append(d[2])
self.players[row].add_player(d[2])
else:
print("DETECTING PLAYER OUTSIDE ROW: %s" % row)

@@ -29,3 +37,65 @@ def get_row(self, row):
return self.players[row]
else:
return None


class Row:
"""
This class represents a row. Since players can't switch place in a row,
we create an array that holds the players. Each player is defined by its position,
so the top player is the first in the array and so on.
When re-detecting the field, we recreate the row, but since the position of the
players can't change, we will get the same configuration again with updated positions.
"""
def __init__(self, row_number, num_players):
self.row_number = row_number
self.num_players = num_players

self.players = []
self.rotation = 0.0

self.x_coordinate = None

def reset(self):
self.players = []
self.rotation = 0.0

def add_player(self, player):
"""
Adds player to the row. This function ensures that the players array is always
sorted by y-position of the players starting from top.
:param player: (X,Y) position
"""
if len(self.players) >= self.num_players:
print("Too many players in row %s!" % self.row_number)

# add player to array
self.players.append(player)
# sort according to y-axis
self.players.sort(key=lambda x: x[1])

if len(self.players) == self.num_players:
self.calculate_rotation()
self.x_coordinate = np.mean([p[0] for p in self.players])

def get_players(self):
return self.players

def calculate_rotation(self):
"""
Since we view from above, we know a player is rotated if the width of the bounding box
is greater than the height. The rotation will go from 0% rotated (standing) to 100% rotated (lying)
We calculate it as an average of the bounding box size of all players in the row.
"""
avg_width = np.mean([p[2] for p in self.players])
avg_height = np.mean([p[3] for p in self.players])

ratio = avg_width / avg_height

self.rotation = min(1,
max(0,
(ratio - ROTATION_MIN) / (ROTATION_MAX - ROTATION_MIN)
)
)
@@ -1,4 +1,6 @@
import cv2
from PIL import ImageFont, ImageDraw, Image
import numpy as np

from tablesoccer.Ball import Ball

@@ -32,21 +34,32 @@ def draw(self, canvas):
:return:
"""
if self.center is not None:
cv2.putText(canvas, "DIR: %s" % self.ball.direction, (0, 80),
cv2.FONT_HERSHEY_COMPLEX_SMALL, 1, (120, 255, 50), 1)
font = ImageFont.truetype("Roboto-Regular.ttf", 12)

canvas = cv2.circle(canvas, tuple(self.center.astype(int)), 2, (120, 255, 255), 2)
im = Image.fromarray(canvas)
draw = ImageDraw.Draw(im)

draw.text((0, 0), "Direction: %s" % self.ball.direction, font=font)

center_bb = (self.center[0] - 2, self.center[1] - 2, self.center[0] + 2, self.center[1] + 2)
draw.ellipse(center_bb, (120, 255, 255, 255))

ball_pos = self.ball.get_position()
if ball_pos is not None:
canvas = cv2.circle(canvas, tuple(ball_pos.astype(int)), 2, (255, 120, 255), 2)
ball_bb = (ball_pos[0] - 2, ball_pos[1] - 2, ball_pos[0] + 2, ball_pos[1] + 2)
draw.ellipse(ball_bb, (255, 12, 255, 255))

players = self.players
if players is not None:
for i, row in enumerate(players.players):
for player in row:
draw.text((int(row.x_coordinate - 20), 25), "%.2f%%" % (row.rotation * 100), font=font)

for player in row.get_players():
if len(player) > 0:
x, y = player[0], player[1]
canvas = cv2.circle(canvas, (int(x), int(y)), 2, (255, 255, 120), 2)
bb = (x - 2, y - 2, x + 2, y + 2)
draw.ellipse(bb, (255, 255, 120, 255))

canvas = np.array(im)

return canvas

0 comments on commit 1b1f723

Please sign in to comment.
You can’t perform that action at this time.