Add bounding box annotations

README.md

@@ -127,3 +127,14 @@ The HDR JSON metadata file may have the following fields:
 ```
 
 </details>
+
+
+## BOUNDING BOX STUFF
+
+The current bounding box implementation in this branch does not integrate well with the segmentation side. The HDR must have blobs on the intersections in the following colours:
+
+- L: magenta [255, 0, 255]
+- T: cyan [0, 255, 255]
+- X: darker orange [255, 100, 0]
+
+The goal posts must be solid yellow [255, 255, 0] with just the posts and not the top bar or any other part of the goals.

pbr/config/scene_config.py

@@ -127,7 +127,7 @@
         "mesh_types": [".fbx", ".obj"],
         "path": path.abspath(path.join(res_path, "balls")),
         "mask": {"index": 1, "colour": (1, 0, 0, 1)},
-        "radius": 0.045,
+        "radius": 0.09,
         "standard_deviation": 0.005,  # for randomising ball position, large value means more random, zero means no random
     },
     "environment": {
@@ -152,6 +152,32 @@
         "mask": {"index": 4, "colour": (0, 1, 0, 1), "line_colour": (1, 1, 1, 1)},
     },
     "goal": {"mask": {"index": 2, "colour": (1, 1, 0, 1)}},
+    "bounding_boxes": {
+        "enabled": True,
+        "output_format": "yolo",
+        "mask": {"index": 6, "colour": (0.5, 0.5, 0.5, 1)},  # Dummy mask config for compatibility
+        "classes": {
+            "ball": 0,
+            "goal_post": 1,
+            "robot": 2,
+            "L_intersection": 3,
+            "T_intersection": 4,
+            "X_intersection": 5
+        },
+        "min_bbox_size": 8,  # minimum bounding box size in pixels
+        "max_bbox_size": 800,  # maximum bounding box size in pixels
+        "intersection_base_sizes": {
+            "L": 0.15,  # meters - base size for L intersection
+            "T": 0.20,  # meters - base size for T intersection  
+            "X": 0.25   # meters - base size for X intersection
+        },
+        "goal_post_detection": {
+            "enabled": True,
+            "width": 0.12,  # meters - actual goal post width
+            "min_height_ratio": 0.3,  # minimum visible height ratio
+            "use_segmentation": True  # extract from goal segmentation mask
+        }
+    },
 }
 
 
@@ -231,7 +257,7 @@ def configure_scene():
                 **random.choice(
                     [
                         {"type": "EQUISOLID", "focal_length": 10.5, "fov": pi},
-                        {"type": "RECTILINEAR", "fov": 0.857},
+                        {"type": "RECTILINEAR", "fov": 1.6},
                     ]
                 ),
                 "stereo_camera_distance": 0.1,

pbr/pbr.py

@@ -366,8 +366,51 @@ def main():
                     os.path.join(out_cfg.depth_dir, filename) + ".exr",
                 )
 
-        # Check that the rotation matrix of the main camera is valid
-        print(f"Rotation matrix of {cam_l.obj.name}: \n", cam_l.obj.matrix_world)
+        ##############################################
+        ##          BOUNDING BOX GENERATION         ##
+        ##############################################
+
+        annotations = []
+
+        # Ball annotations
+        ball_annotations = [util.write_annotations(ball.obj)]
+        annotations += [ann for ann in ball_annotations if ann is not None]
+
+        # Goal annotations  
+        # goal_annotations = [util.write_annotations(goal.obj, 1) for goal in goals]
+        # annotations += [ann for ann in goal_annotations if ann is not None]
+
+        # Goal post annotations (from rendered segmentation mask)
+        rendered_mask_path = os.path.join(out_cfg.mask_dir, "{}.png".format(filename))
+        goalpost_annotations = util.write_goal_post_annotations_from_mask(
+            rendered_mask_path, bpy.context.scene
+        )
+        annotations += goalpost_annotations
+
+        # Robot annotations (exclude the camera robot r0)
+        robot_annotations = [util.write_annotations(robot.obj, 2) for robot in robots[1:]]  # Skip robots[0] which is the camera robot
+        annotations += [ann for ann in robot_annotations if ann is not None]
+
+        # Misc robot annotations
+        misc_annotations = [util.write_annotations(misc_robot.obj, 2) for misc_robot in misc_robots]
+        annotations += [ann for ann in misc_annotations if ann is not None]
+
+        # Intersection annotations (from rendered segmentation mask)
+        intersection_annotations = util.write_intersection_annotations_from_mask(
+            rendered_mask_path, bpy.context.scene
+        )
+        annotations += intersection_annotations
+
+        # Write YOLO format annotations
+        if annotations:
+            os.makedirs(out_cfg.output_dir + "/annotations", exist_ok=True)
+            annotation_file = os.path.join(out_cfg.output_dir + "/annotations", f"{filename}.txt")
+            with open(annotation_file, 'w') as f:
+                for ann in annotations:
+                    f.write(f"{ann[0]} {ann[1]:.6f} {ann[2]:.6f} {ann[3]:.6f} {ann[4]:.6f}\n")
+            print(f"[INFO] Wrote {len(annotations)} annotations to {annotation_file}")
+        else:
+            print(f"[INFO] No annotations generated for frame {filename}")
 
         # Generate meta file
         with open(