pre-commit run -a

happypose/pose_estimators/cosypose/cosypose/training/train_detector.py

@@ -231,7 +231,9 @@ def make_datasets(dataset_names):
         logger.info(f"Using pretrained model from {pretrain_path}.")
         model.load_state_dict(torch.load(pretrain_path)["state_dict"])
     elif args.pretrain_coco:
-        state_dict = load_state_dict_from_url('https://download.pytorch.org/models/maskrcnn_resnet50_fpn_coco-bf2d0c1e.pth')
+        state_dict = load_state_dict_from_url(
+            "https://download.pytorch.org/models/maskrcnn_resnet50_fpn_coco-bf2d0c1e.pth"
+        )
 
         def keep(k):
             return "box_predictor" not in k and "mask_predictor" not in k
@@ -289,7 +291,7 @@ def lambd(batch):
     )
     lr_scheduler.last_epoch = start_epoch - 1
     # This led to a warning in newer version of PyTorch?
-    #lr_scheduler.step()
+    # lr_scheduler.step()
 
     for epoch in range(start_epoch, end_epoch):
         meters_train = defaultdict(AverageValueMeter)

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

@@ -439,9 +439,7 @@ def train_epoch():
                     max_norm=args.clip_grad_norm,
                     norm_type=2,
                 )
-                meters_train["grad_norm"].add(
-                    torch.as_tensor(total_grad_norm).item()
-                )
+                meters_train["grad_norm"].add(torch.as_tensor(total_grad_norm).item())
 
                 optimizer.step()
                 meters_time["backward"].add(time.time() - t)
@@ -463,7 +461,6 @@ def validation():
                 loss = h(data=sample, meters=meters_val)
                 meters_val["loss_total"].add(loss.item())
 
-
         @torch.no_grad()
         def test():
             model.eval()

happypose/pose_estimators/megapose/evaluation/bop.py

@@ -49,6 +49,7 @@
 # Third Party
 from bop_toolkit_lib import inout  # noqa
 
+
 def main():
     parser = argparse.ArgumentParser("Bop evaluation")
     parser.add_argument("--results_path", default="", type=str)
