[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

happypose/pose_estimators/cosypose/cosypose/datasets/detection_dataset.py

@@ -1,6 +1,5 @@
 import random
 from dataclasses import dataclass
-from typing import List, Optional
 
 import numpy as np
 import torch
@@ -29,20 +28,18 @@
     PillowSharpness,
     VOCBackgroundAugmentation,
 )
+from happypose.toolbox.datasets.augmentations import (
+    SceneObservationAugmentation as SceneObsAug,
+)
 
 # HappyPose
 from happypose.toolbox.datasets.scene_dataset import (
     IterableSceneDataset,
-    ObjectData,
     SceneDataset,
     SceneObservation,
 )
-
-from happypose.toolbox.datasets.augmentations import (
-    SceneObservationAugmentation as SceneObsAug,
-)
-
 from happypose.toolbox.datasets.scene_dataset_wrappers import remove_invisible_objects
+
 # from happypose.toolbox.datasets.scene_dataset_wrappers import VisibilityWrapper
 """
 from .augmentations import (
@@ -58,6 +55,7 @@
 from .wrappers.visibility_wrapper import VisibilityWrapper
 """
 
+
 def collate_fn(batch):
     print("inside collate fn")
     rgbs, targets = zip(*batch)
@@ -92,7 +90,8 @@ def collate_fn(batch):
     # Return the batch data as a dictionary
     return {"rgbs": rgbs, "targets": target}
 
