Skip to content

Commit

Permalink
Add param file for extraction
Browse files Browse the repository at this point in the history
  • Loading branch information
RobinSchmid7 committed Sep 12, 2023
1 parent d0d119d commit 8696d56
Show file tree
Hide file tree
Showing 4 changed files with 30 additions and 5 deletions.
14 changes: 10 additions & 4 deletions scripts/dataset_generation/extract_binary_maps.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
import rosparam
from sensor_msgs.msg import Image, CameraInfo, CompressedImage
import rosbag
from wild_visual_navigation.cfg import ExtractionParams

from postprocessing_tools_ros.merging import merge_bags_single, merge_bags_all

Expand Down Expand Up @@ -64,11 +65,16 @@ def do(n, dry_run):

s = os.path.join(ROOT_DIR, d["name"])

valid_topics = ["/state_estimator/anymal_state", "/wide_angle_camera_front/img_out"]
extraction_cfg = ExtractionParams()

rosbags = ["/home/rschmid/RosBags/6_proc/images.bag",
"/home/rschmid/RosBags/6_proc/2023-03-02-11-13-08_anymal-d020-lpc_mission_0.bag",
"/home/rschmid/RosBags/6_proc/2023-03-02-11-13-08_anymal-d020-lpc_mission_1.bag"]
valid_topics = extraction_cfg.wvn_topics
rosbags = extraction_cfg.wvn_bags

# valid_topics = ["/state_estimator/anymal_state", "/wide_angle_camera_front/img_out", "/depth_camera_front_upper/point_cloud_self_filtered"]
#
# rosbags = ["/home/rschmid/RosBags/6_proc/images.bag",
# "/home/rschmid/RosBags/6_proc/2023-03-02-11-13-08_anymal-d020-lpc_mission_0.bag",
# "/home/rschmid/RosBags/6_proc/2023-03-02-11-13-08_anymal-d020-lpc_mission_1.bag"]

output_bag_wvn = s + "_wvn.bag"
output_bag_tf = s + "_tf.bag"
Expand Down
1 change: 1 addition & 0 deletions wild_visual_navigation/cfg/__init__.py
Original file line number Diff line number Diff line change
@@ -1 +1,2 @@
from .experiment_params import ExperimentParams
from .extraction_params import ExtractionParams
11 changes: 11 additions & 0 deletions wild_visual_navigation/cfg/extraction_params.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
from dataclasses import dataclass, field, asdict
from typing import Tuple, Dict, List, Optional, Any

@dataclass
class ExtractionParams:
wvn_topics: List[str] = field(default_factory=lambda: ["/state_estimator/anymal_state", "/wide_angle_camera_front/img_out", "/depth_camera_front_upper/point_cloud_self_filtered"])
wvn_bags: List[str] = field(default_factory=lambda: ["/home/rschmid/RosBags/6_proc/images.bag",
"/home/rschmid/RosBags/6_proc/2023-03-02-11-13-08_anymal-d020-lpc_mission_0.bag",
"/home/rschmid/RosBags/6_proc/2023-03-02-11-13-08_anymal-d020-lpc_mission_1.bag"])

data: ExtractionParams = ExtractionParams()
Original file line number Diff line number Diff line change
Expand Up @@ -406,12 +406,19 @@ def add_proprio_node(self, pnode: ProprioceptionNode, projection_mode: str = "im
if self._mode == WVNMode.EXTRACT_LABELS:
p = os.path.join(
self._extraction_store_folder,
"supervision_mask",
"data_mask",
str(mnode.timestamp).replace(".", "_") + ".pt",
)
store = torch.nan_to_num(mnode.supervision_mask.nanmean(axis=0)) != 0
torch.save(store, p)

if self._mode == WVNMode.EXTRACT_LABELS:
p = os.path.join(
self._extraction_store_folder,
"data_image",
str(mnode.timestamp).replace(".", "_") + ".pt",
)
torch.save(mnode.image, p)

# if self._anomaly_detection:
# # Visualize supervision mask
Expand Down

0 comments on commit 8696d56

Please sign in to comment.