Merge pull request #74 from bit-bots/feature/generate_goals_in_maps

Generate maps: Goal back area and variable blur

humanoid_league_map_generator/generate_maps.py

@@ -1,19 +1,24 @@
 #!/usr/bin/env python3
 
+import argparse
+import math
+import os
+
 import cv2
 import numpy as np
 from scipy import ndimage
-import math
-import os
-import argparse
 
 # Generates .png files for localization
 # Default color scheme: black on white background
 # Scale: 1 px = 1 cm.
 
 parser = argparse.ArgumentParser(description="Generate maps for localization")
-parser.add_argument('output', help="output folder where the models should be saved")
-parser.add_argument('-p', '--package', dest="package", help="ros package where the models maps be saved")
+parser.add_argument('output',
+                    help="output folder where the models should be saved")
+parser.add_argument('-p', '--package', dest="package",
+                    help="ros package where the models maps be saved")
+parser.add_argument('-b', '--blur', type=float, default=1.0,
+                    help="amount of applied blurring (between 0 and 1)")
 args = parser.parse_args()
 
 lines = True
@@ -27,24 +32,27 @@
 tcrossings_blobs = False
 crosses_blobs = False
 
-penalty_mark = False
-center_point = False
+penalty_mark = True
+center_point = True
+goal_back = True  # Draw goal back area
 
-# 2019 field WM
+# V-HL21 Rules
+line_width = 5
 field_length = 900
 field_width = 600
+goal_depth = 60
 goal_width = 260
 goal_area_length = 100
 goal_area_width = 300
-penalty_area_length = 200
-penalty_area_width = 500
 penalty_mark_distance = 150
 center_circle_diameter = 150
 border_strip_width = 100
-line_width = 5
+penalty_area_length = 200
+penalty_area_width = 500
 
 if args.package:
     import rospkg
+
     # Path to store generated models
     path = os.path.join(rospkg.RosPack().get_path(args.package), args.output)
 else:
@@ -53,9 +61,6 @@
 if not os.path.exists(path):
     os.mkdir(path)
 
-# Invert image image to get black on white background
-invert = True
-
 # Choose mark style
 # mark_type = 'point'
 mark_type = 'cross'
@@ -64,54 +69,83 @@
 color = (255, 255, 255)  # white
 
 # Size of complete turf field (field with outside borders)
-image_size = (field_width + border_strip_width * 2, field_length + border_strip_width * 2, 3)
+image_size = (field_width + border_strip_width * 2,
+              field_length + border_strip_width * 2,
+              3)
 
 # Calculate important points on the field
 field_outline_start = (border_strip_width, border_strip_width)