-#TODO : Double check on types and add format documentation
+
+# TODO : Double check on types and add format documentation
 @dataclass
 class DetectionData:
     """rgb: (h, w, 3) uint8
@@ -135,6 +134,7 @@ def pin_memory(self) -> "BatchDetectionData":
             self.depths = self.depths.pin_memory()
         return self
 
+
 class DetectionDataset(torch.utils.data.IterableDataset):
     def __init__(
         self,
@@ -168,13 +168,14 @@ def __init__(
                     SceneObsAug(PillowSharpness(factor_interval=(0.0, 50.0)), p=0.3),
                     SceneObsAug(PillowContrast(factor_interval=(0.2, 50.0)), p=0.3),
                     SceneObsAug(PillowBrightness(factor_interval=(0.1, 6.0)), p=0.5),
-                    SceneObsAug(PillowColor(factor_interval=(0.0, 20.0)), p=0.3),]
+                    SceneObsAug(PillowColor(factor_interval=(0.0, 20.0)), p=0.3),
+                ]
             )
-
         ]
 
         self.label_to_category_id = label_to_category_id
         self.min_area = min_area
+
     """
     def collate_fn(self, list_data: List[DetectionData]) -> BatchDetectionData:
         batch_data = BatchDetectionData(
@@ -194,19 +195,19 @@ def collate_fn(self, list_data: List[DetectionData]) -> BatchDetectionData:
 
         return batch_data
     """
+
     def make_data_from_obs(self, obs: SceneObservation, idx):
-
         obs = remove_invisible_objects(obs)
 
         obs = self.resize_augmentation(obs)
 
         for aug in self.background_augmentations:
             obs = aug(obs)
-            
+
         if self.rgb_augmentations and random.random() < 0.8:
             for aug in self.rgb_augmentations:
                 obs = aug(obs)
-        
+
         assert obs.object_datas is not None
         assert obs.rgb is not None
         assert obs.camera_data is not None
@@ -226,13 +227,13 @@ def make_data_from_obs(self, obs: SceneObservation, idx):
             if obs.binary_masks is not None:
                 binary_mask = torch.tensor(obs.binary_masks[obj_data.unique_id]).float()
                 masks.append(binary_mask)
-                
+
             if obs.segmentation is not None:
                 binary_mask = np.zeros_like(obs.segmentation, dtype=np.bool_)
                 binary_mask[obs.segmentation == obj_data.unique_id] = 1
                 binary_mask = torch.as_tensor(binary_mask).float()
                 masks.append(binary_mask)
-                
+
         masks = np.array(masks)
         masks = masks == obj_ids[:, None, None]
 
@@ -258,17 +259,15 @@ def make_data_from_obs(self, obs: SceneObservation, idx):
         target["image_id"] = image_id
         target["area"] = area
         target["iscrowd"] = iscrowd
-        
+
         return rgb, target
-
 
-
     def __getitem__(self, index: int):
         assert isinstance(self.scene_ds, SceneDataset)
         obs = self.scene_ds[index]
         return self.make_data_from_obs(obs, index)
-    
-      # def find_valid_data(self, iterator: Iterator[SceneObservation]) -> PoseData:
+
+    # def find_valid_data(self, iterator: Iterator[SceneObservation]) -> PoseData:
     def find_valid_data(self, iterator):
         n_attempts = 0
         for idx, obs in enumerate(iterator):
@@ -279,17 +278,13 @@ def find_valid_data(self, iterator):
             if n_attempts > 200:
                 msg = "Cannot find valid image in the dataset"
                 raise ValueError(msg)
-        
+
     def __iter__(self):
         assert isinstance(self.scene_ds, IterableSceneDataset)
         iterator = iter(self.scene_ds)
         while True:
             yield self.find_valid_data(iterator)
-
-
-
-
-
+
     def __init___old(
         self,
         scene_ds,
@@ -322,10 +317,7 @@ def __init___old(
         self.label_to_category_id = label_to_category_id
         self.min_area = min_area
 
-
-
     def get_data_old(self, idx):
-
         print("I am in get_data")
         rgb, mask, state = self.scene_ds[idx]
 
@@ -393,4 +385,3 @@ def __getitem__old(self, index):
                 try_index = random.randint(0, len(self.scene_ds) - 1)
                 n_attempts += 1
         return im, target
-

happypose/pose_estimators/cosypose/cosypose/datasets/synthetic_dataset.py

@@ -14,7 +14,6 @@
     CameraData,
     ObjectData,
     ObservationInfos,
-    SceneDataset,
     SceneObservation,
 )
 from happypose.toolbox.lib3d.transform import Transform
@@ -82,33 +81,33 @@ def __getitem__(self, idx):
             if obj["id_in_segm"] in mask_uniqs:
                 obj["bbox"] = dets_gt[obj["id_in_segm"]].numpy()
                 obj_data = ObjectData(
-                    label= obj['name'],
-                    TWO=Transform(obj['TWO']),
-                    unique_id=obj['body_id'],
-                    bbox_modal=obj['bbox']
+                    label=obj["name"],
+                    TWO=Transform(obj["TWO"]),
+                    unique_id=obj["body_id"],
+                    bbox_modal=obj["bbox"],
                 )
             else:
                 obj_data = ObjectData(
-                    label= obj['name'],
-                    TWO=Transform(obj['TWO']),
-                    unique_id=obj['body_id']
+                    label=obj["name"],
+                    TWO=Transform(obj["TWO"]),
+                    unique_id=obj["body_id"],
                 )
             object_datas.append(obj_data)
-            
+
         state = {
             "camera": cam,
             "objects": objects,
             "frame_info": self.frame_index.iloc[idx].to_dict(),
         }
         image_infos = ObservationInfos(
-            scene_id=self.frame_index.iloc[idx].to_dict()['scene_id'],
-            view_id=self.frame_index.iloc[idx].to_dict()['view_id'])
-
+            scene_id=self.frame_index.iloc[idx].to_dict()["scene_id"],
+            view_id=self.frame_index.iloc[idx].to_dict()["view_id"],
+        )
+
         cam = CameraData(
-            K=cam["K"],
-            resolution=cam["resolution"],
-            TWC=Transform(cam["TWC"]))
-
+            K=cam["K"], resolution=cam["resolution"], TWC=Transform(cam["TWC"])
+        )
+
         observation = SceneObservation(
             rgb=rgb.numpy().astype(np.uint8),
             segmentation=mask.numpy().astype(np.uint32),

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

@@ -94,18 +94,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']}"
             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(
@@ -144,7 +144,7 @@ def run_inference_pipeline(
             "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

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

@@ -14,7 +14,7 @@
 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():
-        pose_predictor = pred_kwargs_n['pose_predictor']
+        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

@@ -180,7 +180,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
+                model_outputs=model_outputs,
             )
 
             TCO_input = TCO_output

happypose/pose_estimators/cosypose/cosypose/scripts/run_cosypose_eval.py

@@ -22,10 +22,6 @@
     RESULTS_DIR,
 )
 from happypose.pose_estimators.cosypose.cosypose.datasets.bop import remap_bop_targets
-from happypose.toolbox.datasets.datasets_cfg import (
-    make_object_dataset,
-    make_scene_dataset,
-)
 from happypose.pose_estimators.cosypose.cosypose.datasets.samplers import ListSampler
 from happypose.pose_estimators.cosypose.cosypose.datasets.wrappers.multiview_wrapper import (
     MultiViewWrapper,
@@ -61,6 +57,10 @@
     init_distributed_mode,
 )
 from happypose.pose_estimators.cosypose.cosypose.utils.logging import get_logger
+from happypose.toolbox.datasets.datasets_cfg import (
+    make_object_dataset,
+    make_scene_dataset,
+)
 from happypose.toolbox.lib3d.transform import Transform
 from happypose.toolbox.renderer.bullet_batch_renderer import BulletBatchRenderer
 

happypose/pose_estimators/cosypose/cosypose/scripts/run_inference_on_example.py

@@ -85,7 +85,9 @@ def run_inference(
     rgb, depth, camera_data = load_observation_example(example_dir, load_depth=True)
     # TODO: cosypose forward does not work if depth is loaded detection
     # contrary to megapose
-    observation = ObservationTensor.from_numpy(rgb, depth=None, K=camera_data.K).to(device)
+    observation = ObservationTensor.from_numpy(rgb, depth=None, K=camera_data.K).to(
+        device
+    )
 
     # Load models
     pose_estimator = setup_pose_estimator(args.dataset, object_dataset)

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

@@ -11,8 +11,10 @@
 
 
 import warnings
+
 warnings.filterwarnings("ignore")
 
+
 def make_cfg(args):
     cfg = argparse.ArgumentParser("").parse_args([])
     if args.config:

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

@@ -68,7 +68,6 @@ def h_pose(model, mesh_db, data, meters, cfg, n_iterations=1, input_generator="f
         TCO_pred = iter_outputs.TCO_output
         model_outputs = iter_outputs.model_outputs
 
-
         if cfg.loss_disentangled:
             if cfg.n_pose_dims == 9:
                 loss_fn = loss_refiner_CO_disentangled