Scripts for creating a dataset. DVC support (#19)

* feat: added DVC support

uploaded dataset via DVC, created config files

* feat: bgr2raw converter added, dataset creation implemented

implemented and tested dataset creation (undistortion and convertion) from raw images (DVC stored)

implemented and tested naive converter from BGR to arbitrary 2x2 Bayer filter in case one needed

* fix: DVC autostaging enabled

* fix: relocated dvc files and added public DVC

* feat: implemented dataset undistortion with params from file

fix: number of images to proccess increased

* fix: deleted potentially unuseful script

* fix: dataset fail redone

fix: modified parser and fix filenames length

* Separated dataset + masks

* fix: changes in .gitignore

* fix: requested changed implemented

* fix: minor bugs fixed

* feat: separated id-tool from camera package

* fix: codestyle fix

---------

Co-authored-by: Daniel Paleyev <[email protected]>

.dvc/.gitignore

@@ -0,0 +1,3 @@
+/config.local
+/tmp
+/cache

.dvc/config

@@ -0,0 +1,4 @@
+[core]
+    remote = public
+['remote "public"']
+    url = https://the-lab-bucket.storage.yandexcloud.net/dvc/

.gitignore

@@ -0,0 +1,9 @@
+.vscode
+*.png
+*.bin
+*.db3
+
+#ROS2 builds
+packages/build
+packages/install
+packages/log

datasets/.gitignore

@@ -0,0 +1,3 @@
+/ball_dataset_side_1
+/ball_dataset_side_2
+/annotated_12_02_23

datasets/annotated_12_02_23.dvc

@@ -0,0 +1,6 @@
+outs:
+- md5: f06b595f7166d7cc9428201da84d9650.dir
+  size: 5030610519
+  nfiles: 5161
+  hash: md5
+  path: annotated_12_02_23

datasets/ball_dataset_side_1.dvc

@@ -0,0 +1,6 @@
+outs:
+- md5: 302f5078ba8e1db064d2575cd55a97a7.dir
+  size: 13308883120
+  nfiles: 9661
+  hash: md5
+  path: ball_dataset_side_1

datasets/ball_dataset_side_2.dvc

@@ -0,0 +1,6 @@
+outs:
+- md5: 59e511f3282db5b869282b1992e8fd41.dir
+  size: 6306722201
+  nfiles: 4578
+  hash: md5
+  path: ball_dataset_side_2

packages/camera/scripts/dataset_creator.py

@@ -0,0 +1,124 @@
+import cv2
+import numpy as np
+import os
+import argparse
+import yaml
+
+BGR_UNDISTORT_PATH = "bgr_image_undistort"
+BGR_DISTORT_PATH = "bgr_image_distort"
+RAW_UNDISTORT_PATH = "raw_image_undistort"
+
+
+class Intrinsics:
+    def __init__(self, camera_matrix, dist_coefs, image_size):
+        self.camera_matrix = camera_matrix
+        self.dist_coefs = dist_coefs
+        self.image_size = image_size
+        self.mapx, self.mapy = cv2.initUndistortRectifyMap(
+            self.camera_matrix,
+            self.dist_coefs,
+            None,
+            self.camera_matrix,
+            self.image_size[::-1],  # (width, height)
+            5,
+        )
+
+    def undistort_image(self, image):
+        cur_image_size = image.shape[:2]
+        assert (
+            self.image_size == cur_image_size
+        ), f"Images of different sizes were provided: {self.image_size} != {cur_image_size}"
+
+        return cv2.remap(image, self.mapx, self.mapy, cv2.INTER_NEAREST)
+
+    @staticmethod
+    def create_from_yaml(path_to_file):
+        if not os.path.isfile(path_to_file):
+            raise FileNotFoundError
+        with open(path_to_file, "r") as stream:
+            data = yaml.safe_load(stream)
+            camera_matrix = np.array(data["camera_matrix"], dtype=np.double).reshape(
+                (3, 3)
+            )
+            dist_coefs = np.array(data["distorsion_coefs"], dtype=np.double)
+            image_size = tuple(data["image_size"][::-1])  # (height, width)
+            assert len(image_size) == 2
+            return Intrinsics(camera_matrix, dist_coefs, image_size)
+
+
+def init_parser():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--source", help="Directory with raw distorted images")
+    parser.add_argument(
+        "--export", help="Directotory to create folders and store result images"
+    )
+    parser.add_argument("--store-undistort-bgr", action="store_true")
+    parser.add_argument("--store-distort-bgr", action="store_true")
+    parser.add_argument("--store-undistort-raw", action="store_true")
+    parser.add_argument("--params-path", default="")
+
+    return parser
+
+
+def main():
+    intrinsic_params = None
+    parser = init_parser()
+    args = parser.parse_args()
+    if not any([args.store_distort_bgr,
+                args.store_undistort_bgr, args.store_undistort_raw]):  # fmt:skip
+        print("At least one convertation flag must be specified")
+
+    if not os.path.isdir(args.export):
+        print("Invalid export directory provided:", args.export)
+        return
+    if not os.path.isdir(args.source):
+        print("Invalid source directory provided:", args.source)
+        return
+
+    # load params if undistortion is enabled
+    if args.store_undistort_bgr or args.store_undistort_raw:
+        intrinsic_params = Intrinsics.create_from_yaml(args.params_path)
+
+    # create directories if required
+    if args.store_undistort_bgr and \
+        not os.path.isdir(os.path.join(args.export, BGR_UNDISTORT_PATH)):  # fmt:skip
+        os.mkdir(os.path.join(args.export, BGR_UNDISTORT_PATH))
+    if args.store_distort_bgr and \
+        not os.path.isdir(os.path.join(args.export, BGR_DISTORT_PATH)):  # fmt:skip
+        os.mkdir(os.path.join(args.export, BGR_DISTORT_PATH))
+    if args.store_undistort_raw and \
+        not os.path.isdir(os.path.join(args.export, RAW_UNDISTORT_PATH)):  # fmt:skip
+        os.mkdir(os.path.join(args.export, RAW_UNDISTORT_PATH))
+
+    # proccess all images from source directory
+    for filename in os.listdir(args.source):
+        path_to_file = os.path.join(args.source, filename)
+        image = cv2.imread(path_to_file)
+
+        if args.store_undistort_raw:
+            undistorted_raw_img = intrinsic_params.undistort_image(image)
+            path_to_save = os.path.join(args.export, RAW_UNDISTORT_PATH, filename)
+            success = cv2.imwrite(path_to_save, undistorted_raw_img)
+            if not success:
+                print("ERROR", path_to_save, sep="\t")
+
+        if args.store_distort_bgr or args.store_undistort_bgr:
+            bgr_image = cv2.cvtColor(image[:, :, 0], cv2.COLOR_BayerRGGB2BGR)
+            if args.store_distort_bgr:
+                path_to_save = os.path.join(args.export, BGR_DISTORT_PATH, filename)
+                success = cv2.imwrite(path_to_save, bgr_image)
+                if not success:
+                    print("ERROR", path_to_save, sep="\t")
+            if args.store_undistort_bgr:
+                path_to_save = os.path.join(args.export, BGR_UNDISTORT_PATH, filename)
+                undistort_bgr_image = intrinsic_params.undistort_image(bgr_image)
+                success = cv2.imwrite(path_to_save, undistort_bgr_image)
+                if not success:
+                    print("ERROR", path_to_save, sep="\t")
+
+        print("OK", path_to_file, sep="\t")
+
+
+if __name__ == "__main__":
+    main()

packages/camera/scripts/rosbag_parser/rosbag_parser.py → packages/camera/scripts/rosbag_parser.py

@@ -102,7 +102,9 @@ def main():
         msg = deserialize_message(data, get_message(type_by_topic[topic]))
         image = bridge.compressed_imgmsg_to_cv2(msg)
 
-        filename = f"{msg.header.stamp.sec}{msg.header.stamp.nanosec}.png"
+        filename = (
+            f"{msg.header.stamp.sec}{msg.header.stamp.nanosec}".ljust(19, "0") + ".png"
+        )
         joint_path = os.path.join(save_folder, topic_to_dir(topic), filename)
         success = cv2.imwrite(joint_path, image)
         print(f"{'OK' if success else 'ERROR'}   {joint_path}")

packages/camera/scripts/camera_identifier/CMakeLists.txt → packages/camera_identifier/CMakeLists.txt

@@ -6,9 +6,12 @@ add_compile_options(-Wall -Wextra -Wpedantic -Wno-dev)
 add_library(mvsdk SHARED IMPORTED)
 set_target_properties(mvsdk PROPERTIES IMPORTED_LOCATION "/lib/libMVSDK.so")
 
-add_executable(camera_identifier src/identifier.cpp ../../src/camera_status.cpp)
+add_executable(camera_identifier src/identifier.cpp ../camera/src/camera_status.cpp)
 target_include_directories(camera_identifier PUBLIC
-  "$<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/../../include>"
+  "$<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/../camera/include>"
   "/usr/include")
 
 target_link_libraries(camera_identifier mvsdk)
+
+install(TARGETS camera_identifier
+   DESTINATION lib/${PROJECT_NAME})

requirements.txt

@@ -0,0 +1,3 @@
+numpy==1.24.4
+opencv-python==4.8.0.76
+PyYAML==5.3.1