Spaces:
Runtime error
Runtime error
| import torch | |
| import os.path as osp | |
| import os | |
| import sys | |
| from src.dataset.dataset import SimpleIterDataset | |
| from src.utils.utils import to_filelist | |
| from pathlib import Path | |
| import pickle | |
| from src.utils.paths import get_path | |
| import argparse | |
| import numpy as np | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--input", type=str) | |
| parser.add_argument("--output", type=str) | |
| parser.add_argument("--overwrite", action="store_true") | |
| parser.add_argument("--dataset-cap", type=int, default=-1) | |
| parser.add_argument("--v2", action="store_true") # V2 means that the dataset also stores parton-level and genParticles | |
| parser.add_argument("--delphes", action="store_true") | |
| args = parser.parse_args() | |
| path = get_path(args.input, "data") | |
| def remove_from_list(lst): | |
| out = [] | |
| for item in lst: | |
| if item in ["hgcal", "data.txt", "test_file.root"]: | |
| continue | |
| out.append(item) | |
| return out | |
| def preprocess_dataset(datasets, output_path, config_file, dataset_cap): | |
| #datasets = os.listdir(path) | |
| #datasets = [os.path.join(path, x) for x in datasets] | |
| class Args: | |
| def __init__(self): | |
| self.data_train = datasets | |
| self.data_val = datasets | |
| #self.data_train = files_train | |
| self.data_config = config_file | |
| self.extra_selection = None | |
| self.train_val_split = 1.0 | |
| self.data_fraction = 1 | |
| self.file_fraction = 1 | |
| self.fetch_by_files = False | |
| self.fetch_step = 1 | |
| self.steps_per_epoch = None | |
| self.in_memory = False | |
| self.local_rank = None | |
| self.copy_inputs = False | |
| self.no_remake_weights = False | |
| self.batch_size = 10 | |
| self.num_workers = 0 | |
| self.demo = False | |
| self.laplace = False | |
| self.diffs = False | |
| self.class_edges = False | |
| args = Args() | |
| train_range = (0, args.train_val_split) | |
| train_file_dict, train_files = to_filelist(args, 'train') | |
| train_data = SimpleIterDataset(train_file_dict, args.data_config, for_training=True, | |
| extra_selection=args.extra_selection, | |
| remake_weights=True, | |
| load_range_and_fraction=(train_range, args.data_fraction), | |
| file_fraction=args.file_fraction, | |
| fetch_by_files=args.fetch_by_files, | |
| fetch_step=args.fetch_step, | |
| infinity_mode=False, | |
| in_memory=args.in_memory, | |
| async_load=False, | |
| name='train', jets=True) | |
| iterator = iter(train_data) | |
| from time import time | |
| t0 = time() | |
| data = [] | |
| count = 0 | |
| while True: | |
| try: | |
| i = next(iterator) | |
| data.append(i) | |
| count += 1 | |
| if dataset_cap > 0 and count >= dataset_cap: | |
| break | |
| except StopIteration: | |
| break | |
| t1 = time() | |
| print("Took", t1-t0, "s -", datasets[0]) | |
| from src.dataset.functions_data import concat_events | |
| events = concat_events(data) # TODO: This can be done in a nicer way, using less memory (?) | |
| result = events.serialize() | |
| dir_name = datasets[0].split("/")[-2] | |
| save_to_dir = os.path.join(output_path, dir_name) | |
| Path(save_to_dir).mkdir(parents=True, exist_ok=True) | |
| for key in result[0]: | |
| with open(osp.join(save_to_dir, key + ".pkl"), "wb") as f: | |
| #pickle.dump(result[0][key], f) #save with torch for mmap | |
| #torch.save(result[0][key], f) | |
| np.save(f, result[0][key].numpy()) | |
| with open(osp.join(save_to_dir, "metadata.pkl"), "wb") as f: | |
| pickle.dump(result[1], f) | |
| print("Saved to", save_to_dir) | |
| print("Finished", dir_name) | |
| ''' | |
| from src.dataset.functions_data import EventCollection, EventJets, Event | |
| from src.dataset.dataset import EventDataset | |
| t2 = time() | |
| data1 = [] | |
| for event in EventDataset(result[0], result[1]): | |
| data1.append(event) | |
| t3 = time() | |
| print("Took", t3-t2, "s") | |
| print("Done") | |
| ''' | |
| output = get_path(args.output, "preprocessed_data") | |
| for dir in os.listdir(path): | |
| if args.overwrite or not os.path.exists(os.path.join(output, dir)): | |
| config = get_path('config_files/config_jets.yaml', 'code') | |
| if args.v2: | |
| delphes_suffix = "" | |
| if args.delphes: | |
| delphes_suffix = "_delphes" | |
| config = get_path(f'config_files/config_jets_2{delphes_suffix}.yaml', 'code') | |
| for i, file in enumerate(sorted(os.listdir(os.path.join(path, dir)))): | |
| print("Preprocessing file", file) | |
| preprocess_dataset([os.path.join(path, dir, file)], output + "_part"+str(i), config_file=config, dataset_cap=args.dataset_cap) | |
| else: | |
| print("Skipping", dir + ", already exists") | |
| # flush | |
| sys.stdout.flush() | |