diff --git a/deeplabcut/modelzoo/model_configs/rtmpose_s.yaml b/deeplabcut/modelzoo/model_configs/rtmpose_s.yaml index 2e7d693ba0..9c80b0583c 100644 --- a/deeplabcut/modelzoo/model_configs/rtmpose_s.yaml +++ b/deeplabcut/modelzoo/model_configs/rtmpose_s.yaml @@ -51,6 +51,10 @@ model: predictor: type: SimCCPredictor simcc_split_ratio: 2.0 + sigma: + - 5.66 + - 5.66 + decode_beta: 150.0 input_size: - 256 - 256 diff --git a/deeplabcut/modelzoo/model_configs/rtmpose_x.yaml b/deeplabcut/modelzoo/model_configs/rtmpose_x.yaml index 9a7df70196..0d1fb8a547 100644 --- a/deeplabcut/modelzoo/model_configs/rtmpose_x.yaml +++ b/deeplabcut/modelzoo/model_configs/rtmpose_x.yaml @@ -102,6 +102,10 @@ model: predictor: type: SimCCPredictor simcc_split_ratio: 2.0 + sigma: + - 6.0 + - 6.93 + decode_beta: 150.0 input_size: - 288 - 384 diff --git a/deeplabcut/pose_estimation_pytorch/config/ctd/ctd_prenet_rtmpose_m.yaml b/deeplabcut/pose_estimation_pytorch/config/ctd/ctd_prenet_rtmpose_m.yaml index ca35c8a3ed..88b708cc39 100644 --- a/deeplabcut/pose_estimation_pytorch/config/ctd/ctd_prenet_rtmpose_m.yaml +++ b/deeplabcut/pose_estimation_pytorch/config/ctd/ctd_prenet_rtmpose_m.yaml @@ -55,6 +55,8 @@ model: predictor: type: SimCCPredictor simcc_split_ratio: 2.0 + sigma: [5.66, 5.66] + decode_beta: 150.0 input_size: [256, 256] in_channels: 768 out_channels: "num_bodyparts" diff --git a/deeplabcut/pose_estimation_pytorch/config/ctd/ctd_prenet_rtmpose_s.yaml b/deeplabcut/pose_estimation_pytorch/config/ctd/ctd_prenet_rtmpose_s.yaml index c1870e8160..6ae8b5364a 100644 --- a/deeplabcut/pose_estimation_pytorch/config/ctd/ctd_prenet_rtmpose_s.yaml +++ b/deeplabcut/pose_estimation_pytorch/config/ctd/ctd_prenet_rtmpose_s.yaml @@ -55,6 +55,8 @@ model: predictor: type: SimCCPredictor simcc_split_ratio: 2.0 + sigma: [5.66, 5.66] + decode_beta: 150.0 input_size: [256, 256] in_channels: 512 out_channels: "num_bodyparts" diff --git a/deeplabcut/pose_estimation_pytorch/config/ctd/ctd_prenet_rtmpose_x.yaml b/deeplabcut/pose_estimation_pytorch/config/ctd/ctd_prenet_rtmpose_x.yaml index 75c7228d21..f809a0c569 100644 --- a/deeplabcut/pose_estimation_pytorch/config/ctd/ctd_prenet_rtmpose_x.yaml +++ b/deeplabcut/pose_estimation_pytorch/config/ctd/ctd_prenet_rtmpose_x.yaml @@ -55,6 +55,8 @@ model: predictor: type: SimCCPredictor simcc_split_ratio: 2.0 + sigma: [6.93, 6.93] + decode_beta: 150.0 input_size: [384, 384] in_channels: 1280 out_channels: "num_bodyparts" diff --git a/deeplabcut/pose_estimation_pytorch/config/ctd/ctd_prenet_rtmpose_x_human.yaml b/deeplabcut/pose_estimation_pytorch/config/ctd/ctd_prenet_rtmpose_x_human.yaml index ad6579fe5e..1d47cb3306 100644 --- a/deeplabcut/pose_estimation_pytorch/config/ctd/ctd_prenet_rtmpose_x_human.yaml +++ b/deeplabcut/pose_estimation_pytorch/config/ctd/ctd_prenet_rtmpose_x_human.yaml @@ -56,6 +56,8 @@ model: predictor: type: SimCCPredictor simcc_split_ratio: 2.0 + sigma: [6., 6.93] + decode_beta: 150.0 input_size: [288, 384] in_channels: 1280 out_channels: "num_bodyparts" diff --git a/deeplabcut/pose_estimation_pytorch/config/rtmpose/rtmpose_m.yaml b/deeplabcut/pose_estimation_pytorch/config/rtmpose/rtmpose_m.yaml index d6bc515f94..d2ed3ae52e 100644 --- a/deeplabcut/pose_estimation_pytorch/config/rtmpose/rtmpose_m.yaml +++ b/deeplabcut/pose_estimation_pytorch/config/rtmpose/rtmpose_m.yaml @@ -49,6 +49,8 @@ model: predictor: type: SimCCPredictor simcc_split_ratio: 2.0 + sigma: [5.66, 5.66] + decode_beta: 150.0 input_size: [256, 256] in_channels: 768 out_channels: "num_bodyparts" diff --git a/deeplabcut/pose_estimation_pytorch/config/rtmpose/rtmpose_s.yaml b/deeplabcut/pose_estimation_pytorch/config/rtmpose/rtmpose_s.yaml index fbc4ff7ed4..463e104654 100644 --- a/deeplabcut/pose_estimation_pytorch/config/rtmpose/rtmpose_s.yaml +++ b/deeplabcut/pose_estimation_pytorch/config/rtmpose/rtmpose_s.yaml @@ -49,6 +49,8 @@ model: predictor: type: SimCCPredictor simcc_split_ratio: 2.0 + sigma: [5.66, 5.66] + decode_beta: 150.0 input_size: [256, 256] in_channels: 512 out_channels: "num_bodyparts" diff --git a/deeplabcut/pose_estimation_pytorch/config/rtmpose/rtmpose_x.yaml b/deeplabcut/pose_estimation_pytorch/config/rtmpose/rtmpose_x.yaml index 0a49baec75..f1d6f61e67 100644 --- a/deeplabcut/pose_estimation_pytorch/config/rtmpose/rtmpose_x.yaml +++ b/deeplabcut/pose_estimation_pytorch/config/rtmpose/rtmpose_x.yaml @@ -49,6 +49,8 @@ model: predictor: type: SimCCPredictor simcc_split_ratio: 2.0 + sigma: [6.93, 6.93] + decode_beta: 150.0 input_size: [384, 384] in_channels: 1280 out_channels: "num_bodyparts" diff --git a/deeplabcut/pose_estimation_pytorch/models/predictors/sim_cc.py b/deeplabcut/pose_estimation_pytorch/models/predictors/sim_cc.py index 25d953c178..b36a9639ab 100644 --- a/deeplabcut/pose_estimation_pytorch/models/predictors/sim_cc.py +++ b/deeplabcut/pose_estimation_pytorch/models/predictors/sim_cc.py @@ -42,21 +42,33 @@ class SimCCPredictor(BasePredictor): def __init__( self, simcc_split_ratio: float = 2.0, - apply_softmax: bool = False, + apply_softmax: bool = True, normalize_outputs: bool = False, + sigma: float | int | tuple[float, ...] = 6.0, + decode_beta: float = 150.0, ) -> None: super().__init__() self.simcc_split_ratio = simcc_split_ratio self.apply_softmax = apply_softmax self.normalize_outputs = normalize_outputs + if isinstance(sigma, (float, int)): + self.sigma = np.array([sigma, sigma]) + else: + self.sigma = np.array(sigma) + self.decode_beta = decode_beta + def forward( self, stride: float, outputs: dict[str, torch.Tensor] ) -> dict[str, torch.Tensor]: x, y = outputs["x"].detach(), outputs["y"].detach() + if self.normalize_outputs: x = get_simcc_normalized(x) y = get_simcc_normalized(y) + else: + x = x * (self.sigma[0] * self.decode_beta) + y = y * (self.sigma[1] * self.decode_beta) keypoints, scores = get_simcc_maximum( x.cpu().numpy(), y.cpu().numpy(), self.apply_softmax