Merge pull request #27 from cameron-a-johnson/dev/optional-bbox-thr

Dev/optional bbox thr

setup.py

@@ -10,7 +10,7 @@
     author_email="",
     url="https://github.com/user/project",
     # install_requires=["lightning", "hydra-core", "hydra-colorlog"],
-    packages=find_packages(include=['tcn_hpl', 'tcn_hpl.*']),
+    packages=find_packages(include=["tcn_hpl", "tcn_hpl.*"]),
     # use this to customize global commands available in the terminal after installing the package
     entry_points={
         "console_scripts": [

tcn_hpl/data/components/augmentations.py

@@ -727,15 +727,16 @@ def forward(self, features):
     def __repr__(self) -> str:
         detail = f"(im_w={self.im_w}, im_h={self.im_h}, feat_version={self.feat_version}, top_k_objects={self.top_k_objects})"
         return f"{self.__class__.__name__}{detail}"
-
 
 
 class DropoutObjects(torch.nn.Module):
-    """Drop out Objects """
+    """Drop out Objects"""
 
-    def __init__(self, dropout_probablity, num_obj_classes, feat_version, top_k_objects):
+    def __init__(
+        self, dropout_probablity, num_obj_classes, feat_version, top_k_objects
+    ):
         """
-        :param dropout_probablity: probablity that a given frame will NOT have an object 
+        :param dropout_probablity: probablity that a given frame will NOT have an object
         """
         super().__init__()
 
@@ -769,8 +770,8 @@ def __init__(self, dropout_probablity, num_obj_classes, feat_version, top_k_obje
                 self.obj_feature_mask.append(0)
 
             if self.use_hand_dist:
-                self.obj_feature_mask += [0]*2*self.num_good_obj_classes
-                ind += 2*self.num_good_obj_classes
+                self.obj_feature_mask += [0] * 2 * self.num_good_obj_classes
+                ind += 2 * self.num_good_obj_classes
 
             if self.use_center_dist:
                 # right hand - image center distance
@@ -785,8 +786,8 @@ def __init__(self, dropout_probablity, num_obj_classes, feat_version, top_k_obje
 
             if self.use_hand_dist:
                 # Left hand distances
-                self.obj_feature_mask += [0]*2*self.num_good_obj_classes
-                ind += 2*self.num_good_obj_classes
+                self.obj_feature_mask += [0] * 2 * self.num_good_obj_classes
+                ind += 2 * self.num_good_obj_classes
 
             if self.use_center_dist:
                 # left hand - image center distance
@@ -825,34 +826,33 @@ def __init__(self, dropout_probablity, num_obj_classes, feat_version, top_k_obje
         if self.use_joint_hand_offset:
             # left hand - joints distances
             ind += 44
-            self.obj_feature_mask += [1]*44
+            self.obj_feature_mask += [1] * 44
 
             # right hand - joints distances
             ind += 44
-            self.obj_feature_mask += [1]*44
+            self.obj_feature_mask += [1] * 44
 
         # OBJS-JOINTS
         if self.use_joint_object_offset:
-            self.obj_feature_mask += [0]*44*self.top_k_objects*self.num_good_obj_classes
-            ind += 44*self.top_k_objects*self.num_good_obj_classes
-
+            self.obj_feature_mask += (
+                [0] * 44 * self.top_k_objects * self.num_good_obj_classes
+            )
+            ind += 44 * self.top_k_objects * self.num_good_obj_classes
+
         self.obj_feature_mask = torch.tensor(self.obj_feature_mask)
 
     def forward(self, features):
         num_frames = features.shape[0]
         # Pick starting location of random mask
-        start = random.randint(0,self.skip_stride)
+        start = random.randint(0, self.skip_stride)
         # Create mask (one element for each frame)
         mask = torch.rand(num_frames) > self.dropout_probablity
 
         if self.dropout_last:
             mask[-1] = 0
 
-
-
         return features
 
     def __repr__(self) -> str:
         detail = f"(im_w={self.im_w}, im_h={self.im_h}, num_obj_classes={self.num_obj_classes}, feat_version={self.feat_version}, top_k_objects={self.top_k_objects})"
         return f"{self.__class__.__name__}{detail}"
-

tcn_hpl/data/utils/pose_generation/rt_pose_generation.py

@@ -5,7 +5,7 @@
 import warnings
 
 
-def predict_single(det_model, pose_model, image: torch.tensor) -> list:
+def predict_single(det_model, pose_model, image: torch.tensor, bbox_thr=0.0) -> list:
 
     keypoints_cats = [
         "nose",
@@ -81,7 +81,7 @@ def predict_single(det_model, pose_model, image: torch.tensor) -> list:
             current_ann["label"] = pred_label
             current_ann["bbox_score"] = f"{scores[box_id] * 100:0.2f}"
 
-            if box_class == 0:
+            if box_class == 0 and float(current_ann["bbox_score"]) > bbox_thr:
                 person_results = [current_ann]
 
                 pose_results, returned_outputs = inference_top_down_pose_model(

tcn_hpl/data/utils/ptg_datagenerator.py

@@ -25,12 +25,18 @@
 from angel_system.data.medical.load_bbn_data import bbn_activity_txt_to_csv
 from angel_system.data.common.kwcoco_utils import add_activity_gt_to_kwcoco
 
+
 def bbn_to_dive(raw_data_root, dive_output_dir, task, label_mapping, label_version=1):
-    used_videos = bbn_activity_txt_to_csv(task=task, root_dir=raw_data_root, 
-                                        output_dir=dive_output_dir, label_mapping=label_mapping,
-                                        label_version=label_version)
+    used_videos = bbn_activity_txt_to_csv(
+        task=task,
+        root_dir=raw_data_root,
+        output_dir=dive_output_dir,
+        label_mapping=label_mapping,
+        label_version=label_version,
+    )
     return used_videos
 
+
 def create_training_data(config_path):
     #####################
     # Input
@@ -56,10 +62,11 @@ def create_training_data(config_path):
     top_k_objects = config["data_gen"]["top_k_objects"]
     pose_repeat_rate = config["data_gen"]["pose_repeat_rate"]
     exp_ext = config["data_gen"]["exp_ext"]
-    raw_data_root = f"{config['data_gen']['raw_data_root']}/{TASK_TO_NAME[task_name]}/Data"
+    raw_data_root = (
+        f"{config['data_gen']['raw_data_root']}/{TASK_TO_NAME[task_name]}/Data"
+    )
     dive_output_root = config["data_gen"]["dive_output_root"]
 
-
     def filter_dset_by_split(split, dataset_to_split):
         # Filter by video names
         video_lookup = dataset_to_split.index.name_to_video
@@ -139,28 +146,34 @@ def filter_dset_by_split(split, dataset_to_split):
             if label_str == "done":
                 continue
             mapping.write(f"{i} {label_str}\n")
-    
+
     print(f"label mapping: {activity_labels_desc_mapping}")
     # exit()
     dset = kwcoco.CocoDataset(config["data_gen"]["dataset_kwcoco"])
     # Check if the dest has activity gt, if it doesn't then add it
     if not "activity_gt" in list(dset.imgs.values())[0].keys():
         print("adding activity ground truth to the dataset")
-        
-        #generate dive files for videos in dataset if it does not exist
+
+        # generate dive files for videos in dataset if it does not exist
         video_id = list(dset.index.videos.keys())[0]
         video = dset.index.videos[video_id]
         video_name = video["name"]
         activity_gt_dir = f"{dive_output_root}/{task_name}_labels/"
         activity_gt_fn = f"{activity_gt_dir}/{video_name}_activity_labels_v1.csv"
         print(f"activity_gt_dir: {activity_gt_dir}, activity_gt_fn: {activity_gt_fn}")
 
-        used_videos = bbn_to_dive(raw_data_root, activity_gt_dir, task_name, activity_labels_desc_mapping)
-        video_ids_to_remove = [vid for vid, value in dset.index.videos.items() if value['name'] not in used_videos]
+        used_videos = bbn_to_dive(
+            raw_data_root, activity_gt_dir, task_name, activity_labels_desc_mapping
+        )
+        video_ids_to_remove = [
+            vid
+            for vid, value in dset.index.videos.items()
+            if value["name"] not in used_videos
+        ]
         dset.remove_videos(video_ids_to_remove)
-        
+
         dset = add_activity_gt_to_kwcoco(topic, task_name, dset, activity_config_fn)
-    
+
     #####################
     # Features,
     # groundtruth and
@@ -223,7 +236,7 @@ def filter_dset_by_split(split, dataset_to_split):
                 ann_by_image,
                 feat_version=feat_version,
                 top_k_objects=top_k_objects,
-                pose_repeat_rate=pose_repeat_rate
+                pose_repeat_rate=pose_repeat_rate,
             )
             print(X.shape)
 

tcn_hpl/models/ptg_module.py

@@ -244,7 +244,7 @@ def compute_loss(self, p, y, mask):
         variation_coef = torch.abs(preds - mode)
         variation_coef = torch.sum(variation_coef, dim=-1)
         gt_variation_coef = torch.zeros_like(variation_coef)
-        
+
         if self.hparams.use_smoothing_loss:
             loss += self.hparams.smoothing_loss * torch.mean(
                 self.mse(
@@ -343,7 +343,7 @@ def on_train_epoch_end(self) -> None:
 
         # acc = self.train_acc.compute()  # get current val acc
         # self.train_acc_best(acc)  # update best so far val acc
-        
+
         all_targets = torch.cat(self.training_step_outputs_target)  # shape: #frames
         all_preds = torch.cat(self.training_step_outputs_pred)  # shape: #frames
         all_probs = torch.cat(
@@ -369,7 +369,7 @@ def on_train_epoch_end(self) -> None:
                 self.train_frames[video[:-4]] = train_fns
 
         per_video_frame_gt_preds = {}
-        
+
         for (gt, pred, source_vid, source_frame) in zip(
             all_targets, all_preds, all_source_vids, all_source_frames
         ):
@@ -476,7 +476,6 @@ def validation_step(
         self.validation_step_outputs_pred.append(preds[inds])
         self.validation_step_outputs_prob.append(probs[inds])
 
-
     def on_validation_epoch_end(self) -> None:
         "Lightning hook that is called when a validation epoch ends."
         acc = self.val_acc.compute()  # get current val acc
@@ -488,7 +487,6 @@ def on_validation_epoch_end(self) -> None:
             best_val_acc = self.val_acc_best.compute()
             self.log("val/acc_best", best_val_acc, sync_dist=True, prog_bar=True)
 
-
         all_targets = torch.cat(self.validation_step_outputs_target)  # shape: #frames
         all_preds = torch.cat(self.validation_step_outputs_pred)  # shape: #frames
         all_probs = torch.cat(
@@ -497,7 +495,6 @@ def on_validation_epoch_end(self) -> None:
         all_source_vids = torch.cat(self.validation_step_outputs_source_vid)
         all_source_frames = torch.cat(self.validation_step_outputs_source_frame)
 
-
         # Load val vidoes
         if self.val_frames is None:
             self.val_frames = {}

tcn_hpl/train.py

@@ -55,7 +55,6 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]:
     log.info(f"Instantiating datamodule <{cfg.data._target_}>")
     datamodule: LightningDataModule = hydra.utils.instantiate(cfg.data)
 
-
     log.info(f"Instantiating model <{cfg.model._target_}>")
     model: LightningModule = hydra.utils.instantiate(cfg.model)
 

tcn_hpl/utils/logging_utils.py

@@ -53,7 +53,6 @@ def log_hyperparameters(object_dict: Dict[str, Any]) -> None:
     hparams["tags"] = cfg.get("tags")
     hparams["ckpt_path"] = cfg.get("ckpt_path")
     hparams["seed"] = cfg.get("seed")
-
 
     hparams["cfg"] = cfg
 

tests/conftest.py

@@ -43,7 +43,9 @@ def cfg_eval_global() -> DictConfig:
     :return: A DictConfig containing a default Hydra configuration for evaluation.
     """
     with initialize(version_base="1.3", config_path="../configs"):
-        cfg = compose(config_name="eval.yaml", return_hydra_config=True, overrides=["ckpt_path=."])
+        cfg = compose(
+            config_name="eval.yaml", return_hydra_config=True, overrides=["ckpt_path=."]
+        )
 
         # set defaults for all tests
         with open_dict(cfg):

tests/helpers/run_if.py

@@ -92,9 +92,7 @@ def __new__(
             reasons.append(f"torch<{max_torch}")
 
         if min_python:
-            py_version = (
-                f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}"
-            )
+            py_version = f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}"
             conditions.append(Version(py_version) < Version(min_python))
             reasons.append(f"python>={min_python}")