diff --git a/src/config/config.py b/src/config/config.py index 242b13b..207bf27 100644 --- a/src/config/config.py +++ b/src/config/config.py @@ -54,7 +54,7 @@ class Run: defaults = [ {"run" : "base_run"}, {"mode" : "train"}, - {"data" : "data"}, + {"data" : "SBND"}, {"network" : "yolo"} ] @@ -71,6 +71,6 @@ class Config: data: Any = MISSING network: Yolo = MISSING output_dir: str = "output/${framework.name}/${network.name}/${run.id}/" - + name: str = "none" cs.store(name="base_config", node=Config) diff --git a/src/config/config.yaml b/src/config/config.yaml index 0fa7c31..01092ca 100644 --- a/src/config/config.yaml +++ b/src/config/config.yaml @@ -4,3 +4,4 @@ defaults: - _self_ output_dir: output/${network.name}//${run.id}/ +name: "SBND" diff --git a/src/config/config_dune.yaml b/src/config/config_dune.yaml new file mode 100644 index 0000000..b191167 --- /dev/null +++ b/src/config/config_dune.yaml @@ -0,0 +1,18 @@ + +defaults: + - base_config + - override data: DUNE + - _self_ + +output_dir: output/${network.name}//${run.id}/ +name: "DUNE" +# +# data: +# data_directory: "/lus/grand/projects/neutrino_osc_ADSP/datasets/DUNE/pixsim_full/" +# file: "merged_sample_chunk8_99.h5" +# aux_file: "merged_sample_chunk8_98.h5" +# image_width: 1536 +# image_height: 1024 +# pitch: 0.4 +# padding_x: 286 +# padding_y: 124 diff --git a/src/config/data.py b/src/config/data.py index 3411dd8..750bf5a 100644 --- a/src/config/data.py +++ b/src/config/data.py @@ -12,13 +12,34 @@ class ImageModeKind(Enum): @dataclass class Data: downsample: int = 1 - data_directory: str = "/lus/grand/projects/neutrino_osc_ADSP/datasets/DUNE/pixsim_full/" - file: str = "merged_sample_chunk8_99.h5" - aux_file: str = "merged_sample_chunk8_98.h5" + data_directory: str = "/lus/grand/projects/neutrino_osc_ADSP/datasets/SBND/VertexID/" + file: str = "sbnd_vertexID_train.h5" + aux_file: str = "sbnd_vertexID_val.h5" image_mode: ImageModeKind = ImageModeKind.dense input_dimension: int = 2 - image_width: int = 1536 - image_height: int = 1024 + image_width: int = 2048 + image_height: int = 1280 + pitch: float = 0.3 + padding_x: int = 0 + padding_y: int = 0 + +@dataclass +class SBND(Data): + data_directory: str = "/lus/grand/projects/neutrino_osc_ADSP/datasets/SBND/VertexID/" + file: str = "sbnd_vertexID_train.h5" + aux_file: str = "sbnd_vertexID_val.h5" + +@dataclass +class DUNE(Data): + data_directory: str = "/lus/grand/projects/neutrino_osc_ADSP/datasets/DUNE/pixsim_full/" + file: str = "merged_sample_chunk8_99.h5" + aux_file: str = "merged_sample_chunk8_98.h5" + image_width: int = 1536 + image_height: int = 1024 + pitch: float = 0.4 + padding_x: int = 286 + padding_y: int = 124 cs = ConfigStore.instance() -cs.store(group="data", name="data", node=Data) +cs.store(group="data", name="SBND", node=SBND) +cs.store(group="data", name="DUNE", node=DUNE) diff --git a/src/config/mode.py b/src/config/mode.py index 8ee718e..ed3727f 100644 --- a/src/config/mode.py +++ b/src/config/mode.py @@ -40,7 +40,7 @@ class Mode: @dataclass class Train(Mode): - checkpoint_iteration: int = 500 + checkpoint_iteration: int = 25 summary_iteration: int = 1 logging_iteration: int = 1 optimizer: Optimizer = Optimizer() diff --git a/src/networks/sparse_yolo.py b/src/networks/sparse_yolo.py index ace0b1f..e7f3ab2 100644 --- a/src/networks/sparse_yolo.py +++ b/src/networks/sparse_yolo.py @@ -164,14 +164,12 @@ class YOLOBlock(nn.Module): A YOLO block ''' - # def __init__(self, inp_dim_w, inp_dim_h, anchors, num_classes, cuda): - def __init__(self, inp_dim_w, inp_dim_h, anchors, num_classes): + def __init__(self, inp_dim_w, inp_dim_h, num_classes): nn.Module.__init__(self) self._inp_dim_w = inp_dim_w self._inp_dim_h = inp_dim_h - self._anchors = anchors self._num_classes = num_classes # self._cuda = cuda @@ -211,7 +209,7 @@ def filter_increase(n_filters): class YOLO(nn.Module): - def __init__(self, input_shape, anchors, args): + def __init__(self, input_shape, args): torch.nn.Module.__init__(self) # All of the parameters are controlled via the args module @@ -220,7 +218,6 @@ def __init__(self, input_shape, anchors, args): self.input_shape = input_shape # self.anchors = args.yolo_anchors - self.anchors = anchors self.num_classes = args.yolo_num_classes self.input_tensor = scn.InputLayer(dimension=3, spatial_size=[3,*input_shape[1:]]) @@ -328,7 +325,6 @@ def __init__(self, input_shape, anchors, args): self.yololayer_1 = YOLOBlock(inp_dim_w=self.input_shape[1], inp_dim_h=self.input_shape[2], - anchors=self.anchors, num_classes=self.num_classes, # cuda=self._cuda ) diff --git a/src/networks/yolo.py b/src/networks/yolo.py index 8af287a..996bf78 100644 --- a/src/networks/yolo.py +++ b/src/networks/yolo.py @@ -179,14 +179,12 @@ class YOLOBlock(nn.Module): A YOLO block ''' - # def __init__(self, inp_dim_w, inp_dim_h, anchors, num_classes, cuda): - def __init__(self, inp_dim_w, inp_dim_h, anchors, num_classes): + def __init__(self, inp_dim_w, inp_dim_h, num_classes): nn.Module.__init__(self) self._inp_dim_w = inp_dim_w self._inp_dim_h = inp_dim_h - self._anchors = anchors self._num_classes = num_classes # self._cuda = cuda @@ -225,7 +223,7 @@ def filter_increase(n_filters): class YOLO(nn.Module): - def __init__(self, input_shape, anchors, args): + def __init__(self, input_shape, args): torch.nn.Module.__init__(self) # All of the parameters are controlled via the args module @@ -234,8 +232,6 @@ def __init__(self, input_shape, anchors, args): self.input_shape = input_shape - # self.anchors = args.yolo_anchors - self.anchors = anchors self.num_classes = args.yolo_num_classes @@ -342,7 +338,6 @@ def __init__(self, input_shape, anchors, args): self.yololayer_1 = YOLOBlock(inp_dim_w=self.input_shape[1], inp_dim_h=self.input_shape[2], - anchors=self.anchors, num_classes=self.num_classes, # cuda=self._cuda ) diff --git a/src/utils/larcv_fetcher.py b/src/utils/larcv_fetcher.py index 88bd32b..e643402 100644 --- a/src/utils/larcv_fetcher.py +++ b/src/utils/larcv_fetcher.py @@ -16,7 +16,7 @@ class larcv_fetcher(object): - def __init__(self, mode, distributed, access_mode, dimension, data_format, downsample_images, seed=None): + def __init__(self, config, mode, distributed, access_mode, dimension, data_format, downsample_images, seed=None): if mode not in ['train', 'inference', 'iotest']: raise Exception("Larcv Fetcher can't handle mode ", mode) @@ -33,6 +33,7 @@ def __init__(self, mode, distributed, access_mode, dimension, data_format, downs self._larcv_interface = queueloader.queue_interface( random_access_mode=access_mode, seed=seed) + self.config = config self.mode = mode self.image_mode = data_format self.input_dimension = dimension @@ -49,7 +50,7 @@ def __del__(self): - def prepare_sample(self, name, input_file, batch_size, color=None, start_index = 0, print_config=False): + def prepare_sample(self, name, input_file, batch_size, color=None, start_index = 0, print_config=True): # If quotations are in the file name, remove them if "\"" in input_file: @@ -78,14 +79,24 @@ def prepare_sample(self, name, input_file, batch_size, color=None, start_index = data_keys = {} data_keys['image'] = name+'data' + # Need to embed the data in SBND: + if self.config.name == "SBND" and self.input_dimension == 2: + cb.add_preprocess( + datatype = "sparse2d", + producer = "sbndwire", + process = "Embed", + OutputProducer = "sbndwire", + TargetSize = [2048,1280] + ) + # Downsampling if self.downsample_images != 0: cb.add_preprocess( datatype = "sparse2d", Product = "sparse2d", - producer = "dunevoxels", + producer = "sbndwire" if self.config.name == 'SBND' else "dunevoxels", process = "Downsample", - OutputProducer = "dunevoxels", + OutputProducer = "sbndwire" if self.config.name == 'SBND' else "dunevoxels", Downsample = 2**self.downsample_images, PoolType = 1 # average, # PoolType = 2 # max @@ -95,12 +106,12 @@ def prepare_sample(self, name, input_file, batch_size, color=None, start_index = # Need to load up on data fillers. if self.input_dimension == 2: cb.add_batch_filler( - datatype = "sparse2d", - producer = "dunevoxels", - name = name+"data", + datatype = "sparse2d", + producer = "sbndwire" if self.config.name == 'SBND' else "dunevoxels", + name = name+"data", MaxVoxels = 20000, - Augment = False, - Channels = [0, 1, 2], + Augment = False, + Channels = [0, 1, 2], ) else: @@ -113,30 +124,39 @@ def prepare_sample(self, name, input_file, batch_size, color=None, start_index = ) # Add something to convert the neutrino particles into bboxes: - cb.add_preprocess( - datatype = "particle", - producer = "neutrino", - process = "BBoxFromParticle", - OutputProducer = "neutrino" - ) - cb.add_batch_filler( - datatype = "bbox3d", - producer = "neutrino", - name = name+"bbox", - MaxBoxes = 2, - Channels = [0,] + if self.config.name == 'DUNE': + cb.add_preprocess( + datatype = "particle", + producer = "neutrino", + process = "BBoxFromParticle", + OutputProducer = "neutrino" ) + cb.add_batch_filler( + datatype = "bbox3d", + producer = "neutrino", + name = name+"bbox", + MaxBoxes = 2, + Channels = [0,] + ) + else: + cb.add_batch_filler( + datatype = "bbox2d", + producer = "bbox_neutrino", + name = name+"bbox", + MaxBoxes = 1, + Channels = [0,1,2] + ) # Add the label configs: - for label_name, l in zip(['neut', 'prot', 'cpi', 'npi'], [3, 3, 2, 2]): - cb.add_batch_filler( - datatype = "PID", - producer = f"{label_name}ID", - name = name+f'label_{label_name}', - PdgClassList = [i for i in range(l)] - ) - data_keys[f'label_{label_name}'] = name+f'label_{label_name}' + # for label_name, l in zip(['neut', 'prot', 'cpi', 'npi'], [3, 3, 2, 2]): + # cb.add_batch_filler( + # datatype = "PID", + # producer = f"{label_name}ID", + # name = name+f'label_{label_name}', + # PdgClassList = [i for i in range(l)] + # ) + # data_keys[f'label_{label_name}'] = name+f'label_{label_name}' if print_config: @@ -238,8 +258,24 @@ def fetch_next_batch(self, name, force_pop=False): # Parse out the vertex info: - minibatch_data['vertex'] = minibatch_data['vertex'][:,:,0,0:3] - minibatch_data['vertex'] = minibatch_data['vertex'].reshape((-1, 3)) + if self.config.name == 'DUNE': + minibatch_data['vertex'] = minibatch_data['vertex'][:,:,0,0:3] + minibatch_data['vertex'] = minibatch_data['vertex'].reshape((-1, 3)) + + else: + minibatch_data['vertex'] = minibatch_data['vertex'][:,:,0,0:2] + + # These should be retrieved by the image meta TODO + width = 614.40 + height = 399.51 + min_x = [-9.6, -9.6, -57.59] + min_y = [1.87, 1.87, 1.87] + + for p in [0, 1, 2]: + minibatch_data['vertex'][:,p,0] = self.config.data.image_width * (minibatch_data['vertex'][:,p,0] - min_x[p]) / width + minibatch_data['vertex'][:,p,1] = self.config.data.image_height * (minibatch_data['vertex'][:,p,1] - min_y[p]) / height + + # print('minibatch_data[vertex]:', minibatch_data['vertex']) # Also, we map the vertex from 0 to 1 across the image. The image size is # [360, 200, 500] and the origin is at [0, -100, 0] @@ -253,8 +289,8 @@ def fetch_next_batch(self, name, force_pop=False): if self.input_dimension == 3: minibatch_data['image'] = data_transforms.larcvsparse_to_dense_3d(minibatch_data['image']) else: - x_dim = int(1536/2**self.downsample_images) - y_dim = int(1024/2**self.downsample_images) + x_dim = int(self.config.data.image_width / 2**self.downsample_images) + y_dim = int(self.config.data.image_height / 2**self.downsample_images) minibatch_data['image'] = data_transforms.larcvsparse_to_dense_2d(minibatch_data['image'], dense_shape=[x_dim, y_dim]) elif self.image_mode == ImageModeKind.sparse: # Have to convert the input image from dense to sparse format: diff --git a/src/utils/trainercore.py b/src/utils/trainercore.py index 1434405..2587903 100644 --- a/src/utils/trainercore.py +++ b/src/utils/trainercore.py @@ -50,6 +50,7 @@ def __init__(self, args): logger.info(f"Access mode: {access_mode}") self.larcv_fetcher = larcv_fetcher.larcv_fetcher( + config = self.args, mode = self.mode, distributed = self.args.run.distributed, access_mode = access_mode, @@ -137,13 +138,12 @@ def init_network(self): # To initialize the network, we see what the name is # and act on that: # if self.args.network == "yolo": - anchors = [(116, 90), (156, 198), (373, 326)] if self.args.data.image_mode == ImageModeKind.dense: from src.networks import yolo - self._net = yolo.YOLO(input_shape, anchors, self.args.network) + self._net = yolo.YOLO(input_shape, self.args.network) elif self.args.data.image_mode == ImageModeKind.sparse: from src.networks import sparse_yolo - self._net = sparse_yolo.YOLO(input_shape, anchors, self.args.network) + self._net = sparse_yolo.YOLO(input_shape, self.args.network) # else: # raise Exception(f"Couldn't identify network {self.args.network.name}") @@ -427,6 +427,14 @@ def _3d_to_2d(self, point, plane, pitch=0.4): def _target_to_yolo(self, target, logits): + + if self.args.name == 'SBND': + return self._target_to_yolo_sbnd(target, logits) + else: + return self._target_to_yolo_dune(target, logits) + + + def _target_to_yolo_dune(self, target, logits): ''' Takes the vertex data from larcv and transforms it to YOLO output. @@ -447,9 +455,9 @@ def _target_to_yolo(self, target, logits): target_out = [] mask = [] - pitch = 0.4 * 2**self.args.data.downsample - padding_x = 286 / 2**self.args.data.downsample - padding_y = 124 / 2**self.args.data.downsample + pitch = self.args.data.pitch * 2**self.args.data.downsample + padding_x = self.args.data.padding_x / 2**self.args.data.downsample + padding_y = self.args.data.padding_y / 2**self.args.data.downsample batch_size = target.size(0) @@ -510,6 +518,76 @@ def _target_to_yolo(self, target, logits): return target_out, mask + def _target_to_yolo_sbnd(self, target, logits): + ''' + Takes the vertex data from larcv and transforms it + to YOLO output. + + arguments: + - target: the data from larcv + - n_channels: the number of channels used during training + - grid_size_w: the image width at the end of the network + - grid_size_h: the image height at the end of the network + + returns: + - the transformed target + - a mask that can mask the entries where there are real objects + ''' + # print('True vertex position:', target) + + with self.default_device_context(): + target_out = [] + mask = [] + + batch_size = target.size(0) + + for plane in range(len(logits)): + + logits_p = logits[plane] + + n_channels = logits_p.size(3) + grid_size_w = logits_p.size(1) + grid_size_h = logits_p.size(2) + + target_out_p = torch.zeros(batch_size, grid_size_w, grid_size_h, n_channels, device=self.default_device()) + mask_p = torch.zeros(batch_size, grid_size_w, grid_size_h, dtype=torch.bool, device=self.default_device()) + + step_w = self.args.data.image_width / grid_size_w + step_h = self.args.data.image_height / grid_size_h + + for batch_id in range(batch_size): + + projected_x = target[batch_id, plane, 0] + projected_y = target[batch_id, plane, 1] + + t_x = projected_x / step_w + t_i = int(t_x) + t_y = projected_y / step_h + t_j = int(t_y) + + if t_j > 39: t_j = 39 + if t_j < 0: t_j = 0 + + target_out_p[batch_id, t_i, t_j, 0] = t_x - t_i + target_out_p[batch_id, t_i, t_j, 1] = t_y - t_j + target_out_p[batch_id, t_i, t_j, 2] = 1. + target_out_p[batch_id, t_i, t_j, 3] = 1. + + mask_p[batch_id, t_i, t_j] = 1 + + # print('Batch', batch_id, 't_i', t_i, 't_j', t_j) + + # print('target_out_p', target_out_p) + if self._global_step == 0: + if not self.args.run.distributed or self._rank == 0: + numpy.save(f'yolo_tgt_{plane}', target_out_p.cpu()) + + target_out.append(target_out_p) + mask.append(mask_p) + + return target_out, mask + + def _calculate_loss(self, target, mask, prediction, full=False): ''' Calculates the loss. @@ -667,16 +745,16 @@ def _calculate_accuracy_per_plane(self, target, prediction): x_targ = target[mask_targ][:,0] y_targ = target[mask_targ][:,1] - # if self._global_step % 25 == 0: - # if not self.args.run.distributed or self._rank == 0: - # numpy.save('xypred', numpy.array([x_pred.detach().cpu().float(), y_pred.detach().cpu().float()])) - # numpy.save('xytarg', numpy.array([x_targ.detach().cpu().float(), y_targ.detach().cpu().float()])) + if self._global_step == 0: + if not self.args.run.distributed or self._rank == 0: + numpy.save('xypred', numpy.array([x_pred.detach().cpu().float(), y_pred.detach().cpu().float()])) + numpy.save('xytarg', numpy.array([x_targ.detach().cpu().float(), y_targ.detach().cpu().float()])) grid_size_w = prediction.size(1) grid_size_h = prediction.size(2) - pitch = 0.4 # cm - a = 1536 / grid_size_w * pitch - b = 1024 / grid_size_h * pitch + # pitch = 0.4 # cm + a = self.args.data.image_width / grid_size_w * self.args.data.pitch + b = self.args.data.image_height / grid_size_h * self.args.data.pitch resolution = 0 with torch.no_grad(): @@ -873,9 +951,9 @@ def train_step(self): minibatch_data = self.larcv_fetcher.fetch_next_batch("primary", force_pop=True) io_end_time = datetime.datetime.now() - # if self._global_step % 25 == 0 and self._rank == 0: - # numpy.save("img",minibatch_data['image']) - # numpy.save("vtx",minibatch_data['vertex']) + if self._global_step == 0 and self._rank == 0: + numpy.save("img",minibatch_data['image']) + numpy.save("vtx",minibatch_data['vertex']) minibatch_data = self.to_torch(minibatch_data)