Update eval

run/train_3d.py

@@ -54,6 +54,7 @@ def save_model(
     load=0,
     name="pointpillars_car",
     config="xyres_16.proto",
+    num_ensembles=None,
 ):
 
     config = os.path.join(config_roots[model_type], config,)
@@ -64,7 +65,8 @@ def save_model(
         device=device,
         checkpoint_after_iter=1000,
         checkpoint_load_iter=load,
-        return_uncertainty=True
+        return_uncertainty=True,
+        num_ensembles=num_ensembles,
     )
 
     learner.save(model_path)
@@ -266,6 +268,35 @@ def test_model(
                 f.write("\n\n")
 
 
+def test_ens_model(
+    model_type,
+    device="cuda:0",
+    name="pointpillars_car",
+    samples_list=[1],
+    config="xyres_16.proto",
+    eval_suffix="classic",
+):
+
+    config = os.path.join(config_roots[model_type], config,)
+    model_path = os.path.join(temp_dir, name)
+
+    results = {}
+    with open(os.path.join(model_path, "eval_" + eval_suffix + ".txt"), "w") as f:
+        for samples in samples_list:
+            print("samples =", samples)
+            learner = VoxelObjectDetection3DLearner(
+                model_config_path=config, device=device, checkpoint_after_iter=1000, return_uncertainty=True,
+                num_ensembles=samples
+            )
+            learner.load(model_path)
+
+            result = learner.eval(dataset, verbose=True, samples=samples)
+            results[samples] = result
+            f.write("samples = " + str(samples) + "\n")
+            f.write("3d AP: " + str(result[2][0, :, 0]))
+            f.write("\n\n")
+
+
 def test_pointpillars(
     device="cuda:0",
     name="pointpillars_car",
@@ -373,7 +404,7 @@ def test_ens_tanet(
     config="ens_xyres_16.proto",
     eval_suffix="vnn",
 ):
-    return test_model(
+    return test_ens_model(
         "tanet",
         device=device,
         name=name,

src/opendr/perception/object_detection_3d/voxel_object_detection_3d/second_detector/pytorch/models/tanet.py

@@ -1103,7 +1103,7 @@ class EnsPSA(nn.Module):
             name=name,
         )
     
-    def forward(self, x, samples):
+    def forward(self, x, samples, combine_predictions=False):
 
         assert samples == self.num_ensembles
 
@@ -1125,8 +1125,12 @@ class EnsPSA(nn.Module):
         for key, value in keys.items():
             result[value] = []
 
-        for output in all_outputs[1:]:
+        for output in all_outputs:
             for key, value in keys.items():
                 result[value].append(output[key])
 
+        if combine_predictions:
+            for key, value in keys.items():
+                result[value] = [torch.mean(torch.stack(result[value]), dim=0)]
+
         return result

src/opendr/perception/object_detection_3d/voxel_object_detection_3d/second_detector/pytorch/models/voxelnet.py

@@ -1113,7 +1113,7 @@ class VoxelNet(nn.Module):
     def get_global_step(self):
         return int(self.global_step.cpu().numpy()[0])
 
-    def forward(self, example, refine_weight=2, samples=1, return_uncertainty=False):
+    def forward(self, example, refine_weight=2, samples=1, return_uncertainty=False, combine_predictions=False):
         """module's forward should always accept dict and return loss.
         """
         voxels = example["voxels"]
@@ -1137,7 +1137,7 @@ class VoxelNet(nn.Module):
                 if self._use_bev:
                     preds_dict = self.rpn(spatial_features, example["bev_map"], samples=samples)
                 else:
-                    preds_dict = self.rpn(spatial_features, samples=samples)
+                    preds_dict = self.rpn(spatial_features, samples=samples, combine_predictions=combine_predictions)
             else:
                 if self._use_bev:
                     raise Exception("Can not use bev in ONNX session")
@@ -1348,11 +1348,11 @@ class VoxelNet(nn.Module):
                 }
         else:
             if self.rpn_class_name in ["PSA", "VPSA", "EnsPSA", "RefineDet"]:
-                coarse_output = self.predict_coarse(example, preds_dict, self.device, samples=samples)
-                refine_output = self.predict_refine(example, preds_dict, self.device, samples=samples)
+                coarse_output = self.predict_coarse(example, preds_dict, self.device, samples=1 if combine_predictions else samples)
+                refine_output = self.predict_refine(example, preds_dict, self.device, samples=1 if combine_predictions else samples)
                 return coarse_output, refine_output
             else:
-                return self.predict_coarse(example, preds_dict, self.device, samples=samples)
+                return self.predict_coarse(example, preds_dict, self.device, samples=1 if combine_predictions else samples)
 
     def compute_predict(
         self,

src/opendr/perception/object_detection_3d/voxel_object_detection_3d/second_detector/run.py

@@ -214,7 +214,7 @@ def train(
                     cls_neg_loss.detach().cpu().numpy())
 
                 ########################################
-                if (model_cfg.rpn.module_class_name in ["PSA", "VPSA", "RefineDet"]):
+                if (model_cfg.rpn.module_class_name in ["PSA", "VPSA", "EnsPSA", "RefineDet"]):
                     coarse_loss = ret_dict["coarse_loss"]
                     refine_loss = ret_dict["refine_loss"]
                     metrics["coarse_loss"] = float(
@@ -265,7 +265,7 @@ def train(
             log(Logger.LOG_WHEN_VERBOSE, "#################################")
             log(Logger.LOG_WHEN_NORMAL, "Generate output labels...")
             t = time.time()
-            if model_cfg.rpn.module_class_name in ["PSA", "VPSA", "RefineDet"]:
+            if model_cfg.rpn.module_class_name in ["PSA", "VPSA", "EnsPSA", "RefineDet"]:
                 dt_annos_coarse = []
                 dt_annos_refine = []
                 prog_bar = ProgressBar()
@@ -335,7 +335,7 @@ def train(
 
         if evaluate:
 
-            if model_cfg.rpn.module_class_name in ["PSA", "VPSA", "RefineDet"]:
+            if model_cfg.rpn.module_class_name in ["PSA", "VPSA", "EnsPSA", "RefineDet"]:
 
                 log(Logger.LOG_WHEN_NORMAL, "Before Refine:")
                 (
@@ -418,7 +418,7 @@ def evaluate(
     net.eval()
     t = time.time()
 
-    if model_cfg.rpn.module_class_name in ["PSA", "VPSA", "RefineDet"]:
+    if model_cfg.rpn.module_class_name in ["PSA", "VPSA", "EnsPSA", "RefineDet"]:
         dt_annos_coarse = []
         dt_annos_refine = []
         log(Logger.LOG_WHEN_NORMAL, "Generate output labels...")
@@ -439,6 +439,7 @@ def evaluate(
                 use_coarse_to_fine=True,
                 global_set=None,
                 image_shape=image_shape,
+                samples=samples,
             )
             dt_annos_coarse += coarse
             dt_annos_refine += refine
@@ -478,7 +479,7 @@ def evaluate(
                 log=lambda *x, **y: log(Logger.LOG_WHEN_NORMAL, *x, **y))
 
     if count is not None:
-        if model_cfg.rpn.module_class_name in ["PSA", "VPSA", "RefineDet"]:
+        if model_cfg.rpn.module_class_name in ["PSA", "VPSA", "EnsPSA", "RefineDet"]:
             gt_annos = gt_annos[:len(dt_annos_refine)]
         else:
             gt_annos = gt_annos[:len(dt_annos)]
@@ -499,7 +500,7 @@ def evaluate(
     )
     if not predict_test:
 
-        if model_cfg.rpn.module_class_name in ["PSA", "VPSA", "RefineDet"]:
+        if model_cfg.rpn.module_class_name in ["PSA", "VPSA", "EnsPSA", "RefineDet"]:
             log(Logger.LOG_WHEN_NORMAL, "Before Refine:")
             result_coarse = get_official_eval_result(gt_annos, dt_annos_coarse,
                                                      class_names)
@@ -666,8 +667,30 @@ def predict_kitti_to_anno(
         [image_shape] * len(example["P2"])
     )
 
+    batch_size = len(example["P2"])
+
+    # def combine_preds(p):
+
+    #     result =  [{} for _ in range(batch_size)]
+
+    #     for element in p:
+    #         for batch_id, batch_element in enumerate(element):
+    #             for key, value in batch_element.items():
+    #                 if key not in result[batch_id]:
+    #                     if key in ["image_idx", "label_preds"]:
+    #                         result[batch_id][key] = value
+    #                     else:
+    #                         result[batch_id][key] = value / batch_size
+    #                 else:
+    #                     if key not in ["image_idx", "label_preds"]:
+    #                         result[batch_id][key] += value / batch_size
+
     if use_coarse_to_fine:
-        predictions_dicts_coarse, predictions_dicts_refine = net(example, samples=samples)
+        predictions_dicts_coarse, predictions_dicts_refine = net(example, samples=samples, combine_predictions=True)
+
+        predictions_dicts_coarse = predictions_dicts_coarse[0]
+        predictions_dicts_refine = predictions_dicts_refine[0]
+
         annos_coarse = comput_kitti_output(
             predictions_dicts_coarse,
             batch_image_shape,

src/opendr/perception/object_detection_3d/voxel_object_detection_3d/voxel_object_detection_3d_learner.py

@@ -528,7 +528,7 @@ class VoxelObjectDetection3DLearner(Learner):
             example, samples=samples, return_uncertainty=self.return_uncertainty,
         )
 
-        if self.model_config.rpn.module_class_name in ["PSA", "VPSA", "RefineDet"]:
+        if self.model_config.rpn.module_class_name in ["PSA", "VPSA", "EnsPSA", "RefineDet"]:
             output = output[-1]
 
         all_annotations = []
@@ -586,7 +586,7 @@ class VoxelObjectDetection3DLearner(Learner):
             self.model.middle_feature_extractor.nx,
         ]
 
-        has_refine = self.model.rpn_class_name in ["PSA", "VPSA", "RefineDet"]
+        has_refine = self.model.rpn_class_name in ["PSA", "VPSA", "EnsPSA", "RefineDet"]
 
         try:
             self.__convert_rpn_to_onnx(