torch: always with cuda

.github/workflows/conda-test.yml

@@ -52,7 +52,7 @@ jobs:
         if: steps.cache.outputs.cache-hit != 'true'
 
       - name: Install happypose
-        run: pip install -e .
+        run: pip install -e .[multiview]
 
       - name: Download pre-trained models required for tests
         run: |

.github/workflows/pip-test.yml

@@ -30,7 +30,7 @@ jobs:
         run: pip install -U pip
 
       - name: Install happypose
-        run: pip install ".[cpu,multiview,pypi]" --extra-index-url https://download.pytorch.org/whl/cpu
+        run: pip install ".[multiview,pypi]"
 
       - name: Download pre-trained models required for tests
         run: |

.github/workflows/poetry-test.yml

@@ -31,7 +31,7 @@ jobs:
           cache: poetry
 
       - name: Install happypose
-        run: poetry install --with dev -E cpu -E pypi -E multiview
+        run: poetry install --with dev -E pypi -E multiview
 
       - name: Download pre-trained models required for tests
         run: |

.pre-commit-config.yaml

@@ -1,6 +1,6 @@
 repos:
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.3.2
+  rev: v0.5.2
   hooks:
   - id: ruff
     args:
@@ -13,7 +13,7 @@ repos:
   - id: toml-sort-fix
     exclude: poetry.lock
 - repo: https://github.com/pre-commit/pre-commit-hooks
-  rev: v4.5.0
+  rev: v4.6.0
   hooks:
   - id: check-added-large-files
   - id: check-ast

README.md

@@ -39,15 +39,13 @@ git clone --branch dev --recurse-submodules https://github.com/agimus-project/ha
 cd happypose
 python -m venv .venv
 source .venv/bin/activate
-pip install .[pypi,cpu] --extra-index-url https://download.pytorch.org/whl/cpu
+pip install .[pypi]
 ```
 
 ### Install extras:
 
-- `cpu`: required to get pytorch CPU from PyPI (don't use this for GPU or with conda)
-- `gpu`: required to get pytorch GPU from PyPI (don't use this for CPU or with conda)
 - `multiview`: installs cosypose c++ extension
-- `pypi`: install pinocchio & opencv from PyPI (don't use this with conda)
+- `pypi`: install torch, pinocchio & opencv from PyPI (don't use this with conda)
 
 ## Create data directory
 

environment.yml

@@ -7,11 +7,9 @@ channels:
    - defaults
 dependencies:
   - pytorch-cuda==12.1
-  - python=3.9
   - pip
-  - pytorch==2.1
-  - torchvision==0.16
-  - mkl==2024.0.0
+  - pytorch
+  - torchvision
   - geckodriver
   - firefox
   - opencv

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: