Merge pull request #56 from bit-bots/feature/simulation-image-conversion

feat(import): correctly handle simulation imports

Author: jaagut

ddlitlab2024/dataset/cli/run.py

@@ -61,12 +61,20 @@ def main():
                         from ddlitlab2024.dataset.imports.strategies.bitbots import BitBotsImportStrategy
 
                         logger.info(f"Trying to import file '{args.file}' to database...")
+
+                        simulated = False
+                        location = "RoboCup2024"
+
+                        if "simulation" in str(args.file) or "simulated" in str(args.file):
+                            simulated = True
+                            location = "Webots Simulator"
+
                         metadata = ImportMetadata(
                             allow_public=True,
                             team_name="Bit-Bots",
                             robot_type="Wolfgang-OP",
-                            location="RoboCup2024",
-                            simulated=False,
+                            location=location,
+                            simulated=simulated,
                         )
                         image_converter = ImageConverter(MaxRateResampler(IMAGE_MAX_RESAMPLE_RATE_HZ))
                         game_state_converter = GameStateConverter(OriginalRateResampler())

ddlitlab2024/dataset/converters/image_converter.py

@@ -37,7 +37,7 @@ def convert_to_model(self, data: InputData, relative_timestamp: float, recording
         return models
 
     def _create_image(self, data, sampling_timestamp: float, recording: Recording) -> Image:
-        img_array = np.frombuffer(data.data, np.uint8).reshape((data.height, data.width, 3))
+        img_array = np.frombuffer(data.data, np.uint8).reshape((data.height, data.width, -1))
 
         will_img_be_upscaled = recording.img_width_scaling > 1.0 or recording.img_height_scaling > 1.0
         interpolation = cv2.INTER_AREA
@@ -50,9 +50,17 @@ def _create_image(self, data, sampling_timestamp: float, recording: Recording) -
                 resized_rgb_img = resized_img
             case "bgr8":
                 resized_rgb_img = cv2.cvtColor(resized_img, cv2.COLOR_BGR2RGB)
+            case "bgra8":
+                resized_rgb_img = cv2.cvtColor(resized_img, cv2.COLOR_BGRA2RGB)
             case _:
                 raise AssertionError(f"Unsupported image encoding: {data.encoding}")
 
+        assert resized_rgb_img.shape == (
+            recording.img_height,
+            recording.img_width,
+            3,
+        ), "Converted image does not have the expected dimensions"
+
         return Image(
             stamp=sampling_timestamp,
             recording=recording,