1st step of fixing training + debug

happypose/pose_estimators/cosypose/cosypose/datasets/augmentations.py

@@ -143,7 +143,6 @@ def __call__(self, im, mask, obs):
 
 class VOCBackgroundAugmentation(BackgroundAugmentation):
     def __init__(self, voc_root, p=0.3):
-        print("voc_root =", voc_root)
         image_dataset = VOCSegmentation(
             root=voc_root,
             year="2012",

happypose/pose_estimators/cosypose/cosypose/evaluation/meters/utils.py

@@ -25,7 +25,7 @@ def add_inst_num(
     group_keys=["scene_id", "view_id", "label"],
     key="pred_inst_num",
 ):
-    inst_num = np.empty(len(infos), dtype=np.int)
+    inst_num = np.empty(len(infos), dtype=int)
     for _group_name, group_ids in infos.groupby(group_keys).groups.items():
         inst_num[group_ids.values] = np.arange(len(group_ids))
     infos[key] = inst_num

happypose/pose_estimators/cosypose/cosypose/evaluation/prediction_runner.py

@@ -95,18 +95,18 @@ def run_inference_pipeline(
 
 
         """
-        if self.inference_cfg.detection_type == "gt":
+        if self.inference_cfg['detection_type'] == "gt":
             detections = gt_detections
             run_detector = False
-        elif self.inference_cfg.detection_type == "detector":
+        elif self.inference_cfg['detection_type'] == "detector":
             detections = None
             run_detector = True
         else:
-            msg = f"Unknown detection type {self.inference_cfg.detection_type}"
+            msg = f"Unknown detection type {self.inference_cfg['detection_type']}"
             raise ValueError(msg)
 
         coarse_estimates = None
-        if self.inference_cfg.coarse_estimation_type == "external":
+        if self.inference_cfg['coarse_estimation_type'] == "external":
             # TODO (ylabbe): This is hacky, clean this for modelnet eval.
             coarse_estimates = initial_estimates
             coarse_estimates = happypose.toolbox.inference.utils.add_instance_id(
@@ -137,15 +137,15 @@ def run_inference_pipeline(
         all_preds = {}
         data_TCO_refiner = extra_data["refiner"]["preds"]
 
-        k_0 = f"refiner/iteration={self.inference_cfg.n_refiner_iterations}"
+        k_0 = f"refiner/iteration={self.inference_cfg['n_refiner_iterations']}"
         all_preds = {
             "final": preds,
             k_0: data_TCO_refiner,
             "refiner/final": data_TCO_refiner,
             "coarse": extra_data["coarse"]["preds"],
         }
 
-        if self.inference_cfg.run_depth_refiner:
+        if self.inference_cfg['run_depth_refiner']:
             all_preds["depth_refiner"] = extra_data["depth_refiner"]["preds"]
 
         # Remove any mask tensors
@@ -174,43 +174,46 @@ def get_predictions(
         """
         predictions_list = defaultdict(list)
         for n, data in enumerate(tqdm(self.dataloader)):
-            # data is a dict
-            rgb = data["rgb"]
-            depth = None
-            K = data["cameras"].K
-            gt_detections = data["gt_detections"].cuda()
-
-            initial_data = None
-            if data["initial_data"]:
-                initial_data = data["initial_data"].cuda()
-
-            obs_tensor = ObservationTensor.from_torch_batched(rgb, depth, K)
-            obs_tensor = obs_tensor.cuda()
-
-            # GPU warmup for timing
-            if n == 0:
+            if n < 3:
+                # data is a dict
+                rgb = data["rgb"]
+                depth = None
+                K = data["cameras"].K
+                gt_detections = data["gt_detections"].cuda()
+
+                initial_data = None
+                if data["initial_data"]:
+                    initial_data = data["initial_data"].cuda()
+
+                obs_tensor = ObservationTensor.from_torch_batched(rgb, depth, K)
+                obs_tensor = obs_tensor.cuda()
+
+                # GPU warmup for timing
+                if n == 0:
+                    with torch.no_grad():
+                        self.run_inference_pipeline(
+                            pose_estimator,
+                            obs_tensor,
+                            gt_detections,
+                            initial_estimates=initial_data,
+                        )
+
+                cuda_timer = CudaTimer()
+                cuda_timer.start()
                 with torch.no_grad():
-                    self.run_inference_pipeline(
+                    all_preds = self.run_inference_pipeline(
                         pose_estimator,
                         obs_tensor,
                         gt_detections,
                         initial_estimates=initial_data,
                     )
+                cuda_timer.end()
+                cuda_timer.elapsed()
 
-            cuda_timer = CudaTimer()
-            cuda_timer.start()
-            with torch.no_grad():
-                all_preds = self.run_inference_pipeline(
-                    pose_estimator,
-                    obs_tensor,
-                    gt_detections,
-                    initial_estimates=initial_data,
-                )
-            cuda_timer.end()
-            cuda_timer.elapsed()
-
-            for k, v in all_preds.items():
-                predictions_list[k].append(v)
+                for k, v in all_preds.items():
+                    predictions_list[k].append(v)
+            else:
+                break
 
         # Concatenate the lists of PandasTensorCollections
         predictions = {}

happypose/pose_estimators/cosypose/cosypose/evaluation/runner_utils.py

@@ -14,8 +14,8 @@
 def run_pred_eval(pred_runner, pred_kwargs, eval_runner, eval_preds=None):
     all_predictions = {}
     for pred_prefix, pred_kwargs_n in pred_kwargs.items():
-        print("Prediction :", pred_prefix)
-        preds = pred_runner.get_predictions(**pred_kwargs_n)
+        pose_predictor = pred_kwargs_n['pose_predictor']
+        preds = pred_runner.get_predictions(pose_predictor)
         for preds_name, preds_n in preds.items():
             all_predictions[f"{pred_prefix}/{preds_name}"] = preds_n
 

happypose/pose_estimators/cosypose/cosypose/models/pose.py

@@ -167,6 +167,7 @@ def forward(self, images, K, labels, TCO, n_iterations=1):
                 K_crop=K_crop,
                 boxes_rend=boxes_rend,
                 boxes_crop=boxes_crop,
+                model_outputs=model_outputs
             )
 
             TCO_input = TCO_output

happypose/pose_estimators/cosypose/cosypose/scripts/run_pose_training.py

@@ -10,6 +10,9 @@
 logger = get_logger(__name__)
 
 
+import warnings
+warnings.filterwarnings("ignore")
+
 def make_cfg(args):
     cfg = argparse.ArgumentParser("").parse_args([])
     if args.config:
@@ -56,7 +59,7 @@ def make_cfg(args):
     cfg.n_pose_dims = 9
     cfg.n_rendering_workers = N_WORKERS
     cfg.refiner_run_id_for_test = None
-    cfg.coarse_run_id_for_test = None
+    cfg.coarse_run_id_for_test = "coarse-bop-ycbv-pbr--724183"
 
     # Optimizer
     cfg.lr = 3e-4
@@ -66,7 +69,7 @@ def make_cfg(args):
     cfg.clip_grad_norm = 0.5
 
     # Training
-    cfg.batch_size = 32
+    cfg.batch_size = 16
     cfg.epoch_size = 115200
     cfg.n_epochs = 700
     cfg.n_dataloader_workers = N_WORKERS

happypose/pose_estimators/cosypose/cosypose/training/pose_forward_loss.py

@@ -63,10 +63,11 @@ def h_pose(model, mesh_db, data, meters, cfg, n_iterations=1, input_generator="f
     losses_TCO_iter = []
     for n in range(n_iterations):
         iter_outputs = outputs[f"iteration={n+1}"]
-        K_crop = iter_outputs["K_crop"]
-        TCO_input = iter_outputs["TCO_input"]
-        TCO_pred = iter_outputs["TCO_output"]
-        model_outputs = iter_outputs["model_outputs"]
+        K_crop = iter_outputs.K_crop
+        TCO_input = iter_outputs.TCO_input
+        TCO_pred = iter_outputs.TCO_output
+        model_outputs = iter_outputs.model_outputs
+
 
         if cfg.loss_disentangled:
             if cfg.n_pose_dims == 9:

happypose/pose_estimators/cosypose/cosypose/training/train_pose.py

@@ -97,7 +97,7 @@ def load_model(run_id):
         if run_id is None:
             return None
         run_dir = EXP_DIR / run_id
-        cfg = yaml.load((run_dir / "config.yaml").read_text(), Loader=yaml.FullLoader)
+        cfg = yaml.load((run_dir / "config.yaml").read_text(), Loader=yaml.Loader)
         cfg = check_update_config(cfg)
         model = (
             create_model_pose(
@@ -422,7 +422,7 @@ def train_epoch():
             iterator = tqdm(ds_iter_train, ncols=80)
             t = time.time()
             for n, sample in enumerate(iterator):
-                if n < 5:
+                if n < 3:
                     if n > 0:
                         meters_time["data"].add(time.time() - t)
 
@@ -453,19 +453,19 @@ def train_epoch():
                         lr_scheduler_warmup.step()
                     t = time.time()
                 else:
-                    continue
+                    break
             if epoch >= args.n_epochs_warmup:
                 lr_scheduler.step()
 
         @torch.no_grad()
         def validation():
             model.eval()
             for n, sample in enumerate(tqdm(ds_iter_val, ncols=80)):
-                if n < 5:
+                if n < 3:
                     loss = h(data=sample, meters=meters_val)
                     meters_val["loss_total"].add(loss.item())
                 else:
-                    continue
+                    break
 
         @torch.no_grad()
         def test():

happypose/pose_estimators/megapose/evaluation/runner_utils.py

@@ -30,7 +30,6 @@
 def run_pred_eval(pred_runner, pred_kwargs, eval_runner, eval_preds=None):
     all_predictions = {}
     for pred_prefix, pred_kwargs_n in pred_kwargs.items():
-        print("Prediction :", pred_prefix)
         preds = pred_runner.get_predictions(**pred_kwargs_n)
         for preds_name, preds_n in preds.items():
             all_predictions[f"{pred_prefix}/{preds_name}"] = preds_n

happypose/pose_estimators/megapose/models/pose_rigid.py

@@ -59,6 +59,7 @@ class PosePredictorOutputCosypose:
     K_crop: torch.Tensor
     boxes_rend: torch.Tensor
     boxes_crop: torch.Tensor
+    model_outputs: torch.Tensor
 
 
 @dataclass

happypose/toolbox/lib3d/rigid_mesh_database.py

@@ -146,6 +146,8 @@ def n_sym_mapping(self):
         return {label: obj["n_sym"] for label, obj in self.infos.items()}
 
     def select(self, labels):
+        print("label to id =", self.label_to_id)
+        print("labels = ", labels)
         ids = [self.label_to_id[label] for label in labels]
         return Meshes(
             infos=[self.infos[label] for label in labels],