VertexID with SBND files

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)

src/config/config.yaml

@@ -4,3 +4,4 @@ defaults:
   - _self_
 
 output_dir: output/${network.name}//${run.id}/
+name: "SBND"

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

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)

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()

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
                                      )

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
                                      )

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: