Fix apply on child

alonet/detr/train.py

@@ -88,7 +88,7 @@ def add_argparse_args(parent_parser, parser=None):
         )
         return parent_parser
 
-    def forward(self, frames: Union[list, Frame]):
+    def forward(self, frames: Union[list, Frame], **kwargs):
         """Run a forward pass through the model.
 
         Parameters
@@ -107,10 +107,10 @@ def forward(self, frames: Union[list, Frame]):
         # Assert inputs content
         self.assert_input(frames, inference=True)
         # Run forward pass
-        m_outputs = self.model(frames)
+        m_outputs = self.model(frames, **kwargs)
         return m_outputs
 
-    def inference(self, m_outputs: dict):
+    def inference(self, m_outputs: dict, **kwargs):
         """Given the model forward outputs, this method will return an
         :mod:`BoundingBoxes2D <aloscene.bounding_boxes_2d>` tensor.
 
@@ -124,7 +124,7 @@ def inference(self, m_outputs: dict):
         List[:mod:`BoundingBoxes2D <aloscene.bounding_boxes_2d>`]
             Set of boxes for each batch
         """
-        return self.model.inference(m_outputs)
+        return self.model.inference(m_outputs, **kwargs)
 
     def training_step(self, frames: Union[list, Frame], batch_idx: int):
         """Train the model for one step
@@ -382,7 +382,7 @@ def run_train(
 
     # Logger config
     logging.basicConfig(
-        level=logging.INFO, format="[%(asctime)s][%(levelname)s] %(message)s", datefmt="%d-%m-%y %H:%M:%S",
+        level=logging.INFO, format="[%(asctime)s][%(levelname)s] %(message)s", datefmt="%d-%m-%y %H:%M:%S"
     )
     logger = logging.getLogger("aloception")
 

alonet/detr_panoptic/utils.py

@@ -44,7 +44,7 @@ def get_mask_queries(
     feat_size = dec_outputs.shape[2:]
 
     dec_outputs = [
-        torch.cat([dec_outputs[b : b + 1, idx], torch.zeros(1, max_size - fs, *feat_size, device=device)], dim=1,)
+        torch.cat([dec_outputs[b : b + 1, idx], torch.zeros(1, max_size - fs, *feat_size, device=device)], dim=1)
         for b, (idx, fs) in enumerate(zip(filters, fsizes))
     ]
     return torch.cat(dec_outputs, dim=0), filters
@@ -72,7 +72,7 @@ def criterion(b):
         b.labels = b.labels[cat]
 
     if isinstance(frames.boxes2d[0].labels, dict):
-        frames.apply_on_label(frames.boxes2d, criterion)
+        frames.apply_on_child(frames.boxes2d, criterion)
     if isinstance(frames.segmentation[0].labels, dict):
-        frames.apply_on_label(frames.segmentation, criterion)
+        frames.apply_on_child(frames.segmentation, criterion)
     return frames