Merge pull request #107 from lbr-stack/feature/roboreg-37/io_refactor

IO refactor 37

Author: mhubii

README.md

@@ -183,7 +183,7 @@ rr-cam-swarm \
     --xacro-path urdf/med7/med7.xacro \
     --root-link-name lbr_link_0 \
     --end-link-name lbr_link_7 \
-    --target-reduction 0.95 \
+    --target-reduction 0.8 \
     --scale 0.1 \
     --n-samples 1 \
     --camera-info-file test/assets/lbr_med7/zed2i/left_camera_info.yaml \

cli/rr_cam_swarm.py

@@ -249,7 +249,6 @@ def main() -> None:
     mask_files = np.array(mask_files)[random_indices].tolist()
     joint_states_files = np.array(joint_states_files)[random_indices].tolist()
     images, joint_states, masks = parse_mono_data(
-        path=args.path,
         image_files=image_files,
         mask_files=mask_files,
         joint_states_files=joint_states_files,

cli/rr_hydra.py

@@ -155,7 +155,6 @@ def main():
     mask_files = find_files(args.path, args.mask_pattern)
     depth_files = find_files(args.path, args.depth_pattern)
     joint_states, masks, depths = parse_hydra_data(
-        path=args.path,
         joint_states_files=joint_states_files,
         mask_files=mask_files,
         depth_files=depth_files,

cli/rr_mono_dr.py

@@ -153,7 +153,6 @@ def main() -> None:
     joint_states_files = find_files(args.path, args.joint_states_pattern)
     mask_files = find_files(args.path, args.mask_pattern)
     images, joint_states, masks = parse_mono_data(
-        path=args.path,
         image_files=image_files,
         joint_states_files=joint_states_files,
         mask_files=mask_files,

cli/rr_render.py

@@ -1,6 +1,6 @@
 import argparse
 import os
-import pathlib
+from pathlib import Path
 
 import cv2
 import numpy as np
@@ -142,7 +142,7 @@ def main():
         num_workers=args.num_workers,
     )
 
-    output_path = pathlib.Path(args.output_path)
+    output_path = Path(args.output_path)
     if not output_path.exists():
         output_path.mkdir(parents=True)
 
@@ -162,13 +162,12 @@ def main():
         images = images.numpy()
         renders = (renders * 255.0).squeeze(-1).cpu().numpy().astype(np.uint8)
         for render, image, image_file in zip(renders, images, image_files):
-            image_stem = pathlib.Path(image_file).stem
-            image_suffix = pathlib.Path(image_file).suffix
+            image_file = Path(image_file)
+            output_file = (
+                output_path / f"overlay_render_{image_file.stem + image_file.suffix}"
+            )
             cv2.imwrite(
-                os.path.join(
-                    str(output_path.absolute()),
-                    f"overlay_render_{image_stem + image_suffix}",
-                ),
+                output_file,
                 overlay_mask(image, render, args.color, scale=1.0),
             )
 

cli/rr_sam2.py

@@ -1,6 +1,4 @@
 import argparse
-import os
-import pathlib
 
 import cv2
 import numpy as np
@@ -48,8 +46,7 @@ def args_factory() -> argparse.Namespace:
 
 def main():
     args = args_factory()
-    path = pathlib.Path(args.path)
-    image_names = find_files(path.absolute(), args.pattern)
+    image_files = find_files(args.path, args.pattern)
 
     # detect
     detector = OpenCVDetector(
@@ -60,23 +57,21 @@ def main():
     # segment
     segmentor = Sam2Segmentor(model_id=args.model_id, device=args.device)
 
-    for image_name in progress.track(image_names, description="Generating masks..."):
-        image_stem = pathlib.Path(image_name).stem
-        image_suffix = pathlib.Path(image_name).suffix
-        img = cv2.imread(os.path.join(path.absolute(), image_name))
+    for image_file in progress.track(image_files, description="Generating masks..."):
+        img = cv2.imread(image_file)
         annotations = False
         if args.pre_annotated:
             try:
                 samples, labels = detector.read(
-                    path=os.path.join(path.absolute(), f"{image_stem}_samples.csv")
+                    path=image_file.parent / f"{image_file.stem}_samples.csv"
                 )
                 annotations = True
             except FileNotFoundError:
                 pass
         if not annotations:
             samples, labels = detector.detect(img)
             detector.write(
-                path=os.path.join(path.absolute(), f"{image_stem}_samples.csv"),
+                path=image_file.parent / f"{image_file.stem}_samples.csv",
                 samples=samples,
                 labels=labels,
             )
@@ -86,15 +81,9 @@ def main():
         overlay = overlay_mask(img, mask, mode="g", scale=1.0)
 
         # write probability and mask
-        probability_path = os.path.join(
-            path.absolute(), f"probability_sam2_{image_stem + image_suffix}"
-        )
-        mask_path = os.path.join(
-            path.absolute(), f"mask_sam2_{image_stem + image_suffix}"
-        )
-        overlay_path = os.path.join(
-            path.absolute(), f"overlay_sam2_{image_stem + image_suffix}"
-        )
+        probability_path = image_file.parent / f"probability_sam2_{image_file.name}"
+        mask_path = image_file.parent / f"mask_sam2_{image_file.name}"
+        overlay_path = image_file.parent / f"overlay_sam2_{image_file.name}"
         cv2.imwrite(probability_path, (probability * 255.0).astype(np.uint8))
         cv2.imwrite(mask_path, mask)
         cv2.imwrite(overlay_path, overlay)

cli/rr_stereo_dr.py

@@ -186,7 +186,6 @@ def main() -> None:
     right_mask_files = find_files(args.path, args.right_mask_pattern)
     left_images, right_images, joint_states, left_masks, right_masks = (
         parse_stereo_data(
-            path=args.path,
             left_image_files=left_image_files,
             right_image_files=right_image_files,
             joint_states_files=joint_states_files,

roboreg/io/__init__.py

@@ -0,0 +1,3 @@
+from .datasets import *
+from .filesystem import *
+from .parsers import *

roboreg/io/datasets.py

@@ -0,0 +1,56 @@
+from pathlib import Path
+from typing import Tuple, Union
+
+import cv2
+import numpy as np
+import rich
+from torch.utils.data import Dataset
+
+from .filesystem import find_files
+
+
+class MonocularDataset(Dataset):
+    def __init__(
+        self,
+        images_path: Union[Path, str],
+        image_pattern: str,
+        joint_states_path: Union[Path, str],
+        joint_states_pattern: str,
+    ):
+        self._image_files = find_files(images_path, image_pattern)
+        self._joint_states_files = find_files(joint_states_path, joint_states_pattern)
+
+        rich.print("Found the following files:")
+        rich.print(f"Images: {[f.name for f in self._image_files]}")
+        rich.print(f"Joint states: {[f.name for f in self._joint_states_files]}")
+
+        if len(self._image_files) != len(self._joint_states_files):
+            raise ValueError(
+                f"Number of images '{len(self._image_files)}' and joint states '{len(self._joint_states_files)}' do not match."
+            )
+
+        if len(self._image_files) == 0:
+            raise ValueError("No images found.")
+
+        if len(self._joint_states_files) == 0:
+            raise ValueError("No joint states found.")
+
+        for image_file, joint_states_file in zip(
+            self._image_files, self._joint_states_files
+        ):
+            image_index = image_file.stem.split("_")[-1]
+            joint_states_index = joint_states_file.stem.split("_")[-1]
+            if image_index != joint_states_index:
+                raise ValueError(
+                    f"Image file index '{image_file.name}' and joint states file index '{joint_states_file.name}' do not match."
+                )
+
+    def __len__(self):
+        return len(self._image_files)
+
+    def __getitem__(self, idx: int) -> Tuple[np.ndarray, np.ndarray, str]:
+        image_file = self._image_files[idx]
+        joint_states_file = self._joint_states_files[idx]
+        image = cv2.imread(image_file)
+        joint_states = np.load(joint_states_file)
+        return image, joint_states, image_file.name

roboreg/io/filesystem.py

@@ -0,0 +1,17 @@
+from pathlib import Path
+from typing import List, Union
+
+
+def find_files(path: Union[Path, str], pattern: str = "image_*.png") -> List[Path]:
+    """Find files in a directory.
+
+    Args:
+        path (Union[Path, str]): Path to the directory.
+        pattern (str): Pattern to match. Must include '_{number}.ext' format.
+
+    Returns:
+        List[Path]: Sorted file paths.
+    """
+    path = Path(path)
+    file_paths = list(path.glob(pattern))
+    return sorted(file_paths, key=lambda x: int(x.stem.split("_")[-1]))