@@ -121,16 +122,16 @@ def convert_results_to_bop(
         print("row =", row)
         obj_id = int(row.label.split("_")[-1])
         if use_pose_score:
-            score = row['pose_score']
+            score = row["pose_score"]
         else:
-            score = row['score']
+            score = row["score"]
         if "time" in row:
-            time = row['time']
+            time = row["time"]
         else:
             time = -1
         pred = dict(
-            scene_id=row['scene_id'],
-            im_id=row['view_id'],
+            scene_id=row["scene_id"],
+            im_id=row["view_id"],
             obj_id=obj_id,
             score=score,
             t=t,
@@ -202,9 +203,7 @@ def run_evaluation(cfg: BOPEvalConfig) -> None:
         csv_path = eval_dir / f"{method}_{cfg.dataset.split('.')[0]}-{cfg.split}.csv"
 
         # pose scores give better AR scores in general
-        convert_results_to_bop(
-            results_path, csv_path, cfg.method, use_pose_score=False
-        )
+        convert_results_to_bop(results_path, csv_path, cfg.method, use_pose_score=False)
 
         if not cfg.convert_only:
             _run_bop_evaluation(csv_path, cfg.eval_dir, eval_detection=False)

happypose/toolbox/utils/download.py

@@ -170,9 +170,7 @@ async def main():
 
     if args.ycbv_tests:
         to_dl.append(("ycbv-debug.zip", DOWNLOAD_DIR))
-        to_unzip.append(
-            (DOWNLOAD_DIR / "ycbv-debug.zip", LOCAL_DATA_DIR / "results")
-        )
+        to_unzip.append((DOWNLOAD_DIR / "ycbv-debug.zip", LOCAL_DATA_DIR / "results"))
 
     if args.cosypose_models:
         for model in args.cosypose_models:

tests/test_batch_renderer_panda3d.py

@@ -115,7 +115,12 @@ def test_batch_renderer(self, device):
         # Renders from 2 identical cams are equals
         assert tr_assert_close(renderings.rgbs[0], renderings.rgbs[1]) is None
         assert tr_assert_close(renderings.normals[0], renderings.normals[1]) is None
-        assert tr_assert_close(renderings.depths[0], renderings.depths[1], atol=1e-3, rtol=1e-3) is None
+        assert (
+            tr_assert_close(
+                renderings.depths[0], renderings.depths[1], atol=1e-3, rtol=1e-3
+            )
+            is None
+        )
         assert (
             tr_assert_close(renderings.binary_masks[0], renderings.binary_masks[1])
             is None

tests/test_cosypose_detector_training.py

@@ -1,32 +1,27 @@
 import os
-import numpy as np
-import pytest 
 
+import numpy as np
+import pytest
 from omegaconf import OmegaConf
 
 from happypose.pose_estimators.cosypose.cosypose.utils.distributed import (
-    get_rank,
     get_world_size,
     init_distributed_mode,
     reduce_dict,
     sync_model,
 )
-
-from happypose.toolbox.utils.logging import get_logger, set_logging_level
+from happypose.toolbox.utils.logging import get_logger
 
 logger = get_logger(__name__)
 
-class TestCosyposeDetectorTraining():
-    
+
+class TestCosyposeDetectorTraining:
     @pytest.fixture(autouse=True)
     def setup(self):
+        args = {"config": "bop-ycbv-synt+real"}
 
-        args = {
-            'config':'bop-ycbv-synt+real'
-        }
-
         cfg_detector = OmegaConf.create({})
-        
+
         logger.info(
             f"Training with config: {args['config']}",
         )
@@ -82,22 +77,22 @@ def setup(self):
         cfg_detector.classifier_alpha = 1
         cfg_detector.mask_alpha = 1
         cfg_detector.box_reg_alpha = 1
-        if "tless" in args['config']:
+        if "tless" in args["config"]:
             cfg_detector.input_resize = (540, 720)
-        elif "ycbv" in args['config']:
+        elif "ycbv" in args["config"]:
             cfg_detector.input_resize = (480, 640)
-        elif "bop-" in args['config']:
+        elif "bop-" in args["config"]:
             cfg_detector.input_resize = None
         else:
             raise ValueError
 
-        if "bop-" in args['config']:
+        if "bop-" in args["config"]:
             from happypose.pose_estimators.cosypose.cosypose.bop_config import (
                 BOP_CONFIG,
                 PBR_DETECTORS,
             )
 
-            bop_name, train_type = args['config'].split("-")[1:]
+            bop_name, train_type = args["config"].split("-")[1:]
             bop_cfg = BOP_CONFIG[bop_name]
             if train_type == "pbr":
                 cfg_detector.train_ds_names = [(bop_cfg["train_pbr_ds_name"][0], 1)]
@@ -112,25 +107,23 @@ def setup(self):
                 cfg_detector.test_ds_names = bop_cfg["test_ds_name"]
 
         else:
-            raise ValueError(args['config'])
+            raise ValueError(args["config"])
         cfg_detector.val_ds_names = cfg_detector.train_ds_names
 
         cfg_detector.run_id = f"detector-{args['config']}-{run_comment}-{N_RAND}"
 
         N_GPUS = int(os.environ.get("N_PROCS", 1))
         cfg_detector.epoch_size = cfg_detector.epoch_size // N_GPUS
         self.cfg_detector = cfg_detector
-    
+
     def test_detector_training(self):
         train_detector(self.cfg_detector)
-        
+
 
 import functools
 import time
 from collections import defaultdict
 
-import numpy as np
-import simplejson as json
 import torch
 import torch.distributed as dist
 import yaml
@@ -144,18 +137,17 @@ def test_detector_training(self):
 from happypose.pose_estimators.cosypose.cosypose.datasets.detection_dataset import (
     DetectionDataset,
 )
-from happypose.pose_estimators.cosypose.cosypose.integrated.detector import Detector
 
 # Evaluation
-from happypose.pose_estimators.cosypose.cosypose.scripts.run_detection_eval import (
-    run_detection_eval,
+from happypose.pose_estimators.cosypose.cosypose.training.detector_models_cfg import (
+    check_update_config,
+    create_model_detector,
 )
-from happypose.pose_estimators.cosypose.cosypose.utils.distributed import (
-    get_rank,
-    get_world_size,
-    init_distributed_mode,
-    reduce_dict,
-    sync_model,
+from happypose.pose_estimators.cosypose.cosypose.training.maskrcnn_forward_loss import (
+    h_maskrcnn,
+)
+from happypose.pose_estimators.cosypose.cosypose.training.train_detector import (
+    collate_fn,
 )
 from happypose.pose_estimators.cosypose.cosypose.utils.logging import get_logger
 from happypose.toolbox.datasets.datasets_cfg import make_scene_dataset
@@ -169,11 +161,6 @@ def test_detector_training(self):
     get_total_memory,
 )
 
-from happypose.pose_estimators.cosypose.cosypose.training.detector_models_cfg import check_update_config, create_model_detector
-from happypose.pose_estimators.cosypose.cosypose.training.maskrcnn_forward_loss import h_maskrcnn
-
-from happypose.pose_estimators.cosypose.cosypose.training.train_detector import collate_fn
-
 cudnn.benchmark = True
 logger = get_logger(__name__)
 
@@ -293,7 +280,9 @@ def make_datasets(dataset_names):
         logger.info(f"Using pretrained model from {pretrain_path}.")
         model.load_state_dict(torch.load(pretrain_path)["state_dict"])
     elif args.pretrain_coco:
-        state_dict = load_state_dict_from_url('https://download.pytorch.org/models/maskrcnn_resnet50_fpn_coco-bf2d0c1e.pth')
+        state_dict = load_state_dict_from_url(
+            "https://download.pytorch.org/models/maskrcnn_resnet50_fpn_coco-bf2d0c1e.pth"
+        )
 
         def keep(k):
             return "box_predictor" not in k and "mask_predictor" not in k
@@ -383,7 +372,9 @@ def train_epoch():
                         max_norm=np.inf,
                         norm_type=2,
                     )
-                    meters_train["grad_norm"].add(torch.as_tensor(total_grad_norm).item())
+                    meters_train["grad_norm"].add(
+                        torch.as_tensor(total_grad_norm).item()
+                    )
 
                     optimizer.step()
                     meters_time["backward"].add(time.time() - t)
@@ -444,4 +435,3 @@ def validation():
         log_dict = reduce_dict(log_dict)
 
         dist.barrier()
-