-field_outline_end = (field_length + border_strip_width, field_width + border_strip_width)
-
-middle_line_start = (field_length // 2 + border_strip_width, border_strip_width)
-middle_line_end = (field_length // 2 + border_strip_width, field_width + border_strip_width)
-
-middle_point = (field_length // 2 + border_strip_width, field_width // 2 + border_strip_width)
-
-penalty_mark_left = (penalty_mark_distance + border_strip_width, field_width // 2 + border_strip_width)
-penalty_mark_right = (image_size[1] - border_strip_width - penalty_mark_distance, field_width // 2 + border_strip_width)
-
-goal_area_left_start = (border_strip_width, border_strip_width + field_width // 2 - goal_area_width // 2)
-goal_area_left_end = (
-border_strip_width + goal_area_length, field_width // 2 + border_strip_width + goal_area_width // 2)
-
-goal_area_right_start = (image_size[1] - goal_area_left_start[0], goal_area_left_start[1])
-goal_area_right_end = (image_size[1] - goal_area_left_end[0], goal_area_left_end[1])
-
-penalty_area_left_start = (border_strip_width, border_strip_width + field_width // 2 - penalty_area_width // 2)
-penalty_area_left_end = (
-border_strip_width + penalty_area_length, field_width // 2 + border_strip_width + penalty_area_width // 2)
-
-penalty_area_right_start = (image_size[1] - penalty_area_left_start[0], penalty_area_left_start[1])
-penalty_area_right_end = (image_size[1] - penalty_area_left_end[0], penalty_area_left_end[1])
-
-goalpost_left_1 = (border_strip_width, border_strip_width + field_width // 2 + goal_width // 2)
-goalpost_left_2 = (border_strip_width, border_strip_width + field_width // 2 - goal_width // 2)
+field_outline_end = (field_length + border_strip_width,
+                     field_width + border_strip_width)
+
+middle_line_start = (field_length // 2 + border_strip_width,
+                     border_strip_width)
+middle_line_end = (field_length // 2 + border_strip_width,
+                   field_width + border_strip_width)
+
+middle_point = (field_length // 2 + border_strip_width,
+                field_width // 2 + border_strip_width)
+
+penalty_mark_left = (penalty_mark_distance + border_strip_width,
+                     field_width // 2 + border_strip_width)
+penalty_mark_right = (image_size[1] - border_strip_width - penalty_mark_distance,
+                      field_width // 2 + border_strip_width)
+
+goal_area_left_start = (border_strip_width,
+                        border_strip_width + field_width // 2 - goal_area_width // 2)
+goal_area_left_end = (border_strip_width + goal_area_length,
+                      field_width // 2 + border_strip_width + goal_area_width // 2)
+
+goal_area_right_start = (image_size[1] - goal_area_left_start[0],
+                         goal_area_left_start[1])
+goal_area_right_end = (image_size[1] - goal_area_left_end[0],
+                       goal_area_left_end[1])
+
+penalty_area_left_start = (border_strip_width,
+                           border_strip_width + field_width // 2 - penalty_area_width // 2)
+penalty_area_left_end = (border_strip_width + penalty_area_length,
+                         field_width // 2 + border_strip_width + penalty_area_width // 2)
+
+penalty_area_right_start = (image_size[1] - penalty_area_left_start[0],
+                            penalty_area_left_start[1])
+penalty_area_right_end = (image_size[1] - penalty_area_left_end[0],
+                          penalty_area_left_end[1])
+
+goalpost_left_1 = (border_strip_width,
+                   border_strip_width + field_width // 2 + goal_width // 2)
+goalpost_left_2 = (border_strip_width,
+                   border_strip_width + field_width // 2 - goal_width // 2)
 
 goalpost_right_1 = (image_size[1] - goalpost_left_1[0], goalpost_left_1[1])
 goalpost_right_2 = (image_size[1] - goalpost_left_2[0], goalpost_left_2[1])
 
+goal_back_corner_left_1 = (goalpost_left_1[0] - goal_depth,
+                           goalpost_left_1[1])
+goal_back_corner_left_2 = (goalpost_left_2[0] - goal_depth,
+                           goalpost_left_2[1])
+
+goal_back_corner_right_1 = (goalpost_right_1[0] + goal_depth,
+                            goalpost_right_1[1])
+goal_back_corner_right_2 = (goalpost_right_2[0] + goal_depth,
+                            goalpost_right_2[1])
 
-def drawCross(img, point, width=5, length=10):
-    # Might need some fine tuning
-    vertical_start = (point[0] - width, point[1] - length)
-    vertical_end = (point[0] + width, point[1] + length)
-    horizontal_start = (point[0] - length, point[1] - width)
-    horizontal_end = (point[0] + length, point[1] + width)
+
+def drawCross(img, point, width=5, length=15):
+    half_width = width // 2 + width % 2
+    vertical_start = (point[0] - half_width, point[1] - length)
+    vertical_end = (point[0] + half_width, point[1] + length)
+    horizontal_start = (point[0] - length, point[1] - half_width)
+    horizontal_end = (point[0] + length, point[1] + half_width)
     img = cv2.rectangle(img, vertical_start, vertical_end, color, -1)
     img = cv2.rectangle(img, horizontal_start, horizontal_end, color, -1)
 
 
-def blurDistance(image, b=5):
+def blurDistance(image, blur_factor=args.blur):
+    if blur_factor <= 0:  # Skip blur
+        return (255 - image) // 2.55
+
     # Calc distances
-    distance_map = 255 - ndimage.morphology.distance_transform_edt(255 - image)  # todo weniger hin und her rechnen
+    distance_map = 255 - ndimage.morphology.distance_transform_edt(255 - image)
 
     # Maximum field distance
     maximum_size = math.sqrt(image.shape[0] ** 2 + image.shape[1] ** 2)
@@ -120,7 +154,7 @@ def blurDistance(image, b=5):
     distance_map = (distance_map / maximum_size)
 
     # Magic value please change it
-    beta = b
+    beta = (1 - blur_factor) * 10
 
     # Activation function
     distance_map = distance_map ** (2 * beta)
@@ -133,8 +167,12 @@ def blurDistance(image, b=5):
     return out_img
 
 
-def blurGaussian(image):
-    out_img = cv2.GaussianBlur(image, (99, 99), cv2.BORDER_DEFAULT)
+def blurGaussian(image, blur_factor=args.blur):
+    if blur_factor <= 0:  # Skip blur
+        return (255 - image) // 2.55
+
+    sigma = blur_factor * 50
+    out_img = cv2.GaussianBlur(image, (99, 99), sigma, borderType=cv2.BORDER_DEFAULT)
     out_img = cv2.normalize(out_img, None, alpha=0, beta=100, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_32F)
     out_img = 100 - out_img
     return out_img
@@ -145,7 +183,6 @@ def blurGaussian(image):
 ################################################################################
 
 if lines:
-
     # Create black image in correct size for lines
     img_lines = np.zeros(image_size, np.uint8)
 
@@ -163,16 +200,16 @@ def blurGaussian(image):
         if mark_type == 'point':
             img_lines = cv2.circle(img_lines, middle_point, line_width * 2, color, -1)
         else:
-            drawCross(img_lines, middle_point)
+            drawCross(img_lines, middle_point, line_width)
 
     # Draw penalty marks
     if penalty_mark:
         if mark_type == 'point':
             img_lines = cv2.circle(img_lines, penalty_mark_left, line_width * 2, color, -1)
             img_lines = cv2.circle(img_lines, penalty_mark_right, line_width * 2, color, -1)
         else:
-            drawCross(img_lines, penalty_mark_left)
-            drawCross(img_lines, penalty_mark_right)
+            drawCross(img_lines, penalty_mark_left, line_width)
+            drawCross(img_lines, penalty_mark_right, line_width)
 
     # Draw goal area
     img_lines = cv2.rectangle(img_lines, goal_area_left_start, goal_area_left_end, color, line_width)
@@ -182,8 +219,13 @@ def blurGaussian(image):
     img_lines = cv2.rectangle(img_lines, penalty_area_left_start, penalty_area_left_end, color, line_width)
     img_lines = cv2.rectangle(img_lines, penalty_area_right_start, penalty_area_right_end, color, line_width)
 
+    # Draw goal back area
+    if goal_back:
+        img_lines = cv2.rectangle(img_lines, goalpost_left_1, goal_back_corner_left_2, color, line_width)
+        img_lines = cv2.rectangle(img_lines, goalpost_right_1, goal_back_corner_right_2, color, line_width)
+
     # blur and write
-    cv2.imwrite(os.path.join(path, 'lines.png'), blurDistance(img_lines, 5))
+    cv2.imwrite(os.path.join(path, 'lines.png'), blurDistance(img_lines))
 
 #############################################################################
 # goalposts
@@ -215,7 +257,7 @@ def blurGaussian(image):
                                       line_width)
 
     # blur and write
-    cv2.imwrite(os.path.join(path, 'fieldboundary.png'), blurDistance(img_fieldboundary, 5))  # TODO oder gaussian?
+    cv2.imwrite(os.path.join(path, 'fieldboundary.png'), blurDistance(img_fieldboundary))  # TODO oder gaussian?
 
 ############################################################################
 # features