88from rocker .core import RockerExtensionManager
99from rocker .core import OPERATIONS_NON_INTERACTIVE
1010
11+ from io import BytesIO
12+ from urllib .request import urlopen
13+ from zipfile import ZipFile
14+
15+
16+ def get_bop_template (modelname ):
17+ return f"https://huggingface.co/datasets/bop-benchmark/datasets/resolve/main/{ modelname } /{ modelname } "
18+
19+
20+ available_datasets = {"lm" : get_bop_template ("lm" )}
21+
22+ bop_suffixes = [
23+ "_base.zip" ,
24+ "_models.zip" ,
25+ "_test_all.zip" ,
26+ "_train_pbr.zip" ,
27+ ]
28+
29+
30+ def fetch_bop_dataset (dataset , output_path ):
31+ for suffix in bop_suffixes :
32+
33+ url = get_bop_template (dataset ) + suffix
34+ with urlopen (url ) as zipurlfile :
35+ with ZipFile (BytesIO (zipurlfile .read ())) as zfile :
36+ zfile .extractall (output_path )
37+
1138
1239def main ():
1340
14- parser = argparse .ArgumentParser (
41+ main_parser = argparse .ArgumentParser (
1542 description = "The entry point for the Bin Picking Challenge" ,
1643 formatter_class = argparse .ArgumentDefaultsHelpFormatter ,
1744 )
18- parser .add_argument ("test_image" )
19- parser .add_argument ("dataset_directory" )
20- parser .add_argument (
45+ main_parser .add_argument (
2146 "-v" , "--version" , action = "version" , version = "%(prog)s " + get_rocker_version ()
2247 )
23- parser .add_argument ("--debug-inside" , action = "store_true" )
48+
49+ sub_parsers = main_parser .add_subparsers (title = "test" , dest = "subparser_name" )
50+ test_parser = sub_parsers .add_parser ("test" )
51+
52+ test_parser .add_argument ("estimator_image" )
53+ test_parser .add_argument ("dataset" )
54+ test_parser .add_argument ("--dataset_directory" , action = "store" , default = "." )
55+ test_parser .add_argument ("--debug-inside" , action = "store_true" )
56+
57+ fetch_parser = sub_parsers .add_parser ("fetch" )
58+ fetch_parser .add_argument ("dataset" , choices = ["lm" ])
59+ fetch_parser .add_argument ("--dataset-path" , default = "." )
2460
2561 extension_manager = RockerExtensionManager ()
2662 default_args = {"cuda" : True , "network" : "host" }
27- extension_manager .extend_cli_parser (parser , default_args )
63+ # extension_manager.extend_cli_parser(test_parser , default_args)
2864
29- args = parser .parse_args ()
65+ args = main_parser .parse_args ()
3066 args_dict = vars (args )
67+ if args .subparser_name == "fetch" :
68+ print (f"Fetching dataset { args_dict ['dataset' ]} to { args_dict ['dataset_path' ]} " )
69+ fetch_bop_dataset (args_dict ["dataset" ], args_dict ["dataset_path" ])
70+ print ("Fetch complete" )
71+ return
3172
3273 # Confirm dataset directory is absolute
3374 args_dict ["dataset_directory" ] = os .path .abspath (args_dict ["dataset_directory" ])
@@ -39,10 +80,13 @@ def main():
3980 "network" : "host" ,
4081 "extension_blacklist" : {},
4182 "operating_mode" : OPERATIONS_NON_INTERACTIVE ,
42- "env" : [[f"BOP_PATH:/opt/ros/underlay/install/datasets" ]],
83+ "env" : [
84+ [f"BOP_PATH:/opt/ros/underlay/install/datasets/{ args_dict ['dataset' ]} " ],
85+ [f"DATASET_NAME:{ args_dict ['dataset' ]} " ],
86+ ],
4387 "console_output_file" : "ibpc_test_output.log" ,
4488 "volume" : [
45- [f"{ args_dict ['dataset_directory' ]} :/opt/ros/underlay/install/datasets/lm " ]
89+ [f"{ args_dict ['dataset_directory' ]} :/opt/ros/underlay/install/datasets" ]
4690 ],
4791 }
4892 print ("Buiding tester env" )
@@ -59,11 +103,13 @@ def main():
59103 "extension_blacklist" : {},
60104 "console_output_file" : "ibpc_zenoh_output.log" ,
61105 "operating_mode" : OPERATIONS_NON_INTERACTIVE ,
106+ "volume" : [],
62107 }
108+ zenoh_extensions = extension_manager .get_active_extensions (tester_args )
63109
64110 print ("Buiding zenoh env" )
65111 dig_zenoh = DockerImageGenerator (
66- tester_extensions , tester_args , "eclipse/zenoh:1.1.1"
112+ zenoh_extensions , zenoh_args , "eclipse/zenoh:1.1.1"
67113 )
68114 exit_code = dig_zenoh .build (** zenoh_args )
69115 if exit_code != 0 :
@@ -76,22 +122,14 @@ def run_instance(dig_instance, args):
76122 tester_thread = threading .Thread (target = run_instance , args = (dig_zenoh , zenoh_args ))
77123 tester_thread .start ()
78124
79- # TODO Redirect stdout
80- import time
81-
82- time .sleep (3 )
83-
84125 tester_thread = threading .Thread (
85126 target = run_instance , args = (dig_tester , tester_args )
86127 )
87128 tester_thread .start ()
88- # TODO Redirect stdout
89129
90- import time
91-
92- time .sleep (3 )
93-
94- dig = DockerImageGenerator (active_extensions , args_dict , args_dict ["test_image" ])
130+ dig = DockerImageGenerator (
131+ active_extensions , args_dict , args_dict ["estimator_image" ]
132+ )
95133
96134 exit_code = dig .build (** vars (args ))
97135 if exit_code != 0 :
@@ -104,7 +142,3 @@ def run_instance(dig_instance, args):
104142 result = dig .run (** args_dict )
105143 # TODO clean up threads here
106144 return result
107-
108-
109- if __name__ == "__main__" :
110- main ()
0 commit comments