@@ -157,3 +157,154 @@ def make_aug_short_real_dataset_info(
157157json_string = json .dumps (data_list ,indent = 1 )
158158with open (save_json_path ,'w' )as json_file :
159159json_file .write (json_string )
160+
161+
162+ def oxe_dataset_info ():
163+ dataset_names = [
164+ # {
165+ # "dataset_name": f"bridge_dataset",
166+ # "wrist_image": "Normal",
167+ # "s_ratio": 1.0,
168+ # }, # zheng
169+
170+ # {
171+ # "dataset_name": f"cmu_stretch",
172+ # "wrist_image": "Normal",
173+ # "s_ratio": 1.0,
174+ # }, # zheng
175+
176+ {
177+ "dataset_name" :f"fractal20220817_data" ,
178+ "wrist_image" :"Normal" ,
179+ "s_ratio" :0.54087122203 ,
180+ },# zheng ###
181+
182+ # {
183+ # "dataset_name": f"dlr_edan_shared_control_converted_externally_to_rlds",
184+ # "wrist_image": "Normal",
185+ # "s_ratio": 1.0,
186+ # }, # zheng
187+
188+ # {
189+ # "dataset_name": f"kuka",
190+ # "wrist_image": "Normal",
191+ # "s_ratio": 0.8341046294,
192+ # }, # zheng ###
193+
194+ # {
195+ # "dataset_name": f"roboturk",
196+ # "wrist_image": "Normal",
197+ # "s_ratio": 1.0,
198+ # }, # zheng
199+
200+ # {
201+ # "dataset_name": f"ucsd_kitchen_dataset_converted_externally_to_rlds",
202+ # "wrist_image": "Normal",
203+ # "s_ratio": 1.0,
204+ # }, # zheng
205+
206+ # {
207+ # "dataset_name" : f"berkeley_autolab_ur5",
208+ # "wrist_image": "Flip vertically & horizontally",
209+ # "s_ratio": 1.0,
210+ # }, # fan,
211+
212+ # {
213+ # "dataset_name" : f"berkeley_fanuc_manipulation",
214+ # "wrist_image": "Flip vertically & horizontally",
215+ # "s_ratio": 1.0,
216+ # }, # fan
217+
218+ # {
219+ # "dataset_name" : f"jaco_play",
220+ # "wrist_image": "Flip vertically & horizontally",
221+ # "s_ratio": 1.0,
222+ # }, # fan
223+
224+ # {
225+ # "dataset_name" : f"iamlab_cmu_pickup_insert_converted_externally_to_rlds",
226+ # "wrist_image": "Normal",
227+ # "s_ratio": 1.0,
228+ # }, # zheng
229+
230+ # {
231+ # "dataset_name" : f"viola",
232+ # "wrist_image": "Flip vertically & horizontally",
233+ # "s_ratio": 2.0,
234+ # }, # fan
235+
236+ # {
237+ # "dataset_name" : f"stanford_hydra_dataset_converted_externally_to_rlds",
238+ # "wrist_image": "Flip vertically & horizontally",
239+ # "s_ratio": 2.0,
240+ # }, # fan
241+
242+ # {
243+ # "dataset_name" : f"austin_buds_dataset_converted_externally_to_rlds",
244+ # "wrist_image": "Flip vertically & horizontally",
245+ # "s_ratio": 1.0,
246+ # }, # fan
247+
248+ # {
249+ # "dataset_name" : f"utaustin_mutex",
250+ # "wrist_image": "Normal",
251+ # "s_ratio": 1.0,
252+ # }, # zheng
253+
254+ # {
255+ # "dataset_name" : f"taco_play",
256+ # "wrist_image": "Flip vertically & horizontally",
257+ # "s_ratio": 2.0,
258+ # }, # fan
259+
260+ # {
261+ # "dataset_name" : f"austin_sailor_dataset_converted_externally_to_rlds",
262+ # "wrist_image": "Flip vertically & horizontally",
263+ # "s_ratio": 1.0,
264+ # }, # fan
265+
266+ # {
267+ # "dataset_name" : f"austin_sirius_dataset_converted_externally_to_rlds",
268+ # "wrist_image": "Flip vertically & horizontally",
269+ # "s_ratio": 1.0,
270+ # }, # fan
271+
272+ # {
273+ # "dataset_name" : f"furniture_bench_dataset_converted_externally_to_rlds",
274+ # "wrist_image": "Normal",
275+ # "s_ratio": 0.1,
276+ # }, # zheng
277+ ]
278+
279+ # total_data_list = []
280+
281+ for info in tqdm (dataset_names ):
282+ dataset_name = info ["dataset_name" ]
283+ wrist_image_info = info ["wrist_image" ]
284+ s_ratio = info ["s_ratio" ]
285+ root_path = f"/xxx/preprocess/oxe/{ dataset_name }
286+ save_json_path = f"/xxx/data_info/{ dataset_name }
287+ root_path_list = glob .glob (os .path .join (root_path ,"*" ,"*" ))
288+ root_path_list .sort ()
289+ data_list = []
290+ data_list .append (info )
291+ accumulated_num_steps = 0
292+ for this_path in tqdm (root_path_list ):
293+ exp_id = this_path .split ('/' )[- 2 ]
294+ demo_id = this_path .split ('/' )[- 1 ]
295+ num_step = len (glob .glob (os .path .join (this_path ,"steps" ,"*" )))
296+ if s_ratio >= 1.0 :
297+ for _ in range (int (s_ratio )):
298+ accumulated_num_steps += num_step
299+ data_list .append ([exp_id + '/' + demo_id ,num_step ])
300+ else :
301+ this_p = np .random .random ()
302+ if this_p < s_ratio :
303+ accumulated_num_steps += num_step
304+ data_list .append ([exp_id + '/' + demo_id ,num_step ])
305+
306+ data_list [0 ]["accumulated_num_steps" ]= accumulated_num_steps
307+ json_string = json .dumps (data_list ,indent = 1 )
308+ with open (save_json_path ,'w' )as json_file :
309+ json_file .write (json_string )
